st-bender /
pyeppaurora
| Conditions | 10 |
| Total Lines | 144 |
| Code Lines | 68 |
| 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:
Complex classes like eppaurora.models.ssusiq2023.ssusiq2023() often do a lot of different things. To break such a class down, we need to identify a cohesive component within that class. A common approach to find such a component is to look for fields/methods that share the same prefixes, or suffixes.
Once you have determined the fields that belong together, you can apply the Extract Class refactoring. If the component makes sense as a sub-class, Extract Subclass is also a candidate, and is often faster.
Methods with many parameters are not only hard to understand, but their parameters also often become inconsistent when you need more, or different data.
There are several approaches to avoid long parameter lists:
| 1 | # coding: utf-8 |
||
| 45 | def ssusiq2023( |
||
| 46 | gmlat, |
||
| 47 | mlt, |
||
| 48 | alt, |
||
| 49 | sw_coeffs, |
||
| 50 | coeff_ds=None, |
||
| 51 | interpolate=False, |
||
| 52 | method="linear", |
||
| 53 | return_var=False, |
||
| 54 | ): |
||
| 55 | u""" |
||
| 56 | Parameters |
||
| 57 | ---------- |
||
| 58 | gmlat: float |
||
| 59 | Geomagnetic latitude in [degrees]. |
||
| 60 | mlt: float |
||
| 61 | Magnetic local time in [hours]. |
||
| 62 | alt: float |
||
| 63 | Altitude in [km] |
||
| 64 | sw_coeffs: array_like or `xarray.DataArray` |
||
| 65 | The space weather index values to use (for the requested time(s)), |
||
| 66 | should be of shape (N, M) with N = number of proxies, currently 5: |
||
| 67 | [Kp, PC, Ap, log(f10.7_81ctr_obs), log(v_plasma)]. |
||
| 68 | The `xarray.DataArray` should have a dimension named "proxy" with |
||
| 69 | matching coordinates: |
||
| 70 | ["Kp", "PC", "Ap", "log_f107_81ctr_obs", "log_v_plasma"] |
||
| 71 | All the other dimensions will be broadcasted. |
||
| 72 | coeff_ds: `xarray.Dataset`, optional (default: None) |
||
| 73 | Dataset with the model coefficients, `None` uses the packaged version. |
||
| 74 | interpolate: bool, optional (default: False) |
||
| 75 | If `True`, uses bilinear interpolate in MLT and geomagnetic latitude, |
||
| 76 | using periodic (24h) boundary conditions in MLT. Otherwise, the closest |
||
| 77 | MLT/geomagnetic latitude bin will be selected. |
||
| 78 | method: str, optional (default: "linear") |
||
| 79 | Interpolation method to use, see `scipy.interpolate.interpn` for options. |
||
| 80 | Only used if `interpolate` is `True`. |
||
| 81 | return_var: bool, optional (default: False) |
||
| 82 | If `True`, returns the predicted variance in addition to the values, |
||
| 83 | otherwise only the mean prediction is returned. |
||
| 84 | |||
| 85 | Returns |
||
| 86 | ------- |
||
| 87 | q: `xarray.DataArray` |
||
| 88 | log(q), where q is the ionization rate in [cm⁻³ s⁻¹] |
||
| 89 | if `return_var` is False. |
||
| 90 | q, var(q): tuple of `xarray.DataArray`s |
||
| 91 | log(q) and var(log(q)) where q is the ionization rate in [cm⁻³ s⁻¹] |
||
| 92 | if `return_var` is True. |
||
| 93 | """ |
||
| 94 | coeff_ds = coeff_ds or xr.open_dataset( |
||
| 95 | COEFF_PATH, decode_times=False, engine="h5netcdf" |
||
| 96 | ) |
||
| 97 | coeff_sel = coeff_ds.sel(altitude=alt) |
||
| 98 | if interpolate: |
||
| 99 | _ds_m = coeff_sel.assign_coords(mlt=coeff_sel.mlt - 24) |
||
| 100 | _ds_p = coeff_sel.assign_coords(mlt=coeff_sel.mlt + 24) |
||
| 101 | _ds_mp = xr.concat([_ds_m, coeff_sel, _ds_p], dim="mlt") |
||
| 102 | # square the standard deviation for interpolation |
||
| 103 | _ds_mp["beta_var"] = _ds_mp["beta_std"]**2 |
||
| 104 | coeff_sel = _interp( |
||
| 105 | _ds_mp, |
||
| 106 | latitude=gmlat, mlt=mlt, |
||
| 107 | method=method, |
||
| 108 | ) |
||
| 109 | # and square root back to get the standard deviation |
||
| 110 | coeff_sel["beta_std"] = np.sqrt(coeff_sel["beta_var"]) |
||
| 111 | else: |
||
| 112 | coeff_sel = coeff_sel.sel(latitude=gmlat, mlt=mlt, method="nearest") |
||
| 113 | |||
| 114 | # Determine if `xarray` read bytes or strings to |
||
| 115 | # match the correct name in the proxy names. |
||
| 116 | # Default is plain strings. |
||
| 117 | offset = "offset" |
||
| 118 | if isinstance(coeff_sel.proxy.values[0], bytes): |
||
| 119 | offset = offset.encode() |
||
| 120 | have_offset = offset in coeff_sel.proxy.values |
||
| 121 | |||
| 122 | # prepare the coefficients (array) as a `xarray.DataArray` |
||
| 123 | if isinstance(sw_coeffs, xr.DataArray): |
||
| 124 | if have_offset: |
||
| 125 | ones = xr.ones_like(sw_coeffs.isel(proxy=0)) |
||
| 126 | ones = ones.assign_coords(proxy="offset") |
||
| 127 | sw_coeffs = xr.concat([sw_coeffs, ones], dim="proxy") |
||
| 128 | sw_coeffs = sw_coeffs.sel(proxy=coeff_sel.proxy.astype(sw_coeffs.proxy.dtype)) |
||
| 129 | else: |
||
| 130 | sw_coeffs = np.atleast_2d(sw_coeffs) |
||
| 131 | if have_offset: |
||
| 132 | aix = sw_coeffs.shape.index(len(coeff_sel.proxy.values) - 1) |
||
| 133 | if aix != 0: |
||
| 134 | warn( |
||
| 135 | "Automatically changing axis. " |
||
| 136 | "This is ambiguous, to remove the ambiguity, " |
||
| 137 | "make sure that the different indexes (proxies) " |
||
| 138 | "are ordered along the zero-th axis in multi-" |
||
| 139 | "dimensional settings. I.e. each row corresponds " |
||
| 140 | "to a different index, Kp, PC, Ap, etc." |
||
| 141 | ) |
||
| 142 | sw_coeffs = sw_coeffs.swapaxes(aix, 0) |
||
| 143 | sw_coeffs = np.vstack([sw_coeffs, np.ones(sw_coeffs.shape[1])]) |
||
| 144 | else: |
||
| 145 | aix = sw_coeffs.shape.index(len(coeff_sel.proxy.values)) |
||
| 146 | if aix != 0: |
||
| 147 | warn( |
||
| 148 | "Automatically changing axis. " |
||
| 149 | "This is ambiguous, to remove the ambiguity, " |
||
| 150 | "make sure that the different indexes (proxies) " |
||
| 151 | "are ordered along the zero-th axis in multi-" |
||
| 152 | "dimensional settings. I.e. each row corresponds " |
||
| 153 | "to a different index, Kp, PC, Ap, etc." |
||
| 154 | ) |
||
| 155 | sw_coeffs = sw_coeffs.swapaxes(aix, 0) |
||
| 156 | extra_dims = ["dim_{0}".format(_d) for _d in range(sw_coeffs.ndim - 1)] |
||
| 157 | sw_coeffs = xr.DataArray( |
||
| 158 | sw_coeffs, |
||
| 159 | dims=["proxy"] + extra_dims, |
||
| 160 | coords={"proxy": coeff_sel.proxy.values}, |
||
| 161 | ) |
||
| 162 | |||
| 163 | # Calculate model (mean) values from `beta` |
||
| 164 | # fill NaNs with zero for `.dot()` |
||
| 165 | coeffs = coeff_sel.beta.fillna(0.) |
||
| 166 | q = coeffs.dot(sw_coeffs) |
||
| 167 | q = q.rename("log_q") |
||
| 168 | q.attrs = { |
||
| 169 | "long_name": "natural logarithm of ionization rate", |
||
| 170 | "units": "log(cm-3 s-1)", |
||
| 171 | } |
||
| 172 | if not return_var: |
||
| 173 | return q |
||
| 174 | |||
| 175 | # Calculate variance of the model from `beta_std` |
||
| 176 | # fill NaNs with zero for `.dot()` |
||
| 177 | coeffv = coeff_sel.beta_std.fillna(0.)**2 |
||
| 178 | q_var = coeffv.dot(sw_coeffs**2) |
||
| 179 | if "sigma2" in coeff_sel.data_vars: |
||
| 180 | # if available, add the posterior variance |
||
| 181 | # to get the full posterior predictive variance |
||
| 182 | q_var = coeff_sel["sigma2"] + q_var |
||
| 183 | q_var = q_var.rename("var_log_q") |
||
| 184 | q_var.attrs = { |
||
| 185 | "long_name": "variance of the natural logarithm of ionization rate", |
||
| 186 | "units": "1", |
||
| 187 | } |
||
| 188 | return q, q_var |
||
| 189 |
