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,
)