def n_e(self) -> u.m**-3:
"""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
charge_numbers = np.linspace(0, atomic_numb,
number_of_ionization_states)
n_e += np.sum(number_densities[elem] * charge_numbers)
return n_e
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 {particle}."
)
elif not 0 <= int_charge <= atomic_number(particle):
raise ChargeError(
f"{int_charge} is not a valid charge for {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 __setitem__(self, key, value):
errmsg = (
f"Cannot set item for this IonizationStateCollection instance for "
f"key = {repr(key)} and value = {repr(value)}"
)
try:
particle = particle_symbol(key)
self.ionic_fractions[key]
except (ParticleError, 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.n0)
if abundance_is_undefined:
if n_is_defined:
self._pars["abundances"][particle] = new_n_elem / self.n0
elif all_abundances_are_nan:
self.n0 = new_n_elem
self._pars["abundances"][particle] = 1
else:
raise ParticleError(
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 normalized to one."
)
self._ionic_fractions[particle][:] = new_fractions[:]
assert np.isclose(isotopic_abundance("D"), 0.000115)
assert isotopic_abundance("Be-8") == 0.0, "Be-8"
assert isotopic_abundance("Li-8") == 0.0, "Li-8"
with pytest.warns(ParticleWarning):
isotopic_abundance("Og", 294)
with pytest.raises(InvalidIsotopeError):
isotopic_abundance("neutron")
pytest.fail("No exception raised for neutrons.")
with pytest.raises(InvalidParticleError):
isotopic_abundance("Og-2")
isotopic_abundance_elements = (atomic_number(atomic_numb)
for atomic_numb in range(1, 119))
isotopic_abundance_isotopes = (common_isotopes(element)
for element in isotopic_abundance_elements)
isotopic_abundance_sum_table = ((
element, isotopes) for element, isotopes in zip(
isotopic_abundance_elements, isotopic_abundance_isotopes) if isotopes)
@pytest.mark.parametrize("element, isotopes", isotopic_abundance_sum_table)
def test_isotopic_abundances_sum(element, isotopes):
"""Test that the sum of isotopic abundances for each element with
isotopic abundances is one."""
sum_of_iso_abund = sum(isotopic_abundance(isotope) for isotope in isotopes)
