1
|
|
|
# -*- coding: utf-8 -*- |
2
|
|
|
# Copyright (c) 2020 Stefan Bender |
3
|
|
|
# |
4
|
|
|
# This file is part of pynrlmsise00. |
5
|
|
|
# pynrlmsise00 is free software: you can redistribute it or modify |
6
|
|
|
# it under the terms of the GNU General Public License as published |
7
|
|
|
# by the Free Software Foundation, version 2. |
8
|
|
|
# See accompanying LICENSE file or http://www.gnu.org/licenses/gpl-2.0.html. |
9
|
|
|
"""Python 4-D `xarray.dataset` interface to the NRLMSISE-00 model |
10
|
|
|
|
11
|
|
|
""" |
12
|
|
|
from collections import OrderedDict |
13
|
|
|
|
14
|
|
|
import numpy as np |
15
|
|
|
import pandas as pd |
16
|
|
|
import xarray as xr |
17
|
|
|
|
18
|
|
|
from spaceweather import sw_daily |
19
|
|
|
|
20
|
|
|
from ..core import msise_flat |
21
|
|
|
|
22
|
|
|
__all__ = ["msise_4d"] |
23
|
|
|
|
24
|
|
|
MSIS_OUTPUT = [ |
25
|
|
|
# name, long name, units |
26
|
|
|
# number densities |
27
|
|
|
("He", "He number density", "cm^-3"), |
28
|
|
|
("O", "O number density", "cm^-3"), |
29
|
|
|
("N2", "N2 number density", "cm^-3"), |
30
|
|
|
("O2", "O2 number density", "cm^-3"), |
31
|
|
|
("Ar", "AR number density", "cm^-3"), |
32
|
|
|
("rho", "total mass density", "g cm^-3"), # includes d[8] in gtd7d |
33
|
|
|
("H", "H number density", "cm^-3"), |
34
|
|
|
("N", "N number density", "cm^-3"), |
35
|
|
|
("AnomO", "Anomalous oxygen number density", "cm^-3"), |
36
|
|
|
# temperature |
37
|
|
|
("Texo", "Exospheric temperature", "K"), |
38
|
|
|
("Talt", "Temperature at alt", "K"), |
39
|
|
|
] |
40
|
|
|
|
41
|
|
|
SW_INDICES = [ |
42
|
|
|
# name, long name, units |
43
|
|
|
("Ap", "Daily Ap index", "nT"), |
44
|
|
|
( |
45
|
|
|
"f107", |
46
|
|
|
"Observed solar f10.7 cm radio flux of the previous day", |
47
|
|
|
"sfu, 10^-22 W m^-2 Hz^-1" |
48
|
|
|
), |
49
|
|
|
( |
50
|
|
|
"f107a", |
51
|
|
|
"Observed 81-day running average of the solar f10.7 cm radio flux centred on day", |
52
|
|
|
"sfu, 10^-22 W m^-2 Hz^-1" |
53
|
|
|
), |
54
|
|
|
] |
55
|
|
|
|
56
|
|
|
|
57
|
|
|
# helper functions |
58
|
|
|
def _check_nd(a, ndim=1): |
59
|
|
|
a = np.atleast_1d(a) |
60
|
|
|
if a.ndim > ndim: |
61
|
|
|
raise ValueError( |
62
|
|
|
"Only scalars and up to {ndim}-D arrays are currently supported, " |
63
|
|
|
"got {a.ndim} dimensional array.".format(ndim=ndim, a=a) |
64
|
|
|
) |
65
|
|
|
return a |
66
|
|
|
|
67
|
|
|
|
68
|
|
|
def _check_gm(gm, dts, df=None): |
69
|
|
|
"""Check that GM indices have the correct shape |
70
|
|
|
|
71
|
|
|
Returns the GM index broadcasted to shape (`dts`,). |
72
|
|
|
""" |
73
|
|
|
def _dfloc(t, df=None): |
74
|
|
|
return df.loc[t.floor("D")] |
75
|
|
|
|
76
|
|
|
if gm is None and df is not None: |
77
|
|
|
# select arbitrary shapes from pandas dataframes |
78
|
|
|
return np.vectorize(_dfloc, excluded=["df"])(dts, df=df) |
79
|
|
|
gm = np.atleast_1d(gm) |
80
|
|
|
if gm.ndim > 1: |
81
|
|
|
raise ValueError( |
82
|
|
|
"GM coefficients must be either scalars or one-dimensional arrays, " |
83
|
|
|
"other shapes are currently not supported." |
84
|
|
|
) |
85
|
|
|
if gm.shape != dts.shape: |
86
|
|
|
gm = np.broadcast_to(gm, dts.shape) |
87
|
|
|
return gm |
88
|
|
|
|
89
|
|
|
|
90
|
|
|
def _check_lst(lst, time, lon): |
91
|
|
|
"""Check that LST has the correct shape |
92
|
|
|
|
93
|
|
|
Returns the LST broadcasted to shape (`time`, `lon`). |
94
|
|
|
""" |
95
|
|
|
lst = np.atleast_1d(lst) |
96
|
|
|
if lst.ndim > 2: |
97
|
|
|
raise ValueError("Only up to 2-D arrays are supported for LST.") |
98
|
|
|
if lst.ndim == 2: |
99
|
|
|
if ( |
100
|
|
|
lst.shape[0] in [1, lon.size] |
101
|
|
|
and lst.shape[1] in [1, time.size] |
102
|
|
|
): |
103
|
|
|
# got shape (lon, time), (1, time), (lon, 1) or (1, 1) |
104
|
|
|
lst = lst.T |
105
|
|
|
else: |
106
|
|
|
# require shape (time, lon), (1, lon), (time, 1), or (1, 1) |
107
|
|
|
assert ( |
108
|
|
|
lst.shape[0] in [1, time.size] |
109
|
|
|
and lst.shape[1] in [1, lon.size] |
110
|
|
|
) |
111
|
|
|
elif lst.ndim == 1: |
112
|
|
|
if lst.size == lon.size: |
113
|
|
|
# longitude takes precedence |
114
|
|
|
lst = lst[None, :] |
115
|
|
|
else: |
116
|
|
|
assert lst.size in [1, time.size] |
117
|
|
|
lst = lst[:, None] |
118
|
|
|
lsts = np.broadcast_to(lst, (time.size, lon.size)) |
119
|
|
|
return lsts |
120
|
|
|
|
121
|
|
|
|
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
|
|
|
|