nrlmsise00.dataset.core._check_gm() - Code Metrics - Inspection of "dataset: Test `nd` exception" - st-bender/pynrlmsise00 - Measure and Improve Code Quality continuously with Scrutinizer

Passed

Push — master ( 76dfc4...822cea )

by Stefan

created 2020-07-03 09:53 UTC

nrlmsise00.dataset.core._check_gm() A

↳ Parent: nrlmsise00.dataset.core

Complexity

Conditions

Size

Total Lines	20
Code Lines	12

Duplication

Lines	20
Ratio	100 %

Importance

Changes

Metric	Value
eloc	12
dl	20
loc	20
rs	9.3333
c	0
b	0
f	0
cc	5
nop	3

# Copyright (c) 2020 Stefan Bender
#
# This file is part of pynrlmsise00.
# pynrlmsise00 is free software: you can redistribute it or modify
# it under the terms of the GNU General Public License as published
# by the Free Software Foundation, version 2.
# See accompanying LICENSE file or http://www.gnu.org/licenses/gpl-2.0.html.
"""Python 4-D `xarray.dataset` interface to the NRLMSISE-00 model

"""
import numpy as np
import pandas as pd
import xarray as xr

from spaceweather import sw_daily

from ..core import msise_flat

__all__ = ["msise_4d"]

MSIS_OUTPUT = [
	# name, long name, units
	# number densities
	("He", "He number density", "cm^-3"),
	("O", "O number density", "cm^-3"),
	("N2", "N2 number density", "cm^-3"),
	("O2", "O2 number density", "cm^-3"),
	("Ar", "AR number density", "cm^-3"),
	("rho", "total mass density", "g cm^-3"),  # includes d[8] in gtd7d
	("H", "H number density", "cm^-3"),
	("N", "N number density", "cm^-3"),
	("AnomO", "Anomalous oxygen number density", "cm^-3"),
	# temperature
	("Texo", "Exospheric temperature", "K"),
	("Talt", "Temperature at alt", "K"),
]

SW_INDICES = [
	# name, long name, units
	("Ap", "Daily Ap index", "nT"),
	(
		"f107",
		"Observed solar f10.7 cm radio flux of the previous day",
		"sfu, 10^-22 W m^-2 Hz^-1"
	),
	(
		"f107a",
		"Observed 81-day running average of the solar f10.7 cm radio flux centred on day",
		"sfu, 10^-22 W m^-2 Hz^-1"
	),
]


# helper functions
def _check_nd(a, ndim=1):
	a = np.atleast_1d(a)
	if a.ndim > ndim:
		raise ValueError(
			"Only scalars and up to {ndim}-D arrays are currently supported, "
			"got {a.ndim} dimensional array.".format(ndim=ndim, a=a)
		)
	return a


def _check_gm(gm, dts, df=None):

	"""Check that GM indices have the correct shape

	Returns the GM index broadcasted to shape (`dts`,).
	"""
	def _dfloc(t, df=None):
		return df.loc[t.floor("D")]

	if gm is None and df is not None:
		# select arbitrary shapes from pandas dataframes
		return np.vectorize(_dfloc, excluded=["df"])(dts, df=df)
	gm = np.atleast_1d(gm)
	if gm.ndim > 1:
		raise ValueError(
			"GM coefficients must be either scalars or one-dimensional arrays, "
			"other shapes are currently not supported."
		)
	if gm.shape != dts.shape:
		gm = np.broadcast_to(gm, dts.shape)
	return gm


def _check_lst(lst, time, lon):

	"""Check that LST has the correct shape

	Returns the LST broadcasted to shape (`time`, `lon`).
	"""
	lst = np.atleast_1d(lst)
	if lst.ndim > 2:
		raise ValueError("Only up to 2-D arrays are supported for LST.")
	if lst.ndim == 2:
		if (
			lst.shape[0] in [1, lon.size]
			and lst.shape[1] in [1, time.size]
		):
			# got shape (lon, time), (1, time), (lon, 1) or (1, 1)
			lst = lst.T
		else:
			# require shape (time, lon), (1, lon), (time, 1), or (1, 1)
			assert (
				lst.shape[0] in [1, time.size]
				and lst.shape[1] in [1, lon.size]
			)
	elif lst.ndim == 1:
		if lst.size == lon.size:
			# longitude takes precedence
			lst = lst[None, :]
		else:
			assert lst.size in [1, time.size]
			lst = lst[:, None]
	lsts = np.broadcast_to(lst, (time.size, lon.size))
	return lsts


def msise_4d(

	time, alt, lat, lon,
	f107a=None, f107=None, ap=None,
	lst=None,
	ap_a=None, flags=None,
	method="gtd7",
):
	"""4-D Xarray Interface to :func:`msise_flat()`.

	4-D MSIS model atmosphere as a `xarray.Dataset` with dimensions
	(time, alt, lat, lon). Only scalars and 1-D arrays for `time`, `alt`,
	`lat`, and `lon` are supported.

	The geomagnetic and Solar flux indices can be acquired using the
	`spaceweather` package (when set to `None`, the default).
	They can also be supplied explicitly as scalars or 1-D arrays with
	the same shape as `time`.

	The local solar time (LST) will usually be calculated from `time` and `lon`,
	but can be set explicitly using a fixed scalar LST, a 1-D array matching
	the shape of either `time` or `lon`, or a 2-D array matching the shape
	of (`time`, `lon`).

	Parameters
	----------
	time: float or 1-d array_like (I,)
		Time as `datetime.datetime`s.
	alt: float or 1-d array_like (J,)
		Altitudes in [km].
	lat: float or 1-d array_like (K,)
		Latitudes in [°N].
	lon: float or 1-d array_like (L,)
		Longitudes in [°E]
	f107a: float or 1-d array_like (I,), optional
		The 81-day running average Solar 10.7cm radio flux,
		centred at the day(s) of `time`.
		Set to `None` (default) to use the `spaceweather` package.
	f107: float or 1-d array_like (I,), optional
		The Solar 10.7cm radio flux at the previous day(s) of `time`.
		Set to `None` (default) to use the `spaceweather` package.
	ap: float or 1-d array_like (I,), optional
		The daily geomagnetic Ap index at the day(s) of `time`.
		Set to `None` (default) to use the `spaceweather` package.
	lst: float, 1-d (I,) or (L,) or 2-d array_like (I,L), optional
		The local solar time at `time` and `lon`, to override the
		calculated values.
		Default: `None` (calculate from `time` and `lon`)
	ap_a: list of int (7,), optional
		List of Ap indices, passed to `msise_flat()`,
		broadcasting is currently not supported.
	flags: list of int (23,), optional
		List of flags, passed to `msise_flat()`,
		broadcasting is currently not supported.
	method: str, optional, default "gtd7"
		Select MSISE-00 method, changes the output of "rho",
		the atmospheric mass density.

	Returns
	-------
	msise_4d: `xarray.Dataset`
		The MSIS atmosphere with dimensions ("time", "alt", "lat", "lon")
		and shape (I, J, K, L) containing the data arrays:
		"He", "O", "N2", "O2", "Ar", "rho", "H", "N", "AnomO", "Texo", "Talt",
		as well as the local solar times "lst" (I,L), and the
		values used for "Ap" (I,), "f107" (I,), "f107a" (I,).

	Example
	-------
	Using "internal" setting of the geomagnetic and Solar indices
	and automatic conversion to 4-D looks like this
	("time" can be an array of `datetime.datetime`s as well):

	>>> from datetime import datetime
	>>> from nrlmsise00.dataset import msise_4d
	>>> alts = np.arange(200, 401, 100.)  # = [200, 300, 400] [km]
	>>> lats = np.arange(60, 71, 10.)  # = [60, 70] [°N]
	>>> lons = np.arange(-70., 71., 35.)  # = [-70, -35,  0, 35, 70] [°E]
	>>> # broadcasting is done internally
	>>> ds = msise_4d(datetime(2009, 6, 21, 8, 3, 20), alts, lats, lons)
	>>> ds
	<xarray.Dataset>
	Dimensions:  (alt: 3, lat: 2, lon: 5, time: 1)
	Coordinates:
	  * time     (time) datetime64[ns] 2009-06-21T08:03:20
	  * alt      (alt) float64 200.0 300.0 400.0
	  * lat      (lat) float64 60.0 70.0
	  * lon      (lon) float64 -70.0 -35.0 0.0 35.0 70.0
	Data variables:
	    He       (time, alt, lat, lon) float64 8.597e+05 1.063e+06 ... 4.936e+05
	    O        (time, alt, lat, lon) float64 1.248e+09 1.46e+09 ... 2.635e+07
	    N2       (time, alt, lat, lon) float64 2.555e+09 2.654e+09 ... 1.667e+06
	    O2       (time, alt, lat, lon) float64 2.1e+08 2.062e+08 ... 3.471e+04
	    Ar       (time, alt, lat, lon) float64 3.16e+06 3.287e+06 ... 76.55 67.16
	    rho      (time, alt, lat, lon) float64 1.635e-13 1.736e-13 ... 7.984e-16
	    H        (time, alt, lat, lon) float64 3.144e+05 3.02e+05 ... 1.237e+05
	    N        (time, alt, lat, lon) float64 9.095e+06 1.069e+07 ... 6.765e+05
	    AnomO    (time, alt, lat, lon) float64 1.173e-08 1.173e-08 ... 1.101e+04
	    Texo     (time, alt, lat, lon) float64 805.2 823.7 807.1 ... 818.7 821.2
	    Talt     (time, alt, lat, lon) float64 757.9 758.7 766.4 ... 818.7 821.1
	    lst      (time, lon) float64 3.389 5.722 8.056 10.39 12.72
	    Ap       (time) int32 6
	    f107     (time) float64 66.7
	    f107a    (time) float64 69.0

	See also
	--------
	msise_flat
	"""

	time = _check_nd(time)
	alt = _check_nd(alt)
	lat = _check_nd(lat)
	lon = _check_nd(lon)

	sw = sw_daily()
	# convert arbitrary shapes
	dts = np.vectorize(pd.to_datetime)(time, utc=True)
	dtps = dts - pd.to_timedelta("1d")

	ap = _check_gm(ap, dts, df=sw[["Apavg"]])
	f107 = _check_gm(f107, dtps, df=sw[["f107_obs"]])
	f107a = _check_gm(f107a, dts, df=sw[["f107_81ctr_obs"]])

	# expand dimensions to 4d
	ts = time[:, None, None, None]
	alts = alt[None, :, None, None]
	lats = lat[None, None, :, None]
	lons = lon[None, None, None, :]

	aps = ap[:, None, None, None]
	f107s = f107[:, None, None, None]
	f107as = f107a[:, None, None, None]

	if lst is not None:
		lsts = _check_lst(lst, time, lon)
		lst = lsts[:, None, None, :]
	else:
		lsts = np.array([
			t.hour + t.minute / 60. + t.second / 3600 + lon / 15.
			for t in time
		])

	msis_data = msise_flat(
		ts, alts, lats, lons,
		f107as, f107s, aps,
		lst=lst, ap_a=ap_a, flags=flags, method=method,
	)
	ret = xr.Dataset(
		{
			m[0]: (
				["time", "alt", "lat", "lon"],
				d,
				{"long_name": m[1], "units": m[2]}
			)
			for m, d in zip(MSIS_OUTPUT, np.rollaxis(msis_data, -1))
		},
		coords={"time": time, "alt": alt, "lat": lat, "lon": lon},
	)
	ret["lst"] = (
		["time", "lon"], lsts, {"long_name": "Mean Local Solar Time", "units": "h"}
	)
	ret["Ap"] = (["time"], ap)
	ret["f107"] = (["time"], f107)
	ret["f107a"] = (["time"], f107a)
	for _sw in SW_INDICES:
		ret[_sw[0]].attrs.update({"long_name": _sw[1], "units": _sw[2]})
	return ret


1		# Copyright (c) 2020 Stefan Bender
2		#
3		# This file is part of pynrlmsise00.
4		# pynrlmsise00 is free software: you can redistribute it or modify
5		# it under the terms of the GNU General Public License as published
6		# by the Free Software Foundation, version 2.
7		# See accompanying LICENSE file or http://www.gnu.org/licenses/gpl-2.0.html.
8		"""Python 4-D `xarray.dataset` interface to the NRLMSISE-00 model
9
10		"""
11		import numpy as np
12		import pandas as pd
13		import xarray as xr
14
15		from spaceweather import sw_daily
16
17		from ..core import msise_flat
18
19		__all__ = ["msise_4d"]
20
21		MSIS_OUTPUT = [
22		# name, long name, units
23		# number densities
24		("He", "He number density", "cm^-3"),
25		("O", "O number density", "cm^-3"),
26		("N2", "N2 number density", "cm^-3"),
27		("O2", "O2 number density", "cm^-3"),
28		("Ar", "AR number density", "cm^-3"),
29		("rho", "total mass density", "g cm^-3"), # includes d[8] in gtd7d
30		("H", "H number density", "cm^-3"),
31		("N", "N number density", "cm^-3"),
32		("AnomO", "Anomalous oxygen number density", "cm^-3"),
33		# temperature
34		("Texo", "Exospheric temperature", "K"),
35		("Talt", "Temperature at alt", "K"),
36		]
37
38		SW_INDICES = [
39		# name, long name, units
40		("Ap", "Daily Ap index", "nT"),
41		(
42		"f107",
43		"Observed solar f10.7 cm radio flux of the previous day",
44		"sfu, 10^-22 W m^-2 Hz^-1"
45		),
46		(
47		"f107a",
48		"Observed 81-day running average of the solar f10.7 cm radio flux centred on day",
49		"sfu, 10^-22 W m^-2 Hz^-1"
50		),
51		]
52
53
54		# helper functions
55		def _check_nd(a, ndim=1):
56		a = np.atleast_1d(a)
57		if a.ndim > ndim:
58		raise ValueError(
59		"Only scalars and up to {ndim}-D arrays are currently supported, "
60		"got {a.ndim} dimensional array.".format(ndim=ndim, a=a)
61		)
62		return a
63
64
65	View Code Duplication	def _check_gm(gm, dts, df=None):
		0 ignored issues – show Duplication introduced 2020-07-03 09:58 UTC by Report Bug Copy Issue Report This code seems to be duplicated in your project. Loading history...
66		"""Check that GM indices have the correct shape
67
68		Returns the GM index broadcasted to shape (`dts`,).
69		"""
70		def _dfloc(t, df=None):
71		return df.loc[t.floor("D")]
72
73		if gm is None and df is not None:
74		# select arbitrary shapes from pandas dataframes
75		return np.vectorize(_dfloc, excluded=["df"])(dts, df=df)
76		gm = np.atleast_1d(gm)
77		if gm.ndim > 1:
78		raise ValueError(
79		"GM coefficients must be either scalars or one-dimensional arrays, "
80		"other shapes are currently not supported."
81		)
82		if gm.shape != dts.shape:
83		gm = np.broadcast_to(gm, dts.shape)
84		return gm
85
86
87	View Code Duplication	def _check_lst(lst, time, lon):
		0 ignored issues – show Duplication introduced 2020-07-03 09:58 UTC by Report Bug Copy Issue Report This code seems to be duplicated in your project. Loading history...
88		"""Check that LST has the correct shape
89
90		Returns the LST broadcasted to shape (`time`, `lon`).
91		"""
92		lst = np.atleast_1d(lst)
93		if lst.ndim > 2:
94		raise ValueError("Only up to 2-D arrays are supported for LST.")
95		if lst.ndim == 2:
96		if (
97		lst.shape[0] in [1, lon.size]
98		and lst.shape[1] in [1, time.size]
99		):
100		# got shape (lon, time), (1, time), (lon, 1) or (1, 1)
101		lst = lst.T
102		else:
103		# require shape (time, lon), (1, lon), (time, 1), or (1, 1)
104		assert (
105		lst.shape[0] in [1, time.size]
106		and lst.shape[1] in [1, lon.size]
107		)
108		elif lst.ndim == 1:
109		if lst.size == lon.size:
110		# longitude takes precedence
111		lst = lst[None, :]
112		else:
113		assert lst.size in [1, time.size]
114		lst = lst[:, None]
115		lsts = np.broadcast_to(lst, (time.size, lon.size))
116		return lsts
117
118
119	View Code Duplication	def msise_4d(
		0 ignored issues – show Duplication introduced 2020-07-03 09:58 UTC by Report Bug Copy Issue Report This code seems to be duplicated in your project. Loading history...
120		time, alt, lat, lon,
121		f107a=None, f107=None, ap=None,
122		lst=None,
123		ap_a=None, flags=None,
124		method="gtd7",
125		):
126		"""4-D Xarray Interface to :func:`msise_flat()`.
127
128		4-D MSIS model atmosphere as a `xarray.Dataset` with dimensions
129		(time, alt, lat, lon). Only scalars and 1-D arrays for `time`, `alt`,
130		`lat`, and `lon` are supported.
131
132		The geomagnetic and Solar flux indices can be acquired using the
133		`spaceweather` package (when set to `None`, the default).
134		They can also be supplied explicitly as scalars or 1-D arrays with
135		the same shape as `time`.
136
137		The local solar time (LST) will usually be calculated from `time` and `lon`,
138		but can be set explicitly using a fixed scalar LST, a 1-D array matching
139		the shape of either `time` or `lon`, or a 2-D array matching the shape
140		of (`time`, `lon`).
141
142		Parameters
143		----------
144		time: float or 1-d array_like (I,)
145		Time as `datetime.datetime`s.
146		alt: float or 1-d array_like (J,)
147		Altitudes in [km].
148		lat: float or 1-d array_like (K,)
149		Latitudes in [°N].
150		lon: float or 1-d array_like (L,)
151		Longitudes in [°E]
152		f107a: float or 1-d array_like (I,), optional
153		The 81-day running average Solar 10.7cm radio flux,
154		centred at the day(s) of `time`.
155		Set to `None` (default) to use the `spaceweather` package.
156		f107: float or 1-d array_like (I,), optional
157		The Solar 10.7cm radio flux at the previous day(s) of `time`.
158		Set to `None` (default) to use the `spaceweather` package.
159		ap: float or 1-d array_like (I,), optional
160		The daily geomagnetic Ap index at the day(s) of `time`.
161		Set to `None` (default) to use the `spaceweather` package.
162		lst: float, 1-d (I,) or (L,) or 2-d array_like (I,L), optional
163		The local solar time at `time` and `lon`, to override the
164		calculated values.
165		Default: `None` (calculate from `time` and `lon`)
166		ap_a: list of int (7,), optional
167		List of Ap indices, passed to `msise_flat()`,
168		broadcasting is currently not supported.
169		flags: list of int (23,), optional
170		List of flags, passed to `msise_flat()`,
171		broadcasting is currently not supported.
172		method: str, optional, default "gtd7"
173		Select MSISE-00 method, changes the output of "rho",
174		the atmospheric mass density.
175
176		Returns
177		-------
178		msise_4d: `xarray.Dataset`
179		The MSIS atmosphere with dimensions ("time", "alt", "lat", "lon")
180		and shape (I, J, K, L) containing the data arrays:
181		"He", "O", "N2", "O2", "Ar", "rho", "H", "N", "AnomO", "Texo", "Talt",
182		as well as the local solar times "lst" (I,L), and the
183		values used for "Ap" (I,), "f107" (I,), "f107a" (I,).
184
185		Example
186		-------
187		Using "internal" setting of the geomagnetic and Solar indices
188		and automatic conversion to 4-D looks like this
189		("time" can be an array of `datetime.datetime`s as well):
190
191		>>> from datetime import datetime
192		>>> from nrlmsise00.dataset import msise_4d
193		>>> alts = np.arange(200, 401, 100.) # = [200, 300, 400] [km]
194		>>> lats = np.arange(60, 71, 10.) # = [60, 70] [°N]
195		>>> lons = np.arange(-70., 71., 35.) # = [-70, -35, 0, 35, 70] [°E]
196		>>> # broadcasting is done internally
197		>>> ds = msise_4d(datetime(2009, 6, 21, 8, 3, 20), alts, lats, lons)
198		>>> ds
199		<xarray.Dataset>
200		Dimensions: (alt: 3, lat: 2, lon: 5, time: 1)
201		Coordinates:
202		* time (time) datetime64[ns] 2009-06-21T08:03:20
203		* alt (alt) float64 200.0 300.0 400.0
204		* lat (lat) float64 60.0 70.0
205		* lon (lon) float64 -70.0 -35.0 0.0 35.0 70.0
206		Data variables:
207		He (time, alt, lat, lon) float64 8.597e+05 1.063e+06 ... 4.936e+05
208		O (time, alt, lat, lon) float64 1.248e+09 1.46e+09 ... 2.635e+07
209		N2 (time, alt, lat, lon) float64 2.555e+09 2.654e+09 ... 1.667e+06
210		O2 (time, alt, lat, lon) float64 2.1e+08 2.062e+08 ... 3.471e+04
211		Ar (time, alt, lat, lon) float64 3.16e+06 3.287e+06 ... 76.55 67.16
212		rho (time, alt, lat, lon) float64 1.635e-13 1.736e-13 ... 7.984e-16
213		H (time, alt, lat, lon) float64 3.144e+05 3.02e+05 ... 1.237e+05
214		N (time, alt, lat, lon) float64 9.095e+06 1.069e+07 ... 6.765e+05
215		AnomO (time, alt, lat, lon) float64 1.173e-08 1.173e-08 ... 1.101e+04
216		Texo (time, alt, lat, lon) float64 805.2 823.7 807.1 ... 818.7 821.2
217		Talt (time, alt, lat, lon) float64 757.9 758.7 766.4 ... 818.7 821.1
218		lst (time, lon) float64 3.389 5.722 8.056 10.39 12.72
219		Ap (time) int32 6
220		f107 (time) float64 66.7
221		f107a (time) float64 69.0
222
223		See also
224		--------
225		msise_flat
226		"""
227
228		time = _check_nd(time)
229		alt = _check_nd(alt)
230		lat = _check_nd(lat)
231		lon = _check_nd(lon)
232
233		sw = sw_daily()
234		# convert arbitrary shapes
235		dts = np.vectorize(pd.to_datetime)(time, utc=True)
236		dtps = dts - pd.to_timedelta("1d")
237
238		ap = _check_gm(ap, dts, df=sw[["Apavg"]])
239		f107 = _check_gm(f107, dtps, df=sw[["f107_obs"]])
240		f107a = _check_gm(f107a, dts, df=sw[["f107_81ctr_obs"]])
241
242		# expand dimensions to 4d
243		ts = time[:, None, None, None]
244		alts = alt[None, :, None, None]
245		lats = lat[None, None, :, None]
246		lons = lon[None, None, None, :]
247
248		aps = ap[:, None, None, None]
249		f107s = f107[:, None, None, None]
250		f107as = f107a[:, None, None, None]
251
252		if lst is not None:
253		lsts = _check_lst(lst, time, lon)
254		lst = lsts[:, None, None, :]
255		else:
256		lsts = np.array([
257		t.hour + t.minute / 60. + t.second / 3600 + lon / 15.
258		for t in time
259		])
260
261		msis_data = msise_flat(
262		ts, alts, lats, lons,
263		f107as, f107s, aps,
264		lst=lst, ap_a=ap_a, flags=flags, method=method,
265		)
266		ret = xr.Dataset(
267		{
268		m[0]: (
269		["time", "alt", "lat", "lon"],
270		d,
271		{"long_name": m[1], "units": m[2]}
272		)
273		for m, d in zip(MSIS_OUTPUT, np.rollaxis(msis_data, -1))
274		},
275		coords={"time": time, "alt": alt, "lat": lat, "lon": lon},
276		)
277		ret["lst"] = (
278		["time", "lon"], lsts, {"long_name": "Mean Local Solar Time", "units": "h"}
279		)
280		ret["Ap"] = (["time"], ap)
281		ret["f107"] = (["time"], f107)
282		ret["f107a"] = (["time"], f107a)
283		for _sw in SW_INDICES:
284		ret[_sw[0]].attrs.update({"long_name": _sw[1], "units": _sw[2]})
285		return ret
286

st-bender / pynrlmsise00

Push — master ( 76dfc4...822cea )

nrlmsise00.dataset.core._check_gm() A

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like