def hfd_mean(da: DataArray, datetime_coord: str = None) -> float:
"""Calculate mean half-flow duration.
Mean half-flow date (step on which the cumulative discharge since October 1st
reaches half of the annual discharge) [#]_.
Parameters
----------
da : DataArray
Array of flow values.
datetime_coord : str, optional
Datetime coordinate in the passed DataArray. Tried to infer automatically if not specified.
Returns
-------
float
Mean half-flow duration
References
----------
.. [#] Court, A.: Measures of streamflow timing. Journal of Geophysical Research (1896-1977), 1962, 67, 4335--4339,
doi:10.1029/JZ067i011p04335
"""
if datetime_coord is None:
datetime_coord = utils.infer_datetime_coord(da)
# determine the date of the first October 1st in the data period
first_date = da.coords[datetime_coord][0].values.astype(
'datetime64[s]').astype(datetime)
last_date = da.coords[datetime_coord][-1].values.astype(
'datetime64[s]').astype(datetime)
if first_date > datetime.strptime(f'{first_date.year}-10-01', '%Y-%m-%d'):
start_date = datetime.strptime(f'{first_date.year + 1}-10-01',
'%Y-%m-%d')
else:
start_date = datetime.strptime(f'{first_date.year}-10-01', '%Y-%m-%d')
end_date = start_date + relativedelta(years=1) - relativedelta(seconds=1)
doys = []
while end_date < last_date:
# compute cumulative sum for the selected period
data = da.sel({datetime_coord: slice(start_date, end_date)})
cs = data.cumsum(skipna=True)
# find steps with more cumulative discharge than the half annual sum
hf_steps = np.where(
~np.isnan(cs.where(cs > data.sum(skipna=True) / 2).values))[0]
# ignore days without discharge
if len(hf_steps) > 0:
# store the first step in the result array
doys.append(hf_steps[0])
start_date += relativedelta(years=1)
end_date += relativedelta(years=1)
return np.mean(doys)
def hfd_mean(da: DataArray, coord: str = "date") -> float:
# determine the date of the first October 1st in the data period
first_date = da.coords[coord][0].values.astype("datetime64[s]").astype(datetime)
last_date = da.coords[coord][-1].values.astype("datetime64[s]").astype(datetime)
if first_date > datetime.strptime(f"{first_date.year}-10-01", "%Y-%m-%d"):
start_date = datetime.strptime(f"{first_date.year + 1}-10-01", "%Y-%m-%d")
else:
start_date = datetime.strptime(f"{first_date.year}-10-01", "%Y-%m-%d")
end_date = start_date + relativedelta(years=1) - relativedelta(days=1)
doys = []
while end_date < last_date:
# compute cumulative sum for the selected period
data = da.sel({coord: slice(start_date, end_date)})
cs = data.cumsum(skipna=True)
# find days with more cumulative discharge than the half annual sum
days = np.where(~np.isnan(cs.where(cs > data.sum(skipna=True) / 2).values))[0]
# ignore days without discharge
if len(days) > 0:
# store the first day in the result array
doys.append(days[0])
start_date += relativedelta(years=1)
end_date += relativedelta(years=1)
return np.mean(doys)
def low_q_freq(da: DataArray, coord: str = "date", threshold: float = 0.2) -> float:
# determine the date of the first January 1st in the data period
first_date = da.coords[coord][0].values.astype("datetime64[s]").astype(datetime)
last_date = da.coords[coord][-1].values.astype("datetime64[s]").astype(datetime)
if first_date == datetime.strptime(f"{first_date.year}-01-01", "%Y-%m-%d"):
start_date = first_date
else:
start_date = datetime.strptime(f"{first_date.year + 1}-01-01", "%Y-%m-%d")
# end date of the first full year period
end_date = start_date + relativedelta(years=1) - relativedelta(days=1)
# determine the mean flow over the entire period
mean_flow = da.mean(skipna=True)
lqfs = []
while end_date < last_date:
data = da.sel({coord: slice(start_date, end_date)})
# number of days with discharge lower than threshold * median in a one year period
n_days = (data < (threshold * mean_flow)).sum()
lqfs.append(float(n_days))
start_date += relativedelta(years=1)
end_date += relativedelta(years=1)
return np.mean(lqfs)
def stream_elas(da: DataArray, prcp: DataArray, coord: str = 'date') -> float:
# rename precip coordinate name (to avoid problems with 'index' or 'date')
prcp = prcp.rename({list(prcp.coords.keys())[0]: coord})
# slice prcp to the same time window as the discharge
prcp = prcp.sel({coord: slice(da.coords[coord][0], da.coords[coord][-1])})
# determine the date of the first October 1st in the data period
first_date = da.coords[coord][0].values.astype('datetime64[s]').astype(
datetime)
last_date = da.coords[coord][-1].values.astype('datetime64[s]').astype(
datetime)
if first_date > datetime.strptime(f'{first_date.year}-10-01', '%Y-%m-%d'):
start_date = datetime.strptime(f'{first_date.year + 1}-10-01',
'%Y-%m-%d')
else:
start_date = datetime.strptime(f'{first_date.year}-10-01', '%Y-%m-%d')
end_date = start_date + relativedelta(years=1) - relativedelta(days=1)
# mask only valid time steps (only discharge has missing values)
idx = (da >= 0) & (~da.isnull())
da = da[idx]
prcp = prcp[idx]
# calculate long-term means
q_mean_total = da.mean()
p_mean_total = prcp.mean()
values = []
while end_date < last_date:
q = da.sel({coord: slice(start_date, end_date)})
p = prcp.sel({coord: slice(start_date, end_date)})
val = (q.mean() - q_mean_total) / (p.mean() - p_mean_total) * (
p_mean_total / q_mean_total)
values.append(val)
start_date += relativedelta(years=1)
end_date += relativedelta(years=1)
return np.median([float(v) for v in values])
def runoff_ratio(da: DataArray, prcp: DataArray) -> float:
# get precip coordinate name (to avoid problems with 'index' or 'date')
coord_name = list(prcp.coords.keys())[0]
# slice prcp to the same time window as the discharge
prcp = prcp.sel({coord_name: slice(da.coords["date"][0], da.coords["date"][-1])})
# calculate runoff ratio
value = da.mean() / prcp.mean()
return float(value)
def runoff_ratio(da: DataArray,
prcp: DataArray,
datetime_coord: str = None) -> float:
"""Calculate runoff ratio.
Runoff ratio (ratio of mean discharge to mean precipitation) [#]_ (Eq. 2).
Parameters
----------
da : DataArray
Array of flow values.
prcp : DataArray
Array of precipitation values.
datetime_coord : str, optional
Datetime coordinate in the passed DataArray. Tried to infer automatically if not specified.
Returns
-------
float
Runoff ratio.
References
----------
.. [#] Sawicz, K., Wagener, T., Sivapalan, M., Troch, P. A., and Carrillo, G.: Catchment classification: empirical
analysis of hydrologic similarity based on catchment function in the eastern USA.
Hydrology and Earth System Sciences, 2011, 15, 2895--2911, doi:10.5194/hess-15-2895-2011
"""
if datetime_coord is None:
datetime_coord = utils.infer_datetime_coord(da)
# rename precip coordinate name (to avoid problems with 'index' or 'date')
prcp = prcp.rename({list(prcp.coords.keys())[0]: datetime_coord})
# slice prcp to the same time window as the discharge
prcp = prcp.sel({
datetime_coord:
slice(da.coords[datetime_coord][0], da.coords[datetime_coord][-1])
})
# calculate runoff ratio
value = da.mean() / prcp.mean()
return float(value)
def __call__( # type: ignore
self,
asymmetry_parameters: DataArray,
ion_temperature: DataArray,
main_ion: str,
impurity: str,
Zeff: DataArray,
electron_temp: DataArray,
):
"""Calculates the toroidal rotation frequency from the asymmetry parameter.
Parameters
----------
asymmetry_parameters
xarray.DataArray containing asymmetry parameters data. In units of m^-2.
ion_temperature
xarray.DataArray containing ion temperature data. In units of eV.
main_ion
Element symbol of main ion.
impurity
Element symbol of chosen impurity element.
Zeff
xarray.DataArray containing Z-effective data from diagnostics.
electron_temp
xarray.DataArray containing electron temperature data. In units of eV.
Returns
-------
toroidal_rotation
xarray.DataArray containing data for toroidal rotation frequencies
for the given impurity element
"""
input_check(
"asymmetry_parameters",
asymmetry_parameters,
DataArray,
ndim_to_check=3,
greater_than_or_equal_zero=True,
)
input_check(
"ion_temperature",
ion_temperature,
DataArray,
ndim_to_check=3,
greater_than_or_equal_zero=False,
)
input_check("main_ion", main_ion, str)
try:
assert main_ion in list(ELEMENTS.keys())
except AssertionError:
raise ValueError(
f"main_ion must be one of {list(ELEMENTS.keys())}")
input_check("impurity", impurity, str)
try:
assert impurity in list(ELEMENTS.keys())
except AssertionError:
raise ValueError(
f"impurity must be one of {list(ELEMENTS.keys())}")
input_check("Zeff",
Zeff,
DataArray,
ndim_to_check=2,
greater_than_or_equal_zero=True)
input_check(
"electron_temp",
electron_temp,
DataArray,
ndim_to_check=2,
greater_than_or_equal_zero=False,
)
asymmetry_parameter = asymmetry_parameters.sel(element=impurity)
impurity_mass_int = ELEMENTS[impurity][1]
unified_atomic_mass_unit = 931.4941e6 # in eV/c^2
impurity_mass = float(impurity_mass_int) * unified_atomic_mass_unit
mean_charge = ELEMENTS[impurity][0]
main_ion_mass_int = ELEMENTS[main_ion][1]
main_ion_mass = float(main_ion_mass_int) * unified_atomic_mass_unit
ion_temperature = ion_temperature.sel(element=impurity)
# mypy on the github CI suggests that * is an Unsupported operand type
# between float and DataArray, don't know how to fix yet so for now ignored
toroidal_rotation = 2.0 * ion_temperature * asymmetry_parameter # type: ignore
toroidal_rotation /= impurity_mass * (
1.0 - (mean_charge * main_ion_mass * Zeff * electron_temp
) # type: ignore
/ (impurity_mass * (ion_temperature + Zeff * electron_temp)))
toroidal_rotation = toroidal_rotation**0.5
c = 3.0e8 # speed of light in vacuum
toroidal_rotation *= c
return toroidal_rotation
def stream_elas(da: DataArray,
prcp: DataArray,
datetime_coord: str = None) -> float:
"""Calculate stream elasticity.
Streamflow precipitation elasticity (sensitivity of streamflow to changes in precipitation at
the annual time scale) [#]_.
Parameters
----------
da : DataArray
Array of flow values.
prcp : DataArray
Array of precipitation values.
datetime_coord : str, optional
Datetime coordinate in the passed DataArray. Tried to infer automatically if not specified.
Returns
-------
float
Stream elasticity.
References
----------
.. [#] Sankarasubramanian, A., Vogel, R. M., and Limbrunner, J. F.: Climate elasticity of streamflow in the
United States. Water Resources Research, 2001, 37, 1771--1781, doi:10.1029/2000WR900330
"""
if datetime_coord is None:
datetime_coord = utils.infer_datetime_coord(da)
# rename precip coordinate name (to avoid problems with 'index' or 'date')
prcp = prcp.rename({list(prcp.coords.keys())[0]: datetime_coord})
# slice prcp to the same time window as the discharge
prcp = prcp.sel({
datetime_coord:
slice(da.coords[datetime_coord][0], da.coords[datetime_coord][-1])
})
# determine the date of the first October 1st in the data period
first_date = da.coords[datetime_coord][0].values.astype(
'datetime64[s]').astype(datetime)
last_date = da.coords[datetime_coord][-1].values.astype(
'datetime64[s]').astype(datetime)
if first_date > datetime.strptime(f'{first_date.year}-10-01', '%Y-%m-%d'):
start_date = datetime.strptime(f'{first_date.year + 1}-10-01',
'%Y-%m-%d')
else:
start_date = datetime.strptime(f'{first_date.year}-10-01', '%Y-%m-%d')
end_date = start_date + relativedelta(years=1) - relativedelta(seconds=1)
# mask only valid time steps (only discharge has missing values)
idx = (da >= 0) & (~da.isnull())
da = da[idx]
prcp = prcp[idx]
# calculate long-term means
q_mean_total = da.mean()
p_mean_total = prcp.mean()
values = []
while end_date < last_date:
q = da.sel({datetime_coord: slice(start_date, end_date)})
p = prcp.sel({datetime_coord: slice(start_date, end_date)})
val = (q.mean() - q_mean_total) / (p.mean() - p_mean_total) * (
p_mean_total / q_mean_total)
values.append(val)
start_date += relativedelta(years=1)
end_date += relativedelta(years=1)
return np.median([float(v) for v in values])
def low_q_freq(da: DataArray,
datetime_coord: str = None,
threshold: float = 0.2) -> float:
"""Calculate Low-flow frequency.
Frequency of low-flow events (<`threshold` times the median flow) [#]_, [#]_ (Table 2).
Parameters
----------
da : DataArray
Array of flow values.
datetime_coord : str, optional
Datetime coordinate in the passed DataArray. Tried to infer automatically if not specified.
threshold : float, optional
Low-flow threshold. Values below ``threshold * median`` are considered low flows.
Returns
-------
float
Low-flow frequency
References
----------
.. [#] Olden, J. D. and Poff, N. L.: Redundancy and the choice of hydrologic indices for characterizing streamflow
regimes. River Research and Applications, 2003, 19, 101--121, doi:10.1002/rra.700
.. [#] Westerberg, I. K. and McMillan, H. K.: Uncertainty in hydrological signatures.
Hydrology and Earth System Sciences, 2015, 19, 3951--3968, doi:10.5194/hess-19-3951-2015
"""
if datetime_coord is None:
datetime_coord = utils.infer_datetime_coord(da)
# determine the date of the first January 1st in the data period
first_date = da.coords[datetime_coord][0].values.astype(
'datetime64[s]').astype(datetime)
last_date = da.coords[datetime_coord][-1].values.astype(
'datetime64[s]').astype(datetime)
if first_date == datetime.strptime(f'{first_date.year}-01-01', '%Y-%m-%d'):
start_date = first_date
else:
start_date = datetime.strptime(f'{first_date.year + 1}-01-01',
'%Y-%m-%d')
# end date of the first full year period
end_date = start_date + relativedelta(years=1) - relativedelta(seconds=1)
# determine the mean flow over the entire period
mean_flow = da.mean(skipna=True)
lqfs = []
while end_date < last_date:
data = da.sel({datetime_coord: slice(start_date, end_date)})
# number of steps with discharge lower than threshold * median in a one year period
n_steps = (data < (threshold * mean_flow)).sum()
lqfs.append(float(n_steps))
start_date += relativedelta(years=1)
end_date += relativedelta(years=1)
return np.mean(lqfs)
def high_q_freq(da: DataArray,
datetime_coord: str = None,
threshold: float = 9.) -> float:
"""Calculate high-flow frequency.
Frequency of high-flow events (>`threshold` times the median flow) [#]_, [#]_ (Table 2).
Parameters
----------
da : DataArray
Array of flow values.
datetime_coord : str, optional
Datetime coordinate in the passed DataArray. Tried to infer automatically if not specified.
threshold : float, optional
High-flow threshold. Values larger than ``threshold * median`` are considered high flows.
Returns
-------
float
High-flow frequency
References
----------
.. [#] Clausen, B. and Biggs, B. J. F.: Flow variables for ecological studies in temperate streams: groupings based
on covariance. Journal of Hydrology, 2000, 237, 184--197, doi:10.1016/S0022-1694(00)00306-1
.. [#] Westerberg, I. K. and McMillan, H. K.: Uncertainty in hydrological signatures.
Hydrology and Earth System Sciences, 2015, 19, 3951--3968, doi:10.5194/hess-19-3951-2015
"""
if datetime_coord is None:
datetime_coord = utils.infer_datetime_coord(da)
# determine the date of the first January 1st in the data period
first_date = da.coords[datetime_coord][0].values.astype(
'datetime64[s]').astype(datetime)
last_date = da.coords[datetime_coord][-1].values.astype(
'datetime64[s]').astype(datetime)
if first_date == datetime.strptime(f'{first_date.year}-01-01', '%Y-%m-%d'):
start_date = first_date
else:
start_date = datetime.strptime(f'{first_date.year + 1}-01-01',
'%Y-%m-%d')
# end date of the first full year period
end_date = start_date + relativedelta(years=1) - relativedelta(seconds=1)
# determine the median flow over the entire period
median_flow = da.median(skipna=True)
hqfs = []
while end_date < last_date:
data = da.sel({datetime_coord: slice(start_date, end_date)})
# number of steps with discharge higher than threshold * median in a one year period
n_steps = (data > (threshold * median_flow)).sum()
hqfs.append(float(n_steps))
start_date += relativedelta(years=1)
end_date += relativedelta(years=1)
return np.mean(hqfs)
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)
