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)