eppaurora.models.ssusiq2023.ssusiq2023()

Complexity

Conditions

8

Size

Total Lines	110
Code Lines	47

Duplication

Lines	0
Ratio	0 %

Importance

Changes

0

# coding: utf-8
def ssusiq2023(gmlat, mlt, alt, sw_coeffs, coeff_ds=None, 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 5:
		[Kp, PC, Ap, log(f10.7_81ctr_obs), log(v_plasma)].
		The `xarray.DataArray` should have a dimension named "proxy" with
		matching coordinates:
		["Kp", "PC", "Ap", "log_f107_81ctr_obs", "log_v_plasma"]
		All the other dimensions will be broadcasted.
	coeff_ds: `xarray.Dataset`, optional (default: None)
		Dataset with the model coefficients, `None` uses the packaged version.
	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 xr.open_dataset(
		COEFF_PATH, decode_times=False, engine="h5netcdf"
	)
	coeff_sel = coeff_ds.sel(
		altitude=alt, 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_ds.proxy.values[0], bytes):
		offset = offset.encode()
	have_offset = offset in coeff_ds.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")
	else:
		sw_coeffs = np.atleast_2d(sw_coeffs)
		if have_offset:
			aix = sw_coeffs.shape.index(len(coeff_ds.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_ds.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_ds.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)
	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