eppaurora.models.ssusiq2023.ssusiq2023_coeffs() - Code Metrics - Inspection of "ci: Update ubuntu package list before python setup" - st-bender/pyeppaurora - Measure and Improve Code Quality continuously with Scrutinizer

Passed

Push — master ( c8f3eb...6aa53e )

by Stefan

created 2023-10-31 14:02 UTC

eppaurora.models.ssusiq2023.ssusiq2023_coeffs() A

↳ Parent: eppaurora.models.ssusiq2023

Complexity

Conditions

Size

Total Lines	13
Code Lines	3

Duplication

Lines	0
Ratio	0 %

Importance

Changes

Metric	Value
cc	1
eloc	3
nop	0
dl	0
loc	13
rs	10
c	0
b	0
f	0

# coding: utf-8
# Copyright (c) 2023 Stefan Bender
#
# This file is part of pyeppaurora.
# pyeppaurora 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.
"""Empirical model for auroral ionization rates

Implements the empirical model for auroral ionization,
derived from SSUSI UV observations [1]_.

.. [1] Bender et al., in prep., 2023
"""
from logging import warning as warn
from os import path
from pkg_resources import resource_filename

import numpy as np
import xarray as xr

__all__ = [
	"ssusiq2023",
]

COEFF_FILE = "SSUSI_IRgrid_coeffs_f17f18.nc"
COEFF_PATH = resource_filename(__name__, path.join("data", COEFF_FILE))


def _interp(ds, method="linear", method_non_numeric="nearest", **kwargs):
	"""Fix `xarray` interpolation with non-numeric variables
	"""
	v_n = sorted(
		filter(lambda _v: np.issubdtype(ds[_v].dtype, np.number), ds)
	)
	v_nn = sorted(set(ds) - set(v_n))
	ds_n = ds[v_n].interp(method=method, **kwargs)
	ds_nn = ds[v_nn].sel(method=method_non_numeric, **kwargs)
	# override coordinates for `merge()`
	ds_nn = ds_nn.assign_coords(**ds_n.coords)
	return xr.merge([ds_n, ds_nn], join="left")


def ssusiq2023_coeffs():
	"""SSUSI ionization rate model coefficients

	Returns the fitted ionization rate model coefficents as
	read from the coefficient netcdf file.

	Returns
	-------
	coeffs: `xarray.Dataset`
		The default fitted model coefficients as read from the file.
	"""
	return xr.open_dataset(
		COEFF_PATH, decode_times=False, engine="h5netcdf"
	)


def ssusiq2023(
	gmlat,
	mlt,
	alt,
	sw_coeffs,
	coeff_ds=None,
	interpolate=False,
	method="linear",
	return_var=False,
):
	u"""
	Parameters
	----------
	gmlat: float
		Geomagnetic latitude in [degrees].
	mlt: float
		Magnetic local time in [hours].
	alt: float
		Altitude in [km]
	sw_coeffs: array_like or `xarray.DataArray`
		The space weather index values to use (for the requested time(s)),
		should be of shape (N, M) with N = number of proxies, currently 4:
		[Kp, PC, Ap, log(f10.7_81ctr_obs)].
		The `xarray.DataArray` should have a dimension named "proxy" with
		matching coordinates:
		["Kp", "PC", "Ap", "log_f107_81ctr_obs"]
		All the other dimensions will be broadcasted.
	coeff_ds: `xarray.Dataset`, optional (default: None)
		Dataset with the model coefficients, `None` uses the packaged version.
	interpolate: bool, optional (default: False)
		If `True`, uses bilinear interpolate in MLT and geomagnetic latitude,
		using periodic (24h) boundary conditions in MLT. Otherwise, the closest
		MLT/geomagnetic latitude bin will be selected.
	method: str, optional (default: "linear")
		Interpolation method to use, see `scipy.interpolate.interpn` for options.
		Only used if `interpolate` is `True`.
	return_var: bool, optional (default: False)
		If `True`, returns the predicted variance in addition to the values,
		otherwise only the mean prediction is returned.

	Returns
	-------
	q: `xarray.DataArray`
		log(q), where q is the ionization rate in [cm⁻³ s⁻¹]
		if `return_var` is False.
	q, var(q): tuple of `xarray.DataArray`s
		log(q) and var(log(q)) where q is the ionization rate in [cm⁻³ s⁻¹]
		if `return_var` is True.
	"""
	coeff_ds = coeff_ds or ssusiq2023_coeffs()
	coeff_sel = coeff_ds.sel(altitude=alt)
	if interpolate:
		_ds_m = coeff_sel.assign_coords(mlt=coeff_sel.mlt - 24)
		_ds_p = coeff_sel.assign_coords(mlt=coeff_sel.mlt + 24)
		_ds_mp = xr.concat([_ds_m, coeff_sel, _ds_p], dim="mlt")
		# square the standard deviation for interpolation
		_ds_mp["beta_var"] = _ds_mp["beta_std"]**2
		coeff_sel = _interp(
			_ds_mp,
			latitude=gmlat, mlt=mlt,
			method=method,
		)
		# and square root back to get the standard deviation
		coeff_sel["beta_std"] = np.sqrt(coeff_sel["beta_var"])
	else:
		coeff_sel = coeff_sel.sel(latitude=gmlat, mlt=mlt, method="nearest")

	# Determine if `xarray` read bytes or strings to
	# match the correct name in the proxy names.
	# Default is plain strings.
	offset = "offset"
	if isinstance(coeff_sel.proxy.values[0], bytes):
		offset = offset.encode()
	have_offset = offset in coeff_sel.proxy.values

	# prepare the coefficients (array) as a `xarray.DataArray`
	if isinstance(sw_coeffs, xr.DataArray):
		if have_offset:
			ones = xr.ones_like(sw_coeffs.isel(proxy=0))
			ones = ones.assign_coords(proxy="offset")
			sw_coeffs = xr.concat([sw_coeffs, ones], dim="proxy")
		sw_coeffs = sw_coeffs.sel(proxy=coeff_sel.proxy.astype(sw_coeffs.proxy.dtype))
	else:
		sw_coeffs = np.atleast_2d(sw_coeffs)
		if have_offset:
			aix = sw_coeffs.shape.index(len(coeff_sel.proxy.values) - 1)
			if aix != 0:
				warn(
					"Automatically changing axis. "
					"This is ambiguous, to remove the ambiguity, "
					"make sure that the different indexes (proxies) "
					"are ordered along the zero-th axis in multi-"
					"dimensional settings. I.e. each row corresponds "
					"to a different index, Kp, PC, Ap, etc."
				)
				sw_coeffs = sw_coeffs.swapaxes(aix, 0)
			sw_coeffs = np.vstack([sw_coeffs, np.ones(sw_coeffs.shape[1])])
		else:
			aix = sw_coeffs.shape.index(len(coeff_sel.proxy.values))
			if aix != 0:
				warn(
					"Automatically changing axis. "
					"This is ambiguous, to remove the ambiguity, "
					"make sure that the different indexes (proxies) "
					"are ordered along the zero-th axis in multi-"
					"dimensional settings. I.e. each row corresponds "
					"to a different index, Kp, PC, Ap, etc."
				)
				sw_coeffs = sw_coeffs.swapaxes(aix, 0)
		extra_dims = ["dim_{0}".format(_d) for _d in range(sw_coeffs.ndim - 1)]
		sw_coeffs = xr.DataArray(
			sw_coeffs,
			dims=["proxy"] + extra_dims,
			coords={"proxy": coeff_sel.proxy.values},
		)

	# Calculate model (mean) values from `beta`
	# fill NaNs with zero for `.dot()`
	coeffs = coeff_sel.beta.fillna(0.)
	q = coeffs.dot(sw_coeffs)
	q = q.rename("log_q")
	q.attrs = {
		"long_name": "natural logarithm of ionization rate",
		"units": "log(cm-3 s-1)",
	}
	if not return_var:
		return q

	# Calculate variance of the model from `beta_std`
	# fill NaNs with zero for `.dot()`
	coeffv = coeff_sel.beta_std.fillna(0.)**2
	q_var = coeffv.dot(sw_coeffs**2)
	if "sigma2" in coeff_sel.data_vars:
		# if available, add the posterior variance
		# to get the full posterior predictive variance
		q_var = coeff_sel["sigma2"] + q_var
	q_var = q_var.rename("var_log_q")
	q_var.attrs = {
		"long_name": "variance of the natural logarithm of ionization rate",
		"units": "1",
	}
	return q, q_var


1			# coding: utf-8
2			# Copyright (c) 2023 Stefan Bender
3			#
4			# This file is part of pyeppaurora.
5			# pyeppaurora 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			"""Empirical model for auroral ionization rates
10
11			Implements the empirical model for auroral ionization,
12			derived from SSUSI UV observations [1]_.
13
14			.. [1] Bender et al., in prep., 2023
15			"""
16			from logging import warning as warn
17			from os import path
18			from pkg_resources import resource_filename
19
20			import numpy as np
21			import xarray as xr
22
23			__all__ = [
24			"ssusiq2023",
25			]
26
27			COEFF_FILE = "SSUSI_IRgrid_coeffs_f17f18.nc"
28			COEFF_PATH = resource_filename(__name__, path.join("data", COEFF_FILE))
29
30
31			def _interp(ds, method="linear", method_non_numeric="nearest", **kwargs):
32			"""Fix `xarray` interpolation with non-numeric variables
33			"""
34			v_n = sorted(
35			filter(lambda _v: np.issubdtype(ds[_v].dtype, np.number), ds)
36			)
37			v_nn = sorted(set(ds) - set(v_n))
38			ds_n = ds[v_n].interp(method=method, **kwargs)
39			ds_nn = ds[v_nn].sel(method=method_non_numeric, **kwargs)
40			# override coordinates for `merge()`
41			ds_nn = ds_nn.assign_coords(**ds_n.coords)
42			return xr.merge([ds_n, ds_nn], join="left")
43
44
45			def ssusiq2023_coeffs():
46			"""SSUSI ionization rate model coefficients
47
48			Returns the fitted ionization rate model coefficents as
49			read from the coefficient netcdf file.
50
51			Returns
52			-------
53			coeffs: `xarray.Dataset`
54			The default fitted model coefficients as read from the file.
55			"""
56			return xr.open_dataset(
57			COEFF_PATH, decode_times=False, engine="h5netcdf"
58			)
59
60
61			def ssusiq2023(
62			gmlat,
63			mlt,
64			alt,
65			sw_coeffs,
66			coeff_ds=None,
67			interpolate=False,
68			method="linear",
69			return_var=False,
70			):
71			u"""
72			Parameters
73			----------
74			gmlat: float
75			Geomagnetic latitude in [degrees].
76			mlt: float
77			Magnetic local time in [hours].
78			alt: float
79			Altitude in [km]
80			sw_coeffs: array_like or `xarray.DataArray`
81			The space weather index values to use (for the requested time(s)),
82			should be of shape (N, M) with N = number of proxies, currently 4:
83			[Kp, PC, Ap, log(f10.7_81ctr_obs)].
84			The `xarray.DataArray` should have a dimension named "proxy" with
85			matching coordinates:
86			["Kp", "PC", "Ap", "log_f107_81ctr_obs"]
87			All the other dimensions will be broadcasted.
88			coeff_ds: `xarray.Dataset`, optional (default: None)
89			Dataset with the model coefficients, `None` uses the packaged version.
90			interpolate: bool, optional (default: False)
91			If `True`, uses bilinear interpolate in MLT and geomagnetic latitude,
92			using periodic (24h) boundary conditions in MLT. Otherwise, the closest
93			MLT/geomagnetic latitude bin will be selected.
94			method: str, optional (default: "linear")
95			Interpolation method to use, see `scipy.interpolate.interpn` for options.
96			Only used if `interpolate` is `True`.
97			return_var: bool, optional (default: False)
98			If `True`, returns the predicted variance in addition to the values,
99			otherwise only the mean prediction is returned.
100
101			Returns
102			-------
103			q: `xarray.DataArray`
104			log(q), where q is the ionization rate in [cm⁻³ s⁻¹]
105			if `return_var` is False.
106			q, var(q): tuple of `xarray.DataArray`s
107			log(q) and var(log(q)) where q is the ionization rate in [cm⁻³ s⁻¹]
108			if `return_var` is True.
109			"""
110			coeff_ds = coeff_ds or ssusiq2023_coeffs()
111			coeff_sel = coeff_ds.sel(altitude=alt)
112			if interpolate:
113			_ds_m = coeff_sel.assign_coords(mlt=coeff_sel.mlt - 24)
114			_ds_p = coeff_sel.assign_coords(mlt=coeff_sel.mlt + 24)
115			_ds_mp = xr.concat([_ds_m, coeff_sel, _ds_p], dim="mlt")
116			# square the standard deviation for interpolation
117			_ds_mp["beta_var"] = _ds_mp["beta_std"]**2
118			coeff_sel = _interp(
119			_ds_mp,
120			latitude=gmlat, mlt=mlt,
121			method=method,
122			)
123			# and square root back to get the standard deviation
124			coeff_sel["beta_std"] = np.sqrt(coeff_sel["beta_var"])
125			else:
126			coeff_sel = coeff_sel.sel(latitude=gmlat, mlt=mlt, method="nearest")
127
128			# Determine if `xarray` read bytes or strings to
129			# match the correct name in the proxy names.
130			# Default is plain strings.
131			offset = "offset"
132			if isinstance(coeff_sel.proxy.values[0], bytes):
133			offset = offset.encode()
134			have_offset = offset in coeff_sel.proxy.values
135
136			# prepare the coefficients (array) as a `xarray.DataArray`
137			if isinstance(sw_coeffs, xr.DataArray):
138			if have_offset:
139			ones = xr.ones_like(sw_coeffs.isel(proxy=0))
140			ones = ones.assign_coords(proxy="offset")
141			sw_coeffs = xr.concat([sw_coeffs, ones], dim="proxy")
142			sw_coeffs = sw_coeffs.sel(proxy=coeff_sel.proxy.astype(sw_coeffs.proxy.dtype))
143			else:
144			sw_coeffs = np.atleast_2d(sw_coeffs)
145			if have_offset:
146			aix = sw_coeffs.shape.index(len(coeff_sel.proxy.values) - 1)
147			if aix != 0:
148			warn(
149			"Automatically changing axis. "
150			"This is ambiguous, to remove the ambiguity, "
151			"make sure that the different indexes (proxies) "
152			"are ordered along the zero-th axis in multi-"
153			"dimensional settings. I.e. each row corresponds "
154			"to a different index, Kp, PC, Ap, etc."
155			)
156			sw_coeffs = sw_coeffs.swapaxes(aix, 0)
157			sw_coeffs = np.vstack([sw_coeffs, np.ones(sw_coeffs.shape[1])])
158			else:
159			aix = sw_coeffs.shape.index(len(coeff_sel.proxy.values))
160			if aix != 0:
161			warn(
162			"Automatically changing axis. "
163			"This is ambiguous, to remove the ambiguity, "
164			"make sure that the different indexes (proxies) "
165			"are ordered along the zero-th axis in multi-"
166			"dimensional settings. I.e. each row corresponds "
167			"to a different index, Kp, PC, Ap, etc."
168			)
169			sw_coeffs = sw_coeffs.swapaxes(aix, 0)
170			extra_dims = ["dim_{0}".format(_d) for _d in range(sw_coeffs.ndim - 1)]
171			sw_coeffs = xr.DataArray(
172			sw_coeffs,
173			dims=["proxy"] + extra_dims,
174			coords={"proxy": coeff_sel.proxy.values},
175			)
176
177			# Calculate model (mean) values from `beta`
178			# fill NaNs with zero for `.dot()`
179			coeffs = coeff_sel.beta.fillna(0.)
180			q = coeffs.dot(sw_coeffs)
181			q = q.rename("log_q")
182			q.attrs = {
183			"long_name": "natural logarithm of ionization rate",
184			"units": "log(cm-3 s-1)",
185			}
186			if not return_var:
187			return q
188
189			# Calculate variance of the model from `beta_std`
190			# fill NaNs with zero for `.dot()`
191			coeffv = coeff_sel.beta_std.fillna(0.)**2
192			q_var = coeffv.dot(sw_coeffs**2)
193			if "sigma2" in coeff_sel.data_vars:
194			# if available, add the posterior variance
195			# to get the full posterior predictive variance
196			q_var = coeff_sel["sigma2"] + q_var
197			q_var = q_var.rename("var_log_q")
198			q_var.attrs = {
199			"long_name": "variance of the natural logarithm of ionization rate",
200			"units": "1",
201			}
202			return q, q_var
203

st-bender / pyeppaurora

Push — master ( c8f3eb...6aa53e )

eppaurora.models.ssusiq2023.ssusiq2023_coeffs() A

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like