st-bender /
pyeppaurora
| Conditions | 7 |
| Total Lines | 74 |
| Code Lines | 30 |
| Lines | 0 |
| Ratio | 0 % |
| Changes | 0 | ||
Small methods make your code easier to understand, in particular if combined with a good name. Besides, if your method is small, finding a good name is usually much easier.
For example, if you find yourself adding comments to a method's body, this is usually a good sign to extract the commented part to a new method, and use the comment as a starting point when coming up with a good name for this new method.
Commonly applied refactorings include:
If many parameters/temporary variables are present:
| 1 | # coding: utf-8 |
||
| 30 | def ssusiq2023(gmlat, mlt, alt, sw_coeffs, coeff_ds=None, return_var=False): |
||
| 31 | u""" |
||
| 32 | Parameters |
||
| 33 | ---------- |
||
| 34 | gmlat: float |
||
| 35 | Geomagnetic latitude in [degrees]. |
||
| 36 | mlt: float |
||
| 37 | Magnetic local time in [hours]. |
||
| 38 | alt: float |
||
| 39 | Altitude in [km] |
||
| 40 | sw_coeffs: array_like or `xarray.DataArray` |
||
| 41 | The space weather index values to use (for the requested time(s)), |
||
| 42 | should be of shape (N, M) with N = number of proxies, currently 5: |
||
| 43 | [Kp, PC, Ap, log(f10.7_81ctr_obs), log(v_plasma)]. |
||
| 44 | The `xarray.DataArray` should have a dimension named "proxy" with |
||
| 45 | matching coordinates: |
||
| 46 | ["Kp", "PC", "Ap", "log_f107_81ctr_obs", "log_v_plasma"] |
||
| 47 | All the other dimensions will be broadcasted. |
||
| 48 | coeff_ds: `xarray.Dataset`, optional (default: None) |
||
| 49 | Dataset with the model coefficients, `None` uses the packaged version. |
||
| 50 | return_var: bool, optional (default: False) |
||
| 51 | If `True`, returns the predicted variance in addition to the values, |
||
| 52 | otherwise only the mean prediction is returned. |
||
| 53 | |||
| 54 | Returns |
||
| 55 | ------- |
||
| 56 | q: `xarray.DataArray` |
||
| 57 | log(q), where q is the ionization rate in [cm⁻³ s⁻¹] |
||
| 58 | if `return_var` is False. |
||
| 59 | q, var(q): tuple of `xarray.DataArray`s |
||
| 60 | log(q) and var(log(q)) where q is the ionization rate in [cm⁻³ s⁻¹] |
||
| 61 | if `return_var` is True. |
||
| 62 | """ |
||
| 63 | coeff_ds = coeff_ds or xr.open_dataset(COEFF_PATH, decode_times=False) |
||
| 64 | coeff_sel = coeff_ds.sel( |
||
| 65 | altitude=alt, latitude=gmlat, mlt=mlt, method="nearest", |
||
| 66 | ) |
||
| 67 | |||
| 68 | # prepare the coefficients (array) as a `xarray.DataArray` |
||
| 69 | if isinstance(sw_coeffs, xr.DataArray): |
||
| 70 | if "offset" in coeff_ds.proxy.values: |
||
| 71 | ones = xr.ones_like(sw_coeffs.isel(proxy=0)) |
||
| 72 | ones = ones.assign_coords({"proxy": "offset"}) |
||
| 73 | sw_coeffs = xr.concat([sw_coeffs, ones], dim="proxy") |
||
| 74 | else: |
||
| 75 | sw_coeffs = np.atleast_2d(sw_coeffs) |
||
| 76 | if "offset" in coeff_ds.proxy.values: |
||
| 77 | aix = sw_coeffs.shape.index(len(coeff_ds.proxy.values) - 1) |
||
| 78 | if aix != 0: |
||
| 79 | sw_coeffs = sw_coeffs.T |
||
| 80 | sw_coeffs = np.vstack([sw_coeffs, np.ones(sw_coeffs.shape[1])]) |
||
| 81 | else: |
||
| 82 | aix = sw_coeffs.shape.index(len(coeff_ds.proxy.values)) |
||
| 83 | if aix != 0: |
||
| 84 | sw_coeffs = sw_coeffs.T |
||
| 85 | extra_dims = ["dim_{0}".format(_d) for _d in range(sw_coeffs.ndim - 1)] |
||
| 86 | sw_coeffs = xr.DataArray( |
||
| 87 | sw_coeffs, |
||
| 88 | dims=["proxy"] + extra_dims, |
||
| 89 | coords={"proxy": coeff_ds.proxy.values}, |
||
| 90 | ) |
||
| 91 | |||
| 92 | # Calculate model (mean) values from `beta` |
||
| 93 | # fill NaNs with zero for `.dot()` |
||
| 94 | coeffs = coeff_sel.beta.fillna(0.) |
||
| 95 | q = coeffs.dot(sw_coeffs) |
||
| 96 | if not return_var: |
||
| 97 | return q |
||
| 98 | |||
| 99 | # Calculate variance of the model from `beta_std` |
||
| 100 | # fill NaNs with zero for `.dot()` |
||
| 101 | coeffv = coeff_sel.beta_std.fillna(0.)**2 |
||
| 102 | q_var = coeffv.dot(sw_coeffs**2) |
||
| 103 | return q, q_var |
||
| 104 |
