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