eppaurora.models.ssusiq2023.ssusiq2023()

Complexity

Conditions

10

Size

Total Lines	142
Code Lines	67

Duplication

Lines	0
Ratio	0 %

Importance

Changes

0

# coding: utf-8
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