nrlmsise00.dataset.core._check_lst()   B
last analyzed

Complexity

Conditions 7

Size

Total Lines 30
Code Lines 19

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
eloc 19
dl 0
loc 30
rs 8
c 0
b 0
f 0
cc 7
nop 3
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