def transform_to_R_z( self, R_deriv: DataArray, z_deriv: DataArray, extrapolated_smooth_data: DataArray, flux_surfaces: FluxSurfaceCoordinates, ): """Function to transform data from an (rho, theta) grid to a (R, z) grid Parameters ---------- R_deriv Variable describing value of R in every coordinate on a (rho, theta) grid. xarray.DataArray with dimensions (rho, theta, t) (from derive_and_apply_asymmetry) z_deriv Variable describing value of z in every coordinate on a (rho, theta) grid. xarray.DataArray with dimensions (rho, theta, t) (from derive_and_apply_asymmetry) extrapolated_smooth_data Extrapolated and smoothed data to transform onto (R, z) grid. xarray.DataArray with dimensions (rho, theta, t) flux_surfaces FluxSurfaceCoordinates object representing polar coordinate systems using flux surfaces for the radial coordinate. Returns ------- extrapolated_smooth_data Extrapolated and smoothed data on (R, z) grid. xarray.DataArray with dimensions (R, z, t) """ R_arr = np.linspace(np.min(R_deriv[1:]), np.max(R_deriv[1:]), 40) z_arr = np.linspace(np.min(z_deriv), np.max(z_deriv), 40) R_arr = DataArray(R_arr, {"R": R_arr}, ["R"]) # type: ignore z_arr = DataArray(z_arr, {"z": z_arr}, ["z"]) # type: ignore t_arr = extrapolated_smooth_data.coords["t"] rho_derived, theta_derived = flux_surfaces.convert_from_Rz(R_arr, z_arr, t_arr) rho_derived = cast(DataArray, rho_derived).transpose("R", "z", "t") theta_derived = cast(DataArray, theta_derived).transpose("R", "z", "t") rho_derived = abs(rho_derived) extrapolated_smooth_data = extrapolated_smooth_data.indica.interp2d( {"rho_poloidal": rho_derived, "theta": theta_derived}, method="linear", assume_sorted=True, ) extrapolated_smooth_data = extrapolated_smooth_data.fillna(0.0) return extrapolated_smooth_data
def transform_to_rho_theta( self, data_R_z: DataArray, flux_surfaces: FluxSurfaceCoordinates, rho_arr: DataArray, t_arr: DataArray = None, ): """Function to transform data from an (R, z) grid to a (rho_poloidal, theta) grid Parameters ---------- data_R_z xarray.DataArray to be transformed. Dimensions (R, z, t) flux_surfaces FluxSurfaceCoordinates object representing polar coordinate systems using flux surfaces for the radial coordinate. rho_arr 1D xarray.DataArray of rho_poloidal from 0 to 1. t_arr 1D xarray.DataArray of t. Returns ------- data_rho_theta Transformed xarray.DataArray. Dimensions (rho_poloidal, theta, t) R_deriv Variable describing value of R in every coordinate on a (rho, theta) grid. xarray.DataArray with dimensions (rho, theta, t) z_deriv Variable describing value of z in every coordinate on a (rho, theta) grid. xarray.DataArray with dimensions (rho, theta, t) """ if t_arr is None: t_arr = data_R_z.coords["t"] theta_arr = np.linspace(-np.pi, np.pi, 21) # mypy doesn't like re-assignments which changes types. theta_arr = DataArray( # type: ignore data=theta_arr, coords={"theta": theta_arr}, dims=["theta"] ) R_deriv, z_deriv = flux_surfaces.convert_to_Rz(rho_arr, theta_arr, t_arr) R_deriv = cast(DataArray, R_deriv).transpose("rho_poloidal", "theta", "t") z_deriv = cast(DataArray, z_deriv).transpose("rho_poloidal", "theta", "t") data_rho_theta = data_R_z.indica.interp2d( {"R": R_deriv, "z": z_deriv}, method="linear", assume_sorted=True ) data_rho_theta = data_rho_theta.transpose("rho_poloidal", "theta", "t") return data_rho_theta, R_deriv, z_deriv
def test_mean_charge(): """Test MeanCharge.__call__.""" ADAS_file = ADASReader() element = "be" SCD = ADAS_file.get_adf11("scd", element, "89") ACD = ADAS_file.get_adf11("acd", element, "89") t = np.linspace(75.0, 80.0, 5) rho_profile = np.array([0.0, 0.4, 0.8, 0.95, 1.0]) input_Ne = DataArray( data=np.tile(np.array([5.0e19, 4.0e19, 3.0e19, 2.0e19, 1.0e19]), (len(t), 1)).T, coords=[("rho_poloidal", rho_profile), ("t", t)], dims=["rho_poloidal", "t"], ) input_Te = 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"], ) rho = DataArray( data=np.linspace(0.0, 1.0, 20), coords=[("rho_poloidal", np.linspace(0.0, 1.05, 20))], dims=["rho_poloidal"], ) dummy_coordinates = FluxSurfaceCoordinates("poloidal") input_Ne_spline = Spline(input_Ne, "rho_poloidal", dummy_coordinates) input_Ne = broadcast_spline( input_Ne_spline.spline, input_Ne_spline.spline_dims, input_Ne_spline.spline_coords, rho, ) input_Te_spline = Spline(input_Te, "rho_poloidal", dummy_coordinates) input_Te = broadcast_spline( input_Te_spline.spline, input_Te_spline.spline_dims, input_Te_spline.spline_coords, rho, ) example_frac_abundance = FractionalAbundance(SCD, ACD) example_frac_abundance.interpolate_rates(Ne=input_Ne.isel(t=0), Te=input_Te.isel(t=0)) example_frac_abundance.calc_ionisation_balance_matrix(Ne=input_Ne.isel( t=0)) example_frac_abundance.calc_F_z_tinf() example_frac_abundance.calc_eigen_vals_and_vecs() example_frac_abundance.calc_eigen_coeffs() F_z_t0 = np.real(example_frac_abundance.F_z_t0) F_z_t0 = F_z_t0.expand_dims("t", axis=-1) input_check = Exception_Mean_Charge_Test_Case() input_check.call_type_check(F_z_t0.data, element) input_check.call_value_check(F_z_t0 * -1, element) input_check.call_value_check(F_z_t0 * -np.inf, element) input_check.call_value_check(F_z_t0 * np.inf, element) input_check.call_value_check(F_z_t0 * np.nan, element) input_check.call_type_check(F_z_t0, 4) input_check.call_value_check(F_z_t0, "xy") input_check.call_assertion_check(F_z_t0[0:3], element) example_mean_charge = MeanCharge() result = example_mean_charge(F_z_t0, element) assert np.all(result == 0) F_z_tinf = example_frac_abundance.F_z_tinf F_z_tinf = F_z_tinf.expand_dims("t", axis=-1) example_mean_charge = MeanCharge() result = example_mean_charge(F_z_tinf, element) expected = np.zeros((*F_z_tinf.shape[1:], )) expected += 4.0 assert np.allclose(result, expected, rtol=1e-2)
def __call__( # type: ignore self, element: str, Zeff_LoS: DataArray, impurity_densities: DataArray, electron_density: DataArray, mean_charge: DataArray, flux_surfaces: FluxSurfaceCoordinates, t: DataArray = None, ): """Calculates the impurity concentration for the inputted element. Parameters ---------- element String specifying the symbol of the element for which the impurity concentration is desired. Zeff_LoS xarray.DataArray containing the Zeff value/s from Bremsstrahlung (ZEFH/KS3) impurity_densities xarray.DataArray of impurity densities for all impurity elements of interest. electron_density xarray.DataArray of electron density mean_charge xarray.DataArray of mean charge of all impurity elements of interest. This can be provided manually (with dimensions of ["element", "rho_poloidal", "t]), or can be passed as the results of MeanCharge.__call__ flux_surfaces FluxSurfaceCoordinates object that defines the flux surface geometry of the equilibrium of interest. t Optional, time at which the impurity concentration is to be calculated at. Returns ------- concentration xarray.DataArray containing the impurity concentration for the given impurity element. t If ``t`` was not specified as an argument for the __call__ function, return the time the results are given for. Otherwise return the argument. """ input_check( "impurity_densities", impurity_densities, DataArray, greater_than_or_equal_zero=True, ) input_check("element", element, str) elements_list = impurity_densities.coords["element"] try: assert element in elements_list except AssertionError: raise ValueError(f"Please input a single valid element from list:\ {elements_list}") if t is None: t = Zeff_LoS.t else: input_check("t", t, DataArray, ndim_to_check=1, greater_than_or_equal_zero=True) input_check( "Zeff_LoS", Zeff_LoS, DataArray, ndim_to_check=1, greater_than_or_equal_zero=True, ) input_check( "electron_density", electron_density, DataArray, ndim_to_check=2, greater_than_or_equal_zero=False, ) input_check( "mean_charge", mean_charge, DataArray, ndim_to_check=3, greater_than_or_equal_zero=True, ) input_check("flux_surfaces", flux_surfaces, FluxSurfaceCoordinates) Zeff_LoS = Zeff_LoS.interp(t=t, method="nearest") transform = Zeff_LoS.attrs["transform"] x1_name = transform.x1_name x2_name = transform.x2_name x1 = Zeff_LoS.coords[x1_name] x2_arr = np.linspace(0, 1, 300) x2 = DataArray(data=x2_arr, dims=[x2_name]) R_arr, z_arr = transform.convert_to_Rz(x1, x2, t) rho, theta = flux_surfaces.convert_from_Rz(R_arr, z_arr, t) if isinstance(R_arr, (DataArray, np.ndarray)): R_arr = R_arr.squeeze() if isinstance(rho, (DataArray, np.ndarray)): rho = rho.squeeze() if isinstance(rho, DataArray): rho = rho.drop_vars("t") rho = rho.drop_vars("R") rho = rho.drop_vars("z") rho = rho.fillna(2.0) if set(["R", "z"]).issubset(set(list(impurity_densities.dims))): impurity_densities = impurity_densities.indica.interp2d( z=z_arr, R=R_arr, method="cubic", assume_sorted=True, ) elif set(["rho_poloidal", "theta"]).issubset(set(list(impurity_densities.dims))): impurity_densities = impurity_densities.interp(rho_poloidal=rho, theta=theta, method="linear", assume_sorted=True) elif set(["rho_poloidal" ]).issubset(set(list(impurity_densities.dims))): impurity_densities = impurity_densities.interp(rho_poloidal=rho, method="linear", assume_sorted=True) else: raise ValueError( 'Inputted impurity densities does not have any compatible\ coordinates: ["rho_poloidal", "theta"], ["rho_poloidal"]\ or ["R", "z"]') impurity_densities = impurity_densities.interp(t=t, method="linear", assume_sorted=True) electron_density = electron_density.interp(rho_poloidal=rho, method="linear", assume_sorted=True) electron_density = electron_density.interp(t=t, method="linear", assume_sorted=True) mean_charge = mean_charge.interp(rho_poloidal=rho, method="linear", assume_sorted=True) mean_charge = mean_charge.interp(t=t, method="linear", assume_sorted=True) dl = transform.distance(x2_name, DataArray(0), x2[0:2], 0) dl = dl[1] LoS_length = dl * 300 concentration = zeros_like(Zeff_LoS) term_1 = LoS_length * (Zeff_LoS - 1) term_2 = zeros_like(term_1) for k, kdens in enumerate(impurity_densities.coords["element"]): if element == kdens: term_3 = (mean_charge[k]**2 - mean_charge[k]).sum(x2_name) * dl continue term2_integrand = (impurity_densities[k] / electron_density) * ( mean_charge[k]**2 - mean_charge[k]) term_2 += term2_integrand.sum(x2_name) * dl concentration = (term_1 - term_2) / term_3 return concentration, t
def input_data_setup(): """Initial set-up for the majority of the data needed for ExtrapolateImpurityDensity. Returns ------- input_Ne xarray.DataArray of electron density. Dimensions (rho, t) input_Te xarray.DataArray of electron temperature. Dimensions (rho, t) input_Ti xarray.DataArray of ion temperature. Dimensions (elements, rho, t) toroidal_rotations xarray.DataArray of toroidal rotations (needed for calculating the centrifugal asymmetry parameter). Dimensions (elements, rho, t) rho_arr xarray.DataArray of rho values, np.linspace(0, 1, 41). Dimensions (rho) theta_arr xarray.DataArray of theta values, np.linspace(-np.pi, np.pi, 21). Dimensions (theta) flux_surfs FluxSurfaceCoordinates object representing polar coordinate systems using flux surfaces for the radial coordinate. valid_truncation_threshold Truncation threshold (float) for the electron temperature (below this value soft-xray measurements are not valid) Zeff xarray.DataArray of the effective z(atomic number)-value for the plasma. Dimensions (rho, t) base_t xarray.DataArray of time values. Dimensions (t) R_derived Variable describing value of R in every coordinate on a (rho, theta) grid. xarray.DataArray with dimensions (rho, theta, t) R_lfs_values R_derived values at theta = 0 (ie low-field-side of the tokamak). xarray.DataArray with dimensions (rho, t) elements List of element symbols for all impurities. """ base_rho_profile = np.array([0.0, 0.4, 0.8, 0.95, 1.0]) base_t = np.linspace(75.0, 80.0, 20) input_Te = np.array([3.0e3, 1.5e3, 0.5e3, 0.2e3, 0.1e3]) input_Te = np.tile(input_Te, (len(base_t), 1)).T input_Te = DataArray( data=np.tile(np.array([3.0e3, 1.5e3, 0.5e3, 0.2e3, 0.1e3]), (len(base_t), 1)).T, coords={ "rho_poloidal": base_rho_profile, "t": base_t }, dims=["rho_poloidal", "t"], ) input_Ne = np.array([5.0e19, 4.0e19, 3.0e19, 2.0e19, 1.0e19]) input_Ne = np.tile(input_Ne, (len(base_t), 1)).T input_Ne = DataArray( data=np.tile(np.array([5.0e19, 4.0e19, 3.0e19, 2.0e19, 1.0e19]), (len(base_t), 1)).T, coords={ "rho_poloidal": base_rho_profile, "t": base_t }, dims=["rho_poloidal", "t"], ) elements = ["be", "ne", "ni", "w"] input_Ti = np.array([2.0e3, 1.2e3, 0.5e3, 0.2e3, 0.1e3]) input_Ti = np.tile(input_Ti, (len(elements), len(base_t), 1)) input_Ti = np.swapaxes(input_Ti, 1, 2) input_Ti = DataArray( data=input_Ti, coords={ "element": elements, "rho_poloidal": base_rho_profile, "t": base_t }, dims=["element", "rho_poloidal", "t"], ) toroidal_rotations = np.array([200.0e3, 170.0e3, 100.0e3, 30.0e3, 5.0e3]) toroidal_rotations = np.tile(toroidal_rotations, (len(elements), len(base_t), 1)) toroidal_rotations = np.swapaxes(toroidal_rotations, 1, 2) toroidal_rotations = DataArray( data=toroidal_rotations, coords=[ ("element", elements), ("rho_poloidal", base_rho_profile), ("t", base_t), ], dims=["element", "rho_poloidal", "t"], ) expanded_rho = np.linspace(base_rho_profile[0], base_rho_profile[-1], 41) rho_arr = expanded_rho theta_arr = np.linspace(-np.pi, np.pi, 21) rho_arr = DataArray(data=rho_arr, coords={"rho_poloidal": rho_arr}, dims=["rho_poloidal"]) theta_arr = DataArray(data=theta_arr, coords={"theta": theta_arr}, dims=["theta"]) flux_surfs = FluxSurfaceCoordinates("poloidal") offset = MagicMock(return_value=0.02) equilib_dat, Te = equilibrium_dat_and_te() equilib = Equilibrium(equilib_dat, Te, sess=MagicMock(), offset_picker=offset) flux_surfs.set_equilibrium(equilib) R_derived, _ = flux_surfs.convert_to_Rz( DataArray(expanded_rho, {"rho_poloidal": expanded_rho}, dims=["rho_poloidal"]), theta_arr, base_t, ) R_derived = R_derived.transpose("rho_poloidal", "theta", "t") R_lfs_values = R_derived.sel(theta=0, method="nearest") input_Te = input_Te.interp(rho_poloidal=expanded_rho, method="linear") input_Ne = input_Ne.interp(rho_poloidal=expanded_rho, method="linear") input_Ti = input_Ti.interp(rho_poloidal=expanded_rho, method="linear") toroidal_rotations = toroidal_rotations.interp(rho_poloidal=expanded_rho, method="linear") toroidal_rotations /= R_lfs_values.data # re-scale from velocity to frequency valid_truncation_threshold = 1.0e3 Zeff = DataArray( data=1.85 * np.ones((*expanded_rho.shape, len(base_t))), coords=[("rho_poloidal", expanded_rho), ("t", base_t)], dims=["rho_poloidal", "t"], ) return ( input_Ne, input_Te, input_Ti, toroidal_rotations, rho_arr, theta_arr, flux_surfs, valid_truncation_threshold, Zeff, base_t, R_derived, R_lfs_values, elements, )
def asymmetry_from_R_z( data_R_z: DataArray, flux_surfaces: FluxSurfaceCoordinates, rho_arr: DataArray, threshold_rho: DataArray = None, t_arr: DataArray = None, ): """Function to calculate an asymmetry parameter from a given density profile in (R, z, t) coordinates. Parameters ---------- data_R_z High-z density profile which is to be used to calculate the asymmetry parameter. xarray.DataArray with dimensions (R, z, t) flux_surfaces FluxSurfaceCoordinates object representing polar coordinate systems using flux surfaces for the radial coordinate. rho_arr 1D xarray.DataArray of rho from 0 to 1. threshold_rho rho value denoting the cutoff point beyond which soft x-ray diagnostics are invalid. It's also used in setting the derived asymmetry parameter to be flat in the invalid region. xarray.DataArray with dimensions (t) t_arr 1D xarray.DataArray of t. Returns ------- derived_asymmetry_parameter Derived asymmetry parameter. xarray.DataArray with dimensions (rho, t) """ input_check("data_R_z", data_R_z, DataArray, 3, True) input_check("flux_surfaces", flux_surfaces, FluxSurfaceCoordinates) input_check("rho_arr", rho_arr, DataArray, 1, True) if threshold_rho is not None: input_check("threshold_rho", threshold_rho, DataArray, 1, True) if t_arr is None: t_arr = data_R_z.coords["t"] else: input_check("t_arr", t_arr, DataArray, 1, True) theta_arr_ = np.array([0.0, np.pi], dtype=float) theta_arr = DataArray(data=theta_arr_, coords={"theta": theta_arr_}, dims=["theta"]) R_deriv, z_deriv = flux_surfaces.convert_to_Rz(rho_arr, theta_arr) R_deriv = cast(DataArray, R_deriv).interp(t=t_arr, method="linear") z_deriv = cast(DataArray, z_deriv).interp(t=t_arr, method="linear") R_deriv = cast(DataArray, R_deriv).transpose("rho_poloidal", "theta", "t") z_deriv = cast(DataArray, z_deriv).transpose("rho_poloidal", "theta", "t") data_rho_theta = data_R_z.indica.interp2d( {"R": R_deriv, "z": z_deriv}, method="linear", assume_sorted=True ) data_rho_theta = data_rho_theta.transpose("rho_poloidal", "theta", "t") R_lfs_midplane = cast(DataArray, R_deriv).isel(theta=0) # theta = 0.0 R_hfs_midplane = cast(DataArray, R_deriv).isel(theta=1) # theta = np.pi derived_asymmetry_parameter = np.log( data_rho_theta.isel(theta=1) / data_rho_theta.isel(theta=0) ) derived_asymmetry_parameter /= R_hfs_midplane**2 - R_lfs_midplane**2 # Set constant asymmetry parameter for rho<0.1 derived_asymmetry_parameter = derived_asymmetry_parameter.where( derived_asymmetry_parameter.coords["rho_poloidal"] > 0.1, other=derived_asymmetry_parameter.sel({"rho_poloidal": 0.1}, method="nearest"), ) derived_asymmetry_parameter = np.abs(derived_asymmetry_parameter) if threshold_rho is not None: for ind_t, it in enumerate(threshold_rho.coords["t"]): derived_asymmetry_parameter.loc[ threshold_rho[ind_t] :, it # type:ignore ] = derived_asymmetry_parameter.loc[threshold_rho[ind_t], it] return derived_asymmetry_parameter
def input_data_setup(): """Initial set-up for the majority of the data needed for ExtrapolateImpurityDensity. Returns ------- input_Ne xarray.DataArray of electron density. Dimensions (rho, t) input_Te xarray.DataArray of electron temperature. Dimensions (rho, t) flux_surfs FluxSurfaceCoordinates object representing polar coordinate systems using flux surfaces for the radial coordinate. base_t xarray.DataArray of time values. Dimensions (t) elements List of element symbols for all impurities. rho_arr xarray.DataArray of rho values, np.linspace(0, 1, 41). Dimensions (rho) theta_arr xarray.DataArray of theta values, np.linspace(-np.pi, np.pi, 21). Dimensions (theta) """ base_rho_profile = np.array([0.0, 0.4, 0.8, 0.95, 1.0]) base_t = np.linspace(75.0, 80.0, 20) input_Te = np.array([3.0e3, 1.5e3, 0.5e3, 0.2e3, 0.1e3]) input_Te = np.tile(input_Te, (len(base_t), 1)).T input_Te = DataArray( data=np.tile(np.array([3.0e3, 1.5e3, 0.5e3, 0.2e3, 0.1e3]), (len(base_t), 1)).T, coords={ "rho_poloidal": base_rho_profile, "t": base_t }, dims=["rho_poloidal", "t"], ) input_Ne = np.array([5.0e19, 4.0e19, 3.0e19, 2.0e19, 1.0e19]) input_Ne = np.tile(input_Ne, (len(base_t), 1)).T input_Ne = DataArray( data=np.tile(np.array([5.0e19, 4.0e19, 3.0e19, 2.0e19, 1.0e19]), (len(base_t), 1)).T, coords={ "rho_poloidal": base_rho_profile, "t": base_t }, dims=["rho_poloidal", "t"], ) elements = ["be", "ne", "ni"] expanded_rho = np.linspace(base_rho_profile[0], base_rho_profile[-1], 41) rho_arr = expanded_rho theta_arr = np.linspace(-np.pi, np.pi, 21) rho_arr = DataArray(data=rho_arr, coords={"rho_poloidal": rho_arr}, dims=["rho_poloidal"]) theta_arr = DataArray(data=theta_arr, coords={"theta": theta_arr}, dims=["theta"]) flux_surfs = FluxSurfaceCoordinates("poloidal") offset = MagicMock(return_value=0.02) equilib_dat, Te = equilibrium_dat_and_te() equilib = Equilibrium(equilib_dat, Te, sess=MagicMock(), offset_picker=offset) flux_surfs.set_equilibrium(equilib) input_Te = input_Te.interp(rho_poloidal=expanded_rho, method="linear") input_Ne = input_Ne.interp(rho_poloidal=expanded_rho, method="linear") return ( input_Ne, input_Te, flux_surfs, base_t, elements, rho_arr, theta_arr, )