def test_getitem_one_index(self, index):
instance = self.instance
result = instance[index]
if np.all(np.isnan(instance.number_densities[index])):
inputs = instance.ionic_fractions[index]
else:
inputs = instance.number_densities[index]
expected = IonizationState(index, inputs, T_e=instance.T_e, kappa=instance.kappa)
assert isinstance(result, IonizationState)
assert result == expected
def test_average_ion_consistency(
base_particle,
ionic_fractions,
include_neutrals,
use_rms_mass,
use_rms_charge,
physical_type,
):
"""
Make sure that the average ions returned from equivalent `IonizationState`
and `IonizationStateCollection` instances are consistent with each other.
"""
ionization_state = IonizationState(base_particle, ionic_fractions)
ionization_state_collection = IonizationStateCollection(
{base_particle: ionic_fractions}
)
options = {
"include_neutrals": include_neutrals,
"use_rms_charge": use_rms_charge,
"use_rms_mass": use_rms_mass,
}
average_ion_from_ionization_state = ionization_state.average_ion(**options)
average_ion_from_ionization_state_collection = (
ionization_state_collection.average_ion(**options)
)
quantity_from_ion_state = getattr(average_ion_from_ionization_state, physical_type)
quantity_from_ion_collection = getattr(
average_ion_from_ionization_state_collection, physical_type
)
assert_quantity_allclose(
quantity_from_ion_state, quantity_from_ion_collection, rtol=1e-10
)
def __next__(self):
if self._element_index < len(self.base_particles):
particle = self.base_particles[self._element_index]
result = IonizationState(
particle,
self.ionic_fractions[particle],
T_e=self.T_e,
n_elem=np.sum(self.number_densities[particle]),
tol=self.tol,
)
self._element_index += 1
return result
else:
del self._element_index
raise StopIteration
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