def test_spectral_density_input_errors(kwargs, error, msg, single_species_collective_args): """ This test validates errors with invalid argument and keyword arguments in spectral_density """ args = single_species_collective_args # Replace any modified keys for key, value in kwargs.items(): args[key] = value # Separate the arguments into args and kwargs for spectral_density args, kwargs = spectral_density_args_kwargs(args) if error is None: alpha, Skw = thomson.spectral_density(*args, **kwargs) else: with pytest.raises(error) as excinfo: alpha, Skw = thomson.spectral_density(*args, **kwargs) # If msg is not None, check that this string is a subset of the # error message if msg is not None: assert msg in str(excinfo.value)
def test_split_populations(): """ This test makes sure that splitting a single population of ions or electrons into two identical halves returns the same result. """ wavelengths = np.arange(520, 545, 0.01) * u.nm probe_wavelength = 532 * u.nm n = 5e17 * u.cm**-3 probe_vec = np.array([1, 0, 0]) scatter_vec = np.array([0, 1, 0]) # Combined T_e = np.array([10]) * u.eV T_i = np.array([10]) * u.eV ions = ["H+"] ifract = np.array([1.0]) efract = np.array([1.0]) alpha, Skw0 = thomson.spectral_density( wavelengths, probe_wavelength, n, T_e=T_e, T_i=T_i, ifract=ifract, efract=efract, ions=ions, probe_vec=probe_vec, scatter_vec=scatter_vec, ) # Split e and i populations into two parts # this should not change the results since the parts are identical T_e = np.array([10, 10]) * u.eV T_i = np.array([10, 10]) * u.eV ions = ["H+", "H+"] ifract = np.array([0.2, 0.8]) efract = np.array([0.8, 0.2]) alpha, Skw1 = thomson.spectral_density( wavelengths, probe_wavelength, n, T_e=T_e, T_i=T_i, ifract=ifract, efract=efract, ions=ions, probe_vec=probe_vec, scatter_vec=scatter_vec, ) # Calculate the deviation between the two spectra # (any differences should be in the noise) deviation = (Skw0 - Skw1) / Skw0 * 100 assert np.all(deviation < 1e-6), "Failed split populations test"
def test_different_input_types(): # Define some constants wavelengths = np.arange(520, 545, 0.01) * u.nm probe_wavelength = 532 * u.nm ne = 5e17 * u.cm**-3 probe_vec = np.array([1, 0, 0]) scatter_vec = np.array([0, 1, 0]) fract = np.array([1.0]) Te = 10 * u.eV Ti = np.array([10]) * u.eV ion_species = "C-12 5+" # Raise a ValueError with inconsistent ion array lengths with pytest.raises(ValueError): alpha, Skw = thomson.spectral_density( wavelengths, probe_wavelength, ne, Te, Ti, fract=np.array([0.5, 0.5]), ion_species=ion_species, probe_vec=probe_vec, scatter_vec=scatter_vec, ) # Raise a ValueError with inconsistent ion temperature array with pytest.raises(ValueError): alpha, Skw = thomson.spectral_density( wavelengths, probe_wavelength, ne, Te, np.array([5, 5]) * u.eV, fract=fract, ion_species=ion_species, probe_vec=probe_vec, scatter_vec=scatter_vec, ) # Raise a ValueError with empty ion_species with pytest.raises(ValueError): alpha, Skw = thomson.spectral_density( wavelengths, probe_wavelength, ne, Te, Ti, fract=fract, ion_species=[], probe_vec=probe_vec, scatter_vec=scatter_vec, )
def test_spectral_density_minimal_arguments(single_species_collective_args): """ Check that spectral density runs with minimal arguments """ wavelengths = single_species_collective_args["wavelengths"] args, kwargs = spectral_density_args_kwargs(single_species_collective_args) # Delete the arguments that have default values optional_keys = [ "efract", "ifract", "ions", "electron_vel", "ion_vel", "probe_vec", "scatter_vec", "instr_func", ] for key in optional_keys: if key in kwargs.keys(): del kwargs[key] alpha, Skw = thomson.spectral_density(*args, **kwargs) return (alpha, wavelengths, Skw)
def gen_non_collective_spectrum(): """ Generates an example Thomson scattering spectrum in the non-collective regime """ wavelengths = np.arange(500, 570, 0.01) * u.nm probe_wavelength = 532 * u.nm n = 5e15 * u.cm**-3 probe_vec = np.array([1, 0, 0]) scatter_vec = np.array([0, 1, 0]) Te = 100 * u.eV Ti = np.array([10]) * u.eV ion_species = ["H+"] alpha, Skw = thomson.spectral_density( wavelengths, probe_wavelength, n, Te, Ti, ion_species=ion_species, probe_vec=probe_vec, scatter_vec=scatter_vec, ) return alpha, wavelengths, Skw
def gen_multiple_ion_species_spectrum(): """ Generates an example Thomson scattering spectrum for multiple ion species that also have drift velocities. Parameters are set to be in the collective regime where ion species are important. """ wavelengths = np.arange(520, 545, 0.01) * u.nm probe_wavelength = 532 * u.nm n = 5e17 * u.cm**-3 probe_vec = np.array([1, 0, 0]) scatter_vec = np.array([0, 1, 0]) ifract = np.array([0.7, 0.3]) Te = 10 * u.eV Ti = np.array([5, 5]) * u.eV electron_vel = np.array([[300, 0, 0]]) * u.km / u.s ion_vel = np.array([[-500, 0, 0], [0, 500, 0]]) * u.km / u.s # Use this to also test passing in ion species as Particle objects ion_species = [Particle("p+"), Particle("C-12 5+")] alpha, Skw = thomson.spectral_density( wavelengths, probe_wavelength, n, Te, Ti, ifract=ifract, ion_species=ion_species, probe_vec=probe_vec, scatter_vec=scatter_vec, electron_vel=electron_vel, ion_vel=ion_vel, ) return alpha, wavelengths, Skw
def gen_collective_spectrum(): """ Generates an example Thomson scattering spectrum in the collective regime """ wavelengths = np.arange(520, 545, 0.01) * u.nm probe_wavelength = 532 * u.nm ne = 5e17 * u.cm**-3 probe_vec = np.array([1, 0, 0]) scatter_vec = np.array([0, 1, 0]) fract = np.array([1.0]) Te = 10 * u.eV Ti = np.array([10]) * u.eV ion_species = ["C-12 5+"] alpha, Skw = thomson.spectral_density( wavelengths, probe_wavelength, ne, Te, Ti, fract=fract, ion_species=ion_species, probe_vec=probe_vec, scatter_vec=scatter_vec, ) return alpha, wavelengths, Skw
def test_thomson_with_instrument_function(single_species_collective_args): """ Generates an example Thomson scattering spectrum with an instrument function applied """ wavelengths = single_species_collective_args["wavelengths"] args, kwargs = spectral_density_args_kwargs(single_species_collective_args) alpha, Skw_with = thomson.spectral_density(*args, **kwargs, instr_func=example_instr_func) alpha, Skw_without = thomson.spectral_density(*args, **kwargs) # Assert that the instrument function has made the IAW peak wider w1 = width_at_value(wavelengths.value, Skw_with.value, 2e-13) w2 = width_at_value(wavelengths.value, Skw_without.value, 2e-13) assert w1 > w2
def single_species_collective_spectrum(single_species_collective_args): """ Generates an example Thomson scattering spectrum in the collective regime """ wavelengths = single_species_collective_args["wavelengths"] args, kwargs = spectral_density_args_kwargs(single_species_collective_args) alpha, Skw = thomson.spectral_density(*args, **kwargs) return (alpha, wavelengths, Skw)
def test_thomson_with_invalid_instrument_function( instr_func, single_species_collective_args, ): """ Verifies that an exception is raised if the provided instrument function is invalid. """ args, kwargs = spectral_density_args_kwargs(single_species_collective_args) kwargs["instr_func"] = instr_func with pytest.raises(ValueError): alpha, Skw_with = thomson.spectral_density(*args, **kwargs)
def test_single_species_collective_lite(single_species_collective_args): # Make a copy of the input args args_fixture_copy = copy.copy(single_species_collective_args) args, kwargs = spectral_density_args_kwargs(single_species_collective_args) alpha1, Skw1 = thomson.spectral_density(*args, **kwargs) lite_kwargs = args_to_lite_args(args_fixture_copy) args, kwargs = spectral_density_args_kwargs(lite_kwargs) alpha2, Skw2 = thomson.spectral_density.lite(*args, **kwargs) assert np.isclose(alpha1, alpha2) assert np.allclose(Skw1.to(u.s / u.rad).value, Skw2)
def multiple_species_collective_spectrum(multiple_species_collective_args): """ Generates an example Thomson scattering spectrum for multiple ion species that also have drift velocities. Parameters are set to be in the collective regime where ion species are important. """ wavelengths = multiple_species_collective_args["wavelengths"] args, kwargs = spectral_density_args_kwargs( multiple_species_collective_args) alpha, Skw = thomson.spectral_density(*args, **kwargs) return (alpha, wavelengths, Skw)
def run_fit( wavelengths, params, settings, noise_amp=0.05, notch=None, notch_as_nan=False, fit_method="differential_evolution", fit_kws={}, max_iter=None, check_errors=True, require_redchi=1, # If false, don't perform the actual fit but instead just create the Model run_fit=True, ): """ This function takes a Parameters object, generates some synthetic data near it, perturbs the initial values, then tries a fit """ wavelengths = (wavelengths * u.m).to(u.nm) true_params = copy.deepcopy(params) skeys = list(settings.keys()) pkeys = list(params.keys()) # Fill any missing required parameters if "efract_0" not in pkeys: params.add("efract_0", value=1.0, vary=False) if "ifract_0" not in pkeys: params.add("ifract_0", value=1.0, vary=False) if "electron_speed" not in pkeys: params.add("electron_speed_0", value=0.0, vary=False) if "ion_speed" not in pkeys: params.add("ion_speed_0", value=0.0, vary=False) # LOAD FROM PARAMS n = params["n"] T_e = thomson._params_to_array(params, "T_e") T_i = thomson._params_to_array(params, "T_i") efract = thomson._params_to_array(params, "efract") ifract = thomson._params_to_array(params, "ifract") electron_speed = thomson._params_to_array(params, "electron_speed") ion_speed = thomson._params_to_array(params, "ion_speed") if "instr_func" not in skeys: settings["instr_func"] = None # LOAD FROM SETTINGS ions = settings["ions"] probe_vec = settings["probe_vec"] scatter_vec = settings["scatter_vec"] probe_wavelength = settings["probe_wavelength"] instr_func = settings["instr_func"] electron_vdir = settings.get("electron_vdir", np.ones([len(T_e), 3])) ion_vdir = settings.get("ion_vdir", np.ones([len(T_i), 3])) electron_vel = electron_speed[:, np.newaxis] * electron_vdir ion_vel = ion_speed[:, np.newaxis] * ion_vdir # Create the synthetic data alpha, Skw = thomson.spectral_density( wavelengths, probe_wavelength * u.m, n * u.m**-3, T_e=T_e * u.eV, T_i=T_i * u.eV, ifract=ifract, efract=efract, ions=ions, probe_vec=probe_vec, scatter_vec=scatter_vec, electron_vel=electron_vel * u.m / u.s, ion_vel=ion_vel * u.m / u.s, instr_func=instr_func, ) data = Skw if notch is not None: x0 = np.argmin(np.abs(wavelengths.to(u.m).value * 1e9 - notch[0])) x1 = np.argmin(np.abs(wavelengths.to(u.m).value * 1e9 - notch[1])) # Depending on the notch_as_nan keyword, either delete the missing data # or replace with NaN values if notch_as_nan: data[x0:x1] = np.nan else: data = np.delete(data, np.arange(x0, x1)) wavelengths = np.delete(wavelengths, np.arange(x0, x1)) data *= 1 + np.random.normal(loc=0, scale=noise_amp, size=wavelengths.size) data *= 1 / np.nanmax(data) # Randomly choose the starting values of the parameters within the # search space (to make the algorithm do some work!) for p in list(params.keys()): if params[p].vary: params[p].value = np.random.uniform(low=params[p].min, high=params[p].max, size=1) # Make the model, then perform the fit model = thomson.spectral_density_model( wavelengths.to(u.m).value, settings, params) if run_fit: result = model.fit( data, params, wavelengths=wavelengths.to(u.m).value, method=fit_method, max_nfev=max_iter, fit_kws=fit_kws, ) # Assert that the fit reduced chi2 is under the requirement specified assert result.redchi < require_redchi
def test_ifract_sum_error(single_species_collective_args): args, kwargs = spectral_density_args_kwargs(single_species_collective_args) kwargs["ifract"] = np.array([0.5, 1.2]) # Sum is not 1 with pytest.raises(ValueError): alpha, Skw = thomson.spectral_density(*args, **kwargs)
def test_fit_with_minimal_parameters(): # Create example data for fitting probe_wavelength = 532 * u.nm probe_vec = np.array([1, 0, 0]) scattering_angle = np.deg2rad(90) scatter_vec = np.array( [np.cos(scattering_angle), np.sin(scattering_angle), 0]) w0 = probe_wavelength.value wavelengths = np.linspace(w0 - 5, w0 + 5, num=512) * u.nm ions = ["H+"] n = 2e17 * u.cm**-3 T_i = 20 * u.eV T_e = 10 * u.eV alpha, Skw = thomson.spectral_density( wavelengths, probe_wavelength, n, T_e=T_e, T_i=T_i, ions=ions, probe_vec=probe_vec, scatter_vec=scatter_vec, ) data = Skw.value data *= 1 + np.random.normal(loc=0, scale=0.1, size=wavelengths.size) data *= 1 / np.nanmax(data) # Create settings and params using only the minimal parameters # intentionally leave out a few required values to check to make sure an # exception is raised settings = {} settings["probe_vec"] = probe_vec settings["scatter_vec"] = scatter_vec settings["ions"] = ions params = Parameters() params.add("T_e_0", value=T_e.value, vary=False, min=5, max=20) params.add("T_i_0", value=T_i.value, vary=True, min=5, max=70) # Try creating model: will raise exception because required values # are missing in settings, eg. 'probe_wavelength' with pytest.raises(ValueError): model = thomson.spectral_density_model(wavelengths, settings, params) # Add back in the required values settings["probe_wavelength"] = probe_wavelength.to(u.m).value # Still raises an exception because T_e_0 is still missing with pytest.raises(ValueError): model = thomson.spectral_density_model(wavelengths, settings, params) params.add("n", value=n.to(u.m**-3).value, vary=False) # Make the model, then perform the fit model = thomson.spectral_density_model( wavelengths.to(u.m).value, settings, params) result = model.fit( data, params, wavelengths=wavelengths.to(u.m).value, method="differential_evolution", max_nfev=2000, )
def test_different_input_types(): # Define some constants wavelengths = np.arange(520, 545, 0.01) * u.nm probe_wavelength = 532 * u.nm n = 5e17 * u.cm**-3 probe_vec = np.array([1, 0, 0]) scatter_vec = np.array([0, 1, 0]) ifract = np.array([1.0]) Te = np.array([10]) * u.eV Ti = np.array([10]) * u.eV ion_species = "C-12 5+" # Raise a ValueError with inconsistent ion array lengths with pytest.raises(ValueError): alpha, Skw = thomson.spectral_density( wavelengths, probe_wavelength, n, Te, Ti, ifract=np.array([0.5, 0.5]), ion_species=ion_species, probe_vec=probe_vec, scatter_vec=scatter_vec, ) # Raise a ValueError with inconsistent ion temperature array with pytest.raises(ValueError): alpha, Skw = thomson.spectral_density( wavelengths, probe_wavelength, n, Te, np.array([5, 5]) * u.eV, ifract=ifract, ion_species=ion_species, probe_vec=probe_vec, scatter_vec=scatter_vec, ) # Raise a ValueError with empty ion_species with pytest.raises(ValueError): alpha, Skw = thomson.spectral_density( wavelengths, probe_wavelength, n, Te, Ti, ifract=ifract, ion_species=[], probe_vec=probe_vec, scatter_vec=scatter_vec, ) # Raise a Value Error with inconsistent electron array lengths # Te.size != efract.size with pytest.raises(ValueError): alpha, Skw = thomson.spectral_density( wavelengths, probe_wavelength, n, np.array([1, 10]) * u.eV, Ti, efract=np.array([0.5, 0.2, 0.3]), probe_vec=probe_vec, scatter_vec=scatter_vec, ) # Electron vel shape not compatible with efract.size with pytest.raises(ValueError): alpha, Skw = thomson.spectral_density( wavelengths, probe_wavelength, n, Te, Ti, efract=np.array([0.5, 0.5]), electron_vel=np.array([[100, 0, 0]]) * u.km / u.s, probe_vec=probe_vec, scatter_vec=scatter_vec, )