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)