def type_check(self, truncation_threshold=None, electron_temperature=None):
"""Test TypeError for recover_threshold_rho()."""
inputs = [truncation_threshold, electron_temperature]
for i, iinput in enumerate(inputs):
if iinput is None:
inputs[i] = self.nominal_inputs[i]
(truncation_threshold, electron_temperature) = inputs
with self.assertRaises(TypeError):
recover_threshold_rho(*inputs)
def value_check(self, data_rho_theta, flux_surfaces, threshold_rho, t_arr):
"""Test ValueError for asymmetry_from_rho_theta()."""
inputs = [data_rho_theta, flux_surfaces, threshold_rho, t_arr]
for i, iinput in enumerate(inputs):
if iinput is None:
inputs[i] = self.nominal_inputs[i]
(data_rho_theta, flux_surfaces, threshold_rho, t_arr) = inputs
with self.assertRaises(ValueError):
recover_threshold_rho(*inputs)
def test_recover_threshold_rho():
initial_data = input_data_setup()
valid_input_Te = initial_data[1]
valid_truncation_threshold = initial_data[7]
try:
threshold_rho = recover_threshold_rho(valid_truncation_threshold,
valid_input_Te)
except Exception as e:
raise e
assert np.allclose(threshold_rho, 0.575)
example_recover_threshold_rho_test_case = Recover_Threshold_Rho_Test_Case(
valid_truncation_threshold, valid_input_Te)
invalid_input_checks(
example_recover_threshold_rho_test_case,
"truncation_threshold",
valid_truncation_threshold,
True,
)
invalid_input_checks(
example_recover_threshold_rho_test_case,
"electron_temperature",
valid_input_Te,
True,
)
def type_check(
self,
data_R_z=None,
flux_surfaces=None,
rho_arr=None,
threshold_rho=None,
t_arr=None,
):
"""Test TypeError for asymmetry_from_R_z()."""
inputs = [data_R_z, flux_surfaces, rho_arr, threshold_rho, t_arr]
for i, iinput in enumerate(inputs):
if iinput is None:
inputs[i] = self.nominal_inputs[i]
(data_R_z, flux_surfaces, rho_arr, threshold_rho, t_arr) = inputs
with self.assertRaises(TypeError):
recover_threshold_rho(*inputs)
def sxr_data_setup(input_data, sxr_truncation=True):
"""Set-up for soft-x-ray derived data.
Parameters
----------
input_data
Output of input_data_setup().
sxr_truncation
Boolean to determine whether or not to assume constant asymmetry parameter
outside the soft x-ray valid range.
Returns
-------
input_sxr_density_asym_Rz
Ground truth asymmetric impurity density on a (R, z) grid.
xarray.DataArray with dimensions (R, z, t)
input_sxr_density_asym_rho_theta
Ground truth asymmetric impurity density on a (rho, theta) grid.
xarray.DataArray with dimensions (rho, theta, t)
additional_sig
Gaussian signal to be added after extrapolation of the sxr_density
outside of the valid SXR limits.
example_asymmetry
Asymmetry parameter, xarray.DataArray with dimensions (rho, t)
asymmetry_modifier
Asymmetry modifier used to transform a low-field side only rho-profile
of a poloidally asymmetric quantity to a full poloidal cross-sectional
profile ie. (rho, t) -> (rho, theta, t). Also can be defined as:
exp(asymmetry_parameter * (R ** 2 - R_lfs ** 2)), where R is the major
radius as a function of (rho, theta, t) and R_lfs is the low-field-side
major radius as a function of (rho, t). xarray DataArray with dimensions
(rho, theta, t)
"""
(
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,
) = input_data
R_arr = np.linspace(1.83, 3.9, 100)
z_arr = np.linspace(-1.75, 2.0, 100)
R_arr = DataArray(data=R_arr, coords={"R": R_arr}, dims=["R"])
z_arr = DataArray(data=z_arr, coords={"z": z_arr}, dims=["z"])
additional_sig = zeros_like(input_Ne)
num_t = len(additional_sig.coords["t"].data)
for ind_t, it in enumerate(additional_sig.coords["t"].data):
amp = 1.0e17 - (ind_t / num_t) * 0.005e17
width = 0.05 - (ind_t / num_t) * 0.0004
pos = 0.9 - (ind_t / num_t) * 0.002
additional_sig.loc[:, it] = gaussian_perturbation(
(rho_arr, amp, width, pos))
sxr_density_data = 1e-3 * input_Ne.isel(t=0) # 25.0e15 * np.exp(-rho_arr)
sxr_density_data = np.tile(sxr_density_data,
(len(base_t), len(theta_arr), 1))
sxr_density_data = np.transpose(sxr_density_data, [2, 1, 0])
input_sxr_density = DataArray(
data=sxr_density_data,
coords={
"rho_poloidal": rho_arr,
"theta": theta_arr,
"t": base_t
},
dims=["rho_poloidal", "theta", "t"],
)
example_asymmetry_obj = AsymmetryParameter()
example_asymmetry = example_asymmetry_obj(
toroidal_rotations.copy(deep=True), input_Ti, "d", "w", Zeff, input_Te)
example_asymmetry = example_asymmetry.transpose("rho_poloidal", "t")
if sxr_truncation:
threshold_rho = recover_threshold_rho(valid_truncation_threshold,
input_Te)
example_asymmetry.loc[threshold_rho[0]:, :] = example_asymmetry.loc[
threshold_rho[0], :]
rho_derived, theta_derived = flux_surfs.convert_from_Rz(
R_arr, z_arr, base_t)
rho_derived = abs(rho_derived)
rho_derived = rho_derived.transpose("R", "z", "t")
theta_derived = theta_derived.transpose("R", "z", "t")
asymmetry_modifier = np.exp(example_asymmetry *
(R_derived**2 - R_lfs_values**2))
asymmetry_modifier = asymmetry_modifier.transpose("rho_poloidal", "theta",
"t")
input_sxr_density_lfs = input_sxr_density.sel(theta=0) # + additional_sig
additional_sig = additional_sig * asymmetry_modifier
additional_sig = additional_sig.transpose("rho_poloidal", "theta", "t")
input_sxr_density_asym_rho_theta = input_sxr_density_lfs * asymmetry_modifier
input_sxr_density_asym_rho_theta = input_sxr_density_asym_rho_theta.transpose(
"rho_poloidal", "theta", "t")
input_sxr_density_asym_R_z = input_sxr_density_asym_rho_theta.indica.interp2d(
{
"rho_poloidal": rho_derived,
"theta": theta_derived
}, method="linear")
input_sxr_density_asym_R_z = input_sxr_density_asym_R_z.fillna(0.0)
input_sxr_density_asym_R_z = input_sxr_density_asym_R_z.transpose(
"R", "z", "t")
return (
input_sxr_density_asym_R_z,
input_sxr_density_asym_rho_theta,
additional_sig,
example_asymmetry,
asymmetry_modifier,
)
def test_asymmetry_from_profile():
"""Test for asymmetry_from_R_z() and asymmetry_from_rho_theta()."""
initial_data = input_data_setup()
sxr_data = sxr_data_setup(initial_data, False)
input_Te = initial_data[1]
input_Ti = initial_data[2]
orig_toroidal_rotations = initial_data[3]
flux_surfs = initial_data[6]
impurity_sxr_density_asym_Rz = sxr_data[0]
impurity_sxr_density_asym_rho_theta = sxr_data[1]
orig_toroidal_rotations = orig_toroidal_rotations.sel(element="w")
rho_arr = impurity_sxr_density_asym_rho_theta.coords["rho_poloidal"]
t_arr = impurity_sxr_density_asym_rho_theta.coords["t"]
try:
asymmetry_from_R_z_ = asymmetry_from_R_z(impurity_sxr_density_asym_Rz,
flux_surfs,
rho_arr,
t_arr=t_arr)
except Exception as e:
raise e
try:
asymmetry_from_rho_theta_ = asymmetry_from_rho_theta(
impurity_sxr_density_asym_rho_theta, flux_surfs, t_arr=t_arr)
except Exception as e:
raise e
assert all(dim in asymmetry_from_R_z_.dims
for dim in ["rho_poloidal", "t"])
assert all(dim in asymmetry_from_rho_theta_.dims
for dim in ["rho_poloidal", "t"])
asymmetry_from_R_z_ = asymmetry_from_R_z_.expand_dims("element")
asymmetry_from_R_z_ = asymmetry_from_R_z_.assign_coords({"element": ["w"]})
asymmetry_from_rho_theta_ = asymmetry_from_rho_theta_.expand_dims(
"element")
asymmetry_from_rho_theta_ = asymmetry_from_rho_theta_.assign_coords(
{"element": ["w"]})
Zeff = DataArray(
data=1.85 * np.ones((*rho_arr.shape, len(t_arr))),
coords=[("rho_poloidal", rho_arr.data), ("t", t_arr.data)],
dims=["rho_poloidal", "t"],
)
main_ion = "d"
impurity = "w"
derived_toroidal_rotation_rho_theta = ToroidalRotation()
toroidal_rotation_results_rho_theta = derived_toroidal_rotation_rho_theta(
asymmetry_from_rho_theta_, input_Ti, main_ion, impurity, Zeff,
input_Te)
assert np.allclose(orig_toroidal_rotations,
toroidal_rotation_results_rho_theta,
rtol=0.05)
derived_toroidal_rotation_R_z = ToroidalRotation()
toroidal_rotation_results_R_z = derived_toroidal_rotation_R_z(
asymmetry_from_R_z_, input_Ti, main_ion, impurity, Zeff, input_Te)
# Last 3 rho points don't match most likely due to converison from R-z to rho-theta,
# not much that can be done about that.
assert np.allclose(orig_toroidal_rotations[:-3, :],
toroidal_rotation_results_R_z[:-3, :],
rtol=0.1)
valid_truncation_threshold = initial_data[7]
valid_threshold_rho = recover_threshold_rho(valid_truncation_threshold,
input_Te)
example_asymmetry_from_R_z_test_case = Asymmetry_From_R_z_Test_Case(
impurity_sxr_density_asym_Rz,
flux_surfs,
rho_arr,
valid_threshold_rho,
t_arr=t_arr,
)
invalid_input_checks(
example_asymmetry_from_R_z_test_case,
"data_R_z",
impurity_sxr_density_asym_Rz,
False,
)
invalid_input_checks(
example_asymmetry_from_R_z_test_case,
"flux_surfaces",
flux_surfs,
False,
)
invalid_input_checks(example_asymmetry_from_R_z_test_case, "rho_arr",
rho_arr, False)
invalid_input_checks(
example_asymmetry_from_R_z_test_case,
"threshold_rho",
valid_threshold_rho,
False,
)
invalid_input_checks(example_asymmetry_from_R_z_test_case, "t_arr", t_arr,
False)
example_asymmetry_from_rho_theta_test_case = Asymmetry_From_Rho_Theta_Test_Case(
impurity_sxr_density_asym_rho_theta,
flux_surfs,
valid_threshold_rho,
t_arr=t_arr,
)
invalid_input_checks(
example_asymmetry_from_rho_theta_test_case,
"data_rho_theta",
impurity_sxr_density_asym_rho_theta,
False,
)
invalid_input_checks(
example_asymmetry_from_rho_theta_test_case,
"flux_surfaces",
flux_surfs,
False,
)
invalid_input_checks(
example_asymmetry_from_rho_theta_test_case,
"threshold_rho",
valid_threshold_rho,
False,
)
invalid_input_checks(example_asymmetry_from_rho_theta_test_case, "t_arr",
t_arr, False)