def test_iteration(self, test_name: str):
"""Test that `IonizationState` instances iterate impeccably."""
states = list(self.instances[test_name])
charge_numbers = [state.charge_number for state in states]
ionic_fractions = np.array([state.ionic_fraction for state in states])
ionic_symbols = [state.ionic_symbol for state in states]
try:
base_symbol = isotope_symbol(ionic_symbols[0])
except InvalidIsotopeError:
base_symbol = atomic_symbol(ionic_symbols[0])
finally:
atomic_numb = atomic_number(ionic_symbols[1])
errors = []
expected_charges = np.arange(atomic_numb + 1)
if not np.all(charge_numbers == expected_charges):
errors.append(
f"The resulting charge numbers are {charge_numbers}, "
f"which are not equal to the expected charge numbers, "
f"which are {expected_charges}."
)
expected_fracs = test_cases[test_name]["ionic_fractions"]
if isinstance(expected_fracs, u.Quantity):
expected_fracs = (expected_fracs / expected_fracs.sum()).value
if not np.allclose(ionic_fractions, expected_fracs):
errors.append(
f"The resulting ionic fractions are {ionic_fractions}, "
f"which are not equal to the expected ionic fractions "
f"of {expected_fracs}."
)
expected_particles = [
Particle(base_symbol, Z=charge) for charge in charge_numbers
]
expected_symbols = [particle.ionic_symbol for particle in expected_particles]
if ionic_symbols != expected_symbols:
errors.append(
f"The resulting ionic symbols are {ionic_symbols}, "
f"which are not equal to the expected ionic symbols of "
f"{expected_symbols}."
)
if errors:
errors.insert(
0,
(
f"The test of IonizationState named '{test_name}' has "
f"resulted in the following errors when attempting to "
f"iterate."
),
)
errmsg = " ".join(errors)
pytest.fail(errmsg)
def test_attribute_defaults_to_dict_of_nans(self, uninitialized_attribute):
command = f"self.instance.{uninitialized_attribute}"
default_value = eval(command)
assert list(default_value.keys()) == self.elements, "Incorrect base particle keys."
for element in self.elements:
assert len(default_value[element]) == atomic_number(element) + 1, \
f"Incorrect number of ionization levels for {element}."
assert np.all(np.isnan(default_value[element])), (
f"The values do not default to an array of nans for "
f"{element}.")
def test_getitem_element_intcharge(self, test_name):
instance = self.instances[test_name]
for particle in instance.base_particles:
for int_charge in range(0, atomic_number(particle) + 1):
actual = instance[particle, int_charge].ionic_fraction
expected = instance.ionic_fractions[particle][int_charge]
# We only need to check if one is broken
if not np.isnan(actual) and np.isnan(expected):
assert np.isclose(actual,
expected), (f"Indexing broken for:\n"
f" test = '{test_name}'\n"
f" particle = '{particle}'")
def test_that_ionic_fractions_are_set_correctly(self, test_name):
errmsg = ""
elements_actual = self.instances[test_name].base_particles
inputs = tests[test_name]["inputs"]
if isinstance(inputs, dict):
input_keys = list(tests[test_name]["inputs"].keys())
input_keys = sorted(
input_keys,
key=lambda k: (
atomic_number(k),
mass_number(k) if Particle(k).isotope else 0,
),
)
for element, input_key in zip(elements_actual, input_keys):
expected = tests[test_name]["inputs"][input_key]
if isinstance(expected, u.Quantity):
expected = np.array(expected.value /
np.sum(expected.value))
actual = self.instances[test_name].ionic_fractions[element]
if not np.allclose(actual, expected):
errmsg += (
f"\n\nThere is a discrepancy in ionic fractions for "
f"({test_name}, {element}, {input_key})\n"
f" expected = {expected}\n"
f" actual = {actual}")
if not isinstance(actual, np.ndarray) or isinstance(
actual, u.Quantity):
raise ParticleError(
f"\n\nNot a numpy.ndarray: ({test_name}, {element})")
else:
elements_expected = {
particle_symbol(element)
for element in inputs
}
assert set(
self.instances[test_name].base_particles) == elements_expected
for element in elements_expected:
assert all(
np.isnan(
self.instances[test_name].ionic_fractions[element]))
if errmsg:
pytest.fail(errmsg)
def test_nans():
"""
Test that when no ionic fractions or temperature are inputted,
the result is an array full of `~numpy.nan` of the right size.
"""
element = "He"
nstates = atomic_number(element) + 1
instance = IonizationState(element)
assert (len(instance.ionic_fractions) == nstates
), f"Incorrect number of ionization states for {element}"
assert np.all([
np.isnan(instance.ionic_fractions)
]), ("The ionic fractions for IonizationState are not defaulting "
"to numpy.nan when not set by user.")
def n_e(self) -> u.m**-3:
"""
Return the electron number density under the assumption of
quasineutrality.
"""
number_densities = self.number_densities
n_e = 0.0 * u.m**-3
for elem in self.base_particles:
atomic_numb = atomic_number(elem)
number_of_ionization_states = atomic_numb + 1
integer_charges = np.linspace(0, atomic_numb,
number_of_ionization_states)
n_e += np.sum(number_densities[elem] * integer_charges)
return n_e
def test_that_elements_and_isotopes_are_sorted(self, test_name):
elements = self.instances[test_name].base_particles
before_sorting = []
for element in elements:
atomic_numb = atomic_number(element)
try:
mass_numb = mass_number(element)
except InvalidIsotopeError:
mass_numb = 0
before_sorting.append((atomic_numb, mass_numb))
after_sorting = sorted(before_sorting)
assert before_sorting == after_sorting, (
f"Elements/isotopes are not sorted for test='{test_name}':\n"
f" before_sorting = {before_sorting}\n"
f" after_sorting = {after_sorting}\n"
f"where above is (atomic_number, mass_number if isotope else 0)")
def __getitem__(self, *values) -> IonizationState:
errmsg = f"Invalid indexing for IonizationStates instance: {values[0]}"
one_input = not isinstance(values[0], tuple)
two_inputs = len(values[0]) == 2
if not one_input and not two_inputs:
raise IndexError(errmsg)
try:
arg1 = values[0] if one_input else values[0][0]
int_charge = None if one_input else values[0][1]
particle = arg1 if arg1 in self.base_particles else particle_symbol(
arg1)
if int_charge is None:
return IonizationState(
particle=particle,
ionic_fractions=self.ionic_fractions[particle],
T_e=self._pars["T_e"],
n_elem=np.sum(self.number_densities[particle]),
tol=self.tol,
)
else:
if not isinstance(int_charge, Integral):
raise TypeError(
f"{int_charge} is not a valid charge for {base_particle}."
)
elif not 0 <= int_charge <= atomic_number(particle):
raise ChargeError(
f"{int_charge} is not a valid charge for {base_particle}."
)
return State(
integer_charge=int_charge,
ionic_fraction=self.ionic_fractions[particle][int_charge],
ionic_symbol=particle_symbol(particle, Z=int_charge),
number_density=self.number_densities[particle][int_charge],
)
except Exception as exc:
raise IndexError(errmsg) from exc
def test_that_iron_ionic_fractions_are_still_undefined(self):
assert "Fe" in self.instance.ionic_fractions.keys()
iron_fractions = self.instance.ionic_fractions["Fe"]
assert len(iron_fractions) == atomic_number("Fe") + 1
assert np.all(np.isnan(iron_fractions))
def sort_key(k):
return atomic_number(k), mass_number(k) if Particle(
k).isotope else 0
def __setitem__(self, key, value):
errmsg = (f"Cannot set item for this IonizationStates instance for "
f"key = {repr(key)} and value = {repr(value)}")
try:
particle = particle_symbol(key)
self.ionic_fractions[key]
except (AtomicError, TypeError):
raise KeyError(
f"{errmsg} because {repr(key)} is an invalid particle."
) from None
except KeyError:
raise KeyError(
f"{errmsg} because {repr(key)} is not one of the base "
f"particles whose ionization state is being kept track "
f"of.") from None
if isinstance(value,
u.Quantity) and value.unit != u.dimensionless_unscaled:
try:
new_number_densities = value.to(u.m**-3)
except u.UnitConversionError:
raise ValueError(f"{errmsg} because the units of value do not "
f"correspond to a number density.") from None
old_n_elem = np.sum(self.number_densities[particle])
new_n_elem = np.sum(new_number_densities)
density_was_nan = np.all(np.isnan(self.number_densities[particle]))
same_density = u.quantity.allclose(old_n_elem,
new_n_elem,
rtol=self.tol)
if not same_density and not density_was_nan:
raise ValueError(
f"{errmsg} because the old element number density "
f"of {old_n_elem} is not approximately equal to "
f"the new element number density of {new_n_elem}.")
value = (new_number_densities / new_n_elem).to(
u.dimensionless_unscaled)
# If the abundance of this particle has not been defined,
# then set the abundance if there is enough (but not too
# much) information to do so.
abundance_is_undefined = np.isnan(self.abundances[particle])
isnan_of_abundance_values = np.isnan(list(
self.abundances.values()))
all_abundances_are_nan = np.all(isnan_of_abundance_values)
n_is_defined = not np.isnan(self.n)
if abundance_is_undefined:
if n_is_defined:
self._pars['abundances'][particle] = new_n_elem / self.n
elif all_abundances_are_nan:
self.n = new_n_elem
self._pars['abundances'][particle] = 1
else:
raise AtomicError(
f"Cannot set number density of {particle} to "
f"{value * new_n_elem} when the number density "
f"scaling factor is undefined, the abundance "
f"of {particle} is undefined, and some of the "
f"abundances of other elements/isotopes is "
f"defined.")
try:
new_fractions = np.array(value, dtype=np.float64)
except Exception as exc:
raise TypeError(
f"{errmsg} because value cannot be converted into an "
f"array that represents ionic fractions.") from exc
# TODO: Create a separate function that makes sure ionic
# TODO: fractions are valid to reduce code repetition. This
# TODO: would probably best go as a private function in
# TODO: ionization_state.py.
required_nstates = atomic_number(particle) + 1
new_nstates = len(new_fractions)
if new_nstates != required_nstates:
raise ValueError(
f"{errmsg} because value must have {required_nstates} "
f"ionization levels but instead corresponds to "
f"{new_nstates} levels.")
all_nans = np.all(np.isnan(new_fractions))
if not all_nans and (new_fractions.min() < 0
or new_fractions.max() > 1):
raise ValueError(
f"{errmsg} because the new ionic fractions are not "
f"all between 0 and 1.")
normalized = np.isclose(np.sum(new_fractions), 1, rtol=self.tol)
if not normalized and not all_nans:
raise ValueError(f"{errmsg} because the ionic fractions are not "
f"normalized to one.")
self._ionic_fractions[particle][:] = new_fractions[:]