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)