def baseflow_index(da: DataArray,
alpha: float = 0.98,
warmup: int = 30) -> (float, DataArray):
"""Currently just implemented for daily flows (i.e. 3 passes, see Section 2.3 Landson et al. 2013"""
# create numpy array from streamflow and add the mirrored discharge of length 'window' to the start and end
streamflow = np.zeros((da.size + 2 * warmup))
streamflow[warmup:-warmup] = da.values
streamflow[:warmup] = da.values[1:warmup + 1][::-1]
streamflow[-warmup:] = da.values[-warmup - 1:-1][::-1]
# call jit compiled function to calculate baseflow
bf_index, baseflow = _baseflow_index_jit(streamflow, alpha, warmup)
# parse baseflow as a DataArray using the coordinates of the streamflow array
da_baseflow = da.copy()
da_baseflow.data = baseflow
return bf_index, da_baseflow
def baseflow_index(da: DataArray,
alpha: float = 0.98,
warmup: int = 30,
n_passes: int = None,
datetime_coord: str = None) -> Tuple[float, DataArray]:
"""Calculate baseflow index.
Ratio of mean baseflow to mean discharge [#]_. If `da` contains NaN values, the baseflow is calculated for each
consecutive segment of more than `warmup` non-NaN values.
Parameters
----------
da : DataArray
Array of flow values.
alpha : float, optional
alpha filter parameter.
warmup : int, optional
Number of warmup steps.
n_passes : int, optional
Number of passes (alternating forward and backward) to perform. Should be an odd number. If None, will use
3 for daily and 9 for hourly data and fail for all other input frequencies.
datetime_coord : str, optional
Datetime coordinate in the passed DataArray. Tried to infer automatically if not specified. Used to infer the
frequency if `n_passes` is None.
Returns
-------
Tuple[float, DataArray]
Baseflow index and baseflow array. The baseflow array contains NaNs wherever no baseflow was
calculated due to NaNs in `da`.
Raises
------
ValueError
If `da` has a frequency other than daily or hourly and `n_passes` is None.
References
----------
.. [#] Ladson, T. R., Brown, R., Neal, B., and Nathan, R.: A Standard Approach to Baseflow Separation Using The
Lyne and Hollick Filter. Australasian Journal of Water Resources, Taylor & Francis, 2013, 17, 25--34,
doi:10.7158/13241583.2013.11465417
"""
if datetime_coord is None:
datetime_coord = utils.infer_datetime_coord(da)
if n_passes is None:
freq = utils.infer_frequency(da[datetime_coord].values)
if freq == '1D':
n_passes = 3
elif freq == '1H':
n_passes = 9
else:
raise ValueError(
f'For frequencies other than daily or hourly, n_passes must be specified.'
)
if n_passes % 2 != 1:
warnings.warn(
'n_passes should be an even number. The returned baseflow will be reversed.'
)
# call jit compiled function to calculate baseflow
bf_index, baseflow = _baseflow_index_jit(da.values, alpha, warmup,
n_passes)
# parse baseflow as a DataArray using the coordinates of the streamflow array
da_baseflow = da.copy()
da_baseflow.data = baseflow
return bf_index, da_baseflow
def test_centrifugal_asymmetry():
"""Test AsymmetryParameter.__call__ and ToroidalRotation.__call__."""
example_asymmetry = AsymmetryParameter()
t = np.linspace(75.0, 80.0, 5)
rho_profile = np.array([0.0, 0.4, 0.8, 0.95, 1.0])
example_equilib_dat, example_Te = equilibrium_dat_and_te()
electron_temp = DataArray(
data=np.tile(np.array([3.0e3, 1.5e3, 0.5e3, 0.2e3, 0.1e3]),
(len(t), 1)).T,
coords=[("rho_poloidal", rho_profile), ("t", t)],
dims=["rho_poloidal", "t"],
)
offset = MagicMock(return_value=0.02)
example_equilibrium = Equilibrium(
example_equilib_dat,
example_Te,
sess=MagicMock(),
offset_picker=offset,
)
xr_rho_profile = DataArray(data=rho_profile,
coords={"rho_poloidal": rho_profile},
dims=["rho_poloidal"])
R_lfs_values, _ = example_equilibrium.R_lfs(xr_rho_profile)
# be, ne, ni, w
elements = ["be", "ne", "ni", "w"]
toroidal_rotations = np.array([200.0e3, 170.0e3, 100.0e3, 30.0e3, 5.0e3])
toroidal_rotations /= R_lfs_values.data[0, :]
toroidal_rotations = np.tile(toroidal_rotations,
(len(elements), len(t), 1))
toroidal_rotations = np.swapaxes(toroidal_rotations, 1, 2)
toroidal_rotations = DataArray(
data=toroidal_rotations,
coords=[("element", elements), ("rho_poloidal", rho_profile),
("t", t)],
dims=["element", "rho_poloidal", "t"],
)
ion_temperature = np.array([2.0e3, 1.2e3, 0.5e3, 0.2e3, 0.1e3])
ion_temperature = np.tile(ion_temperature, (len(elements), len(t), 1))
ion_temperature = np.swapaxes(ion_temperature, 1, 2)
ion_temperature = DataArray(
data=ion_temperature,
coords=[("element", elements), ("rho_poloidal", rho_profile),
("t", t)],
dims=["element", "rho_poloidal", "t"],
)
Zeff = DataArray(
data=1.85 * np.ones((*rho_profile.shape, len(t))),
coords=[("rho_poloidal", rho_profile), ("t", t)],
dims=["rho_poloidal", "t"],
)
main_ion = "d"
impurity = "be"
# toroidal_rotations has to be deepcopied otherwise it gets modified when
# passed to example_asymmetry.__call__
nominal_inputs = {
"toroidal_rotations": toroidal_rotations.copy(deep=True),
"ion_temperature": ion_temperature,
"main_ion": main_ion,
"impurity": impurity,
"Zeff": Zeff,
"electron_temp": electron_temp,
}
# Checking outputs of AsymmetryParameter() and ToroidalRotation()
asymmetry_parameters = zeros_like(toroidal_rotations)
try:
asymmetry_parameters.data[0] = example_asymmetry(**nominal_inputs)
except Exception as e:
raise e
nominal_inputs["impurity"] = "ne"
try:
asymmetry_parameters.data[1] = example_asymmetry(**nominal_inputs)
except Exception as e:
raise e
nominal_inputs["impurity"] = "ni"
try:
asymmetry_parameters.data[2] = example_asymmetry(**nominal_inputs)
except Exception as e:
raise e
nominal_inputs["impurity"] = "w"
try:
asymmetry_parameters.data[3] = example_asymmetry(**nominal_inputs)
except Exception as e:
raise e
example_toroidal_rotation = ToroidalRotation()
del nominal_inputs["toroidal_rotations"]
nominal_inputs["asymmetry_parameters"] = asymmetry_parameters
try:
output_toroidal_rotation = example_toroidal_rotation(**nominal_inputs)
except Exception as e:
raise e
expected_toroidal_rotation = toroidal_rotations.sel(element="w")
assert np.allclose(output_toroidal_rotation, expected_toroidal_rotation)
# Checking inputs for AsymmetryParameter
nominal_inputs = {
"toroidal_rotations": toroidal_rotations,
"ion_temperature": ion_temperature,
"main_ion": main_ion,
"impurity": impurity,
"Zeff": Zeff,
"electron_temp": electron_temp,
}
test_case_asymmetry = Exception_Asymmetry_Parameter_Test_Case(
**nominal_inputs)
for k, v in nominal_inputs.items():
if k == "impurity" or k == "main_ion":
continue
input_checking(k, test_case_asymmetry, nominal_inputs)
erroneous_input = {"impurity": 4}
test_case_asymmetry.call_type_check(**erroneous_input)
erroneous_input = {"impurity": "u"}
test_case_asymmetry.call_value_check(**erroneous_input)
erroneous_input = {"main_ion": 4}
test_case_asymmetry.call_type_check(**erroneous_input)
erroneous_input = {"main_ion": "u"}
test_case_asymmetry.call_value_check(**erroneous_input)
# Checking inputs for ToroidalRotation
nominal_inputs = {
"asymmetry_parameters": asymmetry_parameters,
"ion_temperature": ion_temperature,
"main_ion": main_ion,
"impurity": impurity,
"Zeff": Zeff,
"electron_temp": electron_temp,
}
test_case_toroidal = Exception_Toroidal_Rotation_Test_Case(
**nominal_inputs)
for k, v in nominal_inputs.items():
if k == "impurity" or k == "main_ion":
continue
input_checking(k, test_case_toroidal, nominal_inputs)
erroneous_input = {"impurity": 4}
test_case_toroidal.call_type_check(**erroneous_input)
erroneous_input = {"impurity": "u"}
test_case_toroidal.call_value_check(**erroneous_input)
erroneous_input = {"main_ion": 4}
test_case_toroidal.call_type_check(**erroneous_input)
erroneous_input = {"main_ion": "u"}
test_case_toroidal.call_value_check(**erroneous_input)