Conditions | 4 |
Total Lines | 181 |
Code Lines | 56 |
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:
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 -*- |
||
122 | def msise_4d( |
||
123 | time, alt, lat, lon, |
||
124 | f107a=None, f107=None, ap=None, |
||
125 | lst=None, |
||
126 | ap_a=None, flags=None, |
||
127 | method="gtd7", |
||
128 | ): |
||
129 | u"""4-D Xarray Interface to :func:`msise_flat()`. |
||
130 | |||
131 | 4-D MSIS model atmosphere as a :class:`xarray.Dataset` with dimensions |
||
132 | (time, alt, lat, lon). Only scalars and 1-D arrays for `time`, `alt`, |
||
133 | `lat`, and `lon` are supported. |
||
134 | |||
135 | The geomagnetic and Solar flux indices can be acquired using the |
||
136 | `spaceweather` package (when set to `None`, the default). |
||
137 | They can also be supplied explicitly as scalars or 1-D arrays with |
||
138 | the same shape as `time`. |
||
139 | |||
140 | The local solar time (LST) will usually be calculated from `time` and `lon`, |
||
141 | but can be set explicitly using a fixed scalar LST, a 1-D array matching |
||
142 | the shape of either `time` or `lon`, or a 2-D array matching the shape |
||
143 | of (`time`, `lon`). |
||
144 | |||
145 | Parameters |
||
146 | ---------- |
||
147 | time: `datetime.datetime`, `pandas` datetime, str, or 1-d array_like (I,) |
||
148 | Time as `datetime.datetime`s, a `pandas` datetime object, a date-time |
||
149 | string supported by `pandas.to_datetime()`, or an array of those. |
||
150 | Will be converted with `pandas.to_datetime()`. |
||
151 | alt: float or 1-d array_like (J,) |
||
152 | Altitudes in [km]. |
||
153 | lat: float or 1-d array_like (K,) |
||
154 | Latitudes in [°N]. |
||
155 | lon: float or 1-d array_like (L,) |
||
156 | Longitudes in [°E] |
||
157 | f107a: float or 1-d array_like (I,), optional |
||
158 | The 81-day running average Solar 10.7cm radio flux, |
||
159 | centred at the day(s) of `time`. |
||
160 | Set to `None` (default) to use the `spaceweather` package. |
||
161 | f107: float or 1-d array_like (I,), optional |
||
162 | The Solar 10.7cm radio flux at the previous day(s) of `time`. |
||
163 | Set to `None` (default) to use the `spaceweather` package. |
||
164 | ap: float or 1-d array_like (I,), optional |
||
165 | The daily geomagnetic Ap index at the day(s) of `time`. |
||
166 | Set to `None` (default) to use the `spaceweather` package. |
||
167 | lst: float, 1-d (I,) or (L,) or 2-d array_like (I,L), optional |
||
168 | The local solar time at `time` and `lon`, to override the |
||
169 | calculated values. |
||
170 | Default: `None` (calculate from `time` and `lon`) |
||
171 | ap_a: list of int (7,), optional |
||
172 | List of Ap indices, passed to `msise_flat()`, |
||
173 | broadcasting is currently not supported. |
||
174 | flags: list of int (23,), optional |
||
175 | List of flags, passed to `msise_flat()`, |
||
176 | broadcasting is currently not supported. |
||
177 | method: str, optional, default "gtd7" |
||
178 | Select MSISE-00 method, changes the output of "rho", |
||
179 | the atmospheric mass density. |
||
180 | |||
181 | Returns |
||
182 | ------- |
||
183 | msise_4d: :class:`xarray.Dataset` |
||
184 | The MSIS atmosphere with dimensions ("time", "alt", "lat", "lon") |
||
185 | and shape (I, J, K, L) containing the data arrays: |
||
186 | "He", "O", "N2", "O2", "Ar", "rho", "H", "N", "AnomO", "Texo", "Talt", |
||
187 | as well as the local solar times "lst" (I,L), and the |
||
188 | values used for "Ap" (I,), "f107" (I,), "f107a" (I,). |
||
189 | |||
190 | Example |
||
191 | ------- |
||
192 | Using "internal" setting of the geomagnetic and Solar indices |
||
193 | and automatic conversion to 4-D looks like this |
||
194 | ("time" can be an array of :class:`datetime.datetime` as well): |
||
195 | |||
196 | >>> from datetime import datetime |
||
197 | >>> from nrlmsise00.dataset import msise_4d |
||
198 | >>> alts = np.arange(200, 401, 100.) # = [200, 300, 400] [km] |
||
199 | >>> lats = np.arange(60, 71, 10.) # = [60, 70] [°N] |
||
200 | >>> lons = np.arange(-70., 71., 35.) # = [-70, -35, 0, 35, 70] [°E] |
||
201 | >>> # broadcasting is done internally |
||
202 | >>> ds = msise_4d(datetime(2009, 6, 21, 8, 3, 20), alts, lats, lons) |
||
203 | >>> ds # doctest: +SKIP |
||
204 | <xarray.Dataset> |
||
205 | Dimensions: (alt: 3, lat: 2, lon: 5, time: 1) |
||
206 | Coordinates: |
||
207 | * time (time) datetime64[ns] 2009-06-21T08:03:20 |
||
208 | * alt (alt) float64 200.0 300.0 400.0 |
||
209 | * lat (lat) float64 60.0 70.0 |
||
210 | * lon (lon) float64 -70.0 -35.0 0.0 35.0 70.0 |
||
211 | Data variables: |
||
212 | He (time, alt, lat, lon) float64 8.597e+05 1.063e+06 ... 4.936e+05 |
||
213 | O (time, alt, lat, lon) float64 1.248e+09 1.46e+09 ... 2.635e+07 |
||
214 | N2 (time, alt, lat, lon) float64 2.555e+09 2.654e+09 ... 1.667e+06 |
||
215 | O2 (time, alt, lat, lon) float64 2.1e+08 2.062e+08 ... 3.471e+04 |
||
216 | Ar (time, alt, lat, lon) float64 3.16e+06 3.287e+06 ... 76.55 67.16 |
||
217 | rho (time, alt, lat, lon) float64 1.635e-13 1.736e-13 ... 7.984e-16 |
||
218 | H (time, alt, lat, lon) float64 3.144e+05 3.02e+05 ... 1.237e+05 |
||
219 | N (time, alt, lat, lon) float64 9.095e+06 1.069e+07 ... 6.765e+05 |
||
220 | AnomO (time, alt, lat, lon) float64 1.173e-08 1.173e-08 ... 1.101e+04 |
||
221 | Texo (time, alt, lat, lon) float64 805.2 823.7 807.1 ... 818.7 821.2 |
||
222 | Talt (time, alt, lat, lon) float64 757.9 758.7 766.4 ... 818.7 821.1 |
||
223 | lst (time, lon) float64 3.389 5.722 8.056 10.39 12.72 |
||
224 | Ap (time) int32 6 |
||
225 | f107 (time) float64 66.7 |
||
226 | f107a (time) float64 69.0 |
||
227 | |||
228 | See also |
||
229 | -------- |
||
230 | msise_flat |
||
231 | """ |
||
232 | |||
233 | time = _check_nd(time) |
||
234 | alt = _check_nd(alt) |
||
235 | lat = _check_nd(lat) |
||
236 | lon = _check_nd(lon) |
||
237 | |||
238 | sw = sw_daily() |
||
239 | if sw.index.tz is None: |
||
240 | sw = sw.tz_localize("utc") |
||
241 | # convert arbitrary shapes |
||
242 | dts = pd.to_datetime(time, utc=True) |
||
243 | # convert to numpy array for further processing |
||
244 | dtsv = dts.to_numpy() |
||
245 | # previous day for f10.7 |
||
246 | dtps = dtsv - pd.to_timedelta("1d") |
||
247 | |||
248 | ap = _check_gm(ap, dtsv, df=sw[["Apavg"]]) |
||
249 | f107 = _check_gm(f107, dtps, df=sw[["f107_obs"]]) |
||
250 | f107a = _check_gm(f107a, dtsv, df=sw[["f107_81ctr_obs"]]) |
||
251 | |||
252 | # expand dimensions to 4d |
||
253 | ts = dtsv[:, None, None, None] |
||
254 | alts = alt[None, :, None, None] |
||
255 | lats = lat[None, None, :, None] |
||
256 | lons = lon[None, None, None, :] |
||
257 | |||
258 | aps = ap[:, None, None, None] |
||
259 | f107s = f107[:, None, None, None] |
||
260 | f107as = f107a[:, None, None, None] |
||
261 | |||
262 | if lst is not None: |
||
263 | lsts = _check_lst(lst, time, lon) |
||
264 | # used for MSIS below |
||
265 | lst = lsts[:, None, None, :] |
||
266 | else: |
||
267 | # calculated for the returned dataset |
||
268 | lsts = np.array([ |
||
269 | t.hour + t.minute / 60. + t.second / 3600. + lon / 15. |
||
270 | for t in dts |
||
271 | ]) |
||
272 | |||
273 | msis_data = msise_flat( |
||
274 | ts, alts, lats, lons, |
||
275 | f107as, f107s, aps, |
||
276 | lst=lst, ap_a=ap_a, flags=flags, method=method, |
||
277 | ) |
||
278 | ret = xr.Dataset( |
||
279 | OrderedDict([( |
||
280 | m[0], ( |
||
281 | ["time", "alt", "lat", "lon"], |
||
282 | d, |
||
283 | {"long_name": m[1], "units": m[2]} |
||
284 | )) |
||
285 | for m, d in zip(MSIS_OUTPUT, np.rollaxis(msis_data, -1)) |
||
286 | ]), |
||
287 | coords=OrderedDict([ |
||
288 | ("time", dts.tz_localize(None)), |
||
289 | ("alt", ("alt", alt, {"long_name": "altitude", "units": "km"})), |
||
290 | ("lat", ("lat", lat, {"long_name": "latitude", "units": "degrees_north"})), |
||
291 | ("lon", ("lon", lon, {"long_name": "longitude", "units": "degrees_east"})), |
||
292 | ]), |
||
293 | ) |
||
294 | ret["lst"] = ( |
||
295 | ["time", "lon"], lsts, {"long_name": "Mean Local Solar Time", "units": "h"} |
||
296 | ) |
||
297 | ret["Ap"] = (["time"], ap) |
||
298 | ret["f107"] = (["time"], f107) |
||
299 | ret["f107a"] = (["time"], f107a) |
||
300 | for _sw in SW_INDICES: |
||
301 | ret[_sw[0]].attrs.update({"long_name": _sw[1], "units": _sw[2]}) |
||
302 | return ret |
||
303 |