def abundances(self, abundances_dict: Optional[Dict]):
"""
Set the elemental (or isotopic) abundances. The elements and
isotopes must be the same as or a superset of the elements whose
ionization states are being tracked.
"""
if abundances_dict is None:
self._pars['abundances'] = {
elem: np.nan
for elem in self.base_particles
}
elif not isinstance(abundances_dict, dict):
raise TypeError(f"The abundances attribute must be a dict with "
f"elements or isotopes as keys and real numbers "
f"representing relative abundances as values.")
else:
old_keys = abundances_dict.keys()
try:
new_keys_dict = {
particle_symbol(old_key): old_key
for old_key in old_keys
}
except Exception:
raise AtomicError(
f"The key {repr(old_key)} in the abundances "
f"dictionary is not a valid element or isotope.")
new_elements = new_keys_dict.keys()
old_elements_set = set(self.base_particles)
new_elements_set = set(new_elements)
if old_elements_set - new_elements_set:
raise AtomicError(
f"The abundances of the following particles are "
f"missing: {old_elements_set - new_elements_set}")
new_abundances_dict = {}
for element in new_elements:
inputted_abundance = abundances_dict[new_keys_dict[element]]
try:
inputted_abundance = float(inputted_abundance)
except Exception:
raise TypeError(
f"The abundance for {element} was provided as"
f"{inputted_abundance}, which cannot be "
f"converted to a real number.") from None
if inputted_abundance < 0:
raise AtomicError(
f"The abundance of {element} is negative.")
new_abundances_dict[element] = inputted_abundance
self._pars['abundances'] = new_abundances_dict
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_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_getitem(self, test_name):
"""
Test that `IonizationState.__getitem__` returns the same value
when using equivalent keys (integer charge, particle symbol, and
`Particle` instance).
For example, if we create
>>> He_states = IonizationState('He', [0.2, 0.3, 0.5])
then this checks to make sure that `He_states[2]`,
`He_states['He 2+']`, and `He_states[Particle('He 2+')]` all
return the same result.
"""
instance = self.instances[test_name]
particle_name = instance.base_particle
integer_charges = np.arange(instance.atomic_number + 1)
symbols = [
particle_symbol(particle_name, Z=Z) for Z in integer_charges
]
particles = instance._particle_instances
errors = []
# In the following loop, instance[key] will return a namedtuple
# or class which may contain Quantity objects with values of
# numpy.nan. Because of the difficulty of comparing nans in
# these objects, we compare the string representations instead
# (see Astropy issue #7901 on GitHub).
for keys in zip(integer_charges, symbols, particles):
set_of_str_values = {str(instance[key]) for key in keys}
if len(set_of_str_values) != 1:
errors.append(
f"\n\n"
f"The following keys in test '{test_name}' did not "
f"produce identical outputs as required: {keys}. "
f"The set containing string representations of"
f"the values is:\n\n{set_of_str_values}")
if errors:
pytest.fail(str.join("", errors))
def test_getitem_two_indices(self, indices):
instance = self.instance
result = instance[indices]
particle = indices[0]
integer_charge = indices[1]
assert isinstance(result, State)
assert result.integer_charge == integer_charge
expected_ionic_fraction = instance.ionic_fractions[particle][integer_charge]
assert np.any([
np.isclose(result.ionic_fraction, expected_ionic_fraction),
np.isnan(result.ionic_fraction) and np.isnan(expected_ionic_fraction),
])
assert result.ionic_symbol == particle_symbol(particle, Z=integer_charge)
def test_getitem_two_indices(self, indices):
instance = self.instance
result = instance[indices]
particle = indices[0]
charge_number = indices[1]
assert isinstance(result, IonicLevel)
assert result.charge_number == charge_number
expected_ionic_fraction = instance.ionic_fractions[particle][charge_number]
assert np.any(
[
np.isclose(result.ionic_fraction, expected_ionic_fraction),
np.isnan(result.ionic_fraction) and np.isnan(expected_ionic_fraction),
]
)
assert result.ionic_symbol == particle_symbol(particle, Z=charge_number)
def ionic_fractions(self, inputs: Union[Dict, List, Tuple]):
"""
Set the ionic fractions.
Notes
-----
The ionic fractions are initialized during instantiation of
`~plasmapy.particles.IonizationStates`. After this, the only way
to reset the ionic fractions via the ``ionic_fractions``
attribute is via a `dict` with elements or isotopes that are a
superset of the previous elements or isotopes. However, you may
use item assignment of the `~plasmapy.particles.IonizationState`
instance to assign new ionic fractions one element or isotope
at a time.
Raises
------
AtomicError
If the ionic fractions cannot be set.
TypeError
If ``inputs`` is not a `list`, `tuple`, or `dict` during
instantiation, or if ``inputs`` is not a `dict` when it is
being set.
"""
# A potential problem is that using item assignment on the
# ionic_fractions attribute could cause the original attributes
# to be overwritten without checks being performed. We might
# eventually want to create a new class or subclass of UserDict
# that goes through these checks. In the meantime, we should
# make it clear to users to set ionic_fractions by using item
# assignment on the IonizationStates instance as a whole. An
# example of the problem is `s = IonizationStates(["He"])` being
# followed by `s.ionic_fractions["He"] = 0.3`.
if hasattr(self, '_ionic_fractions'):
if not isinstance(inputs, dict):
raise TypeError(
"Can only reset ionic_fractions with a dict if "
"ionic_fractions has been set already.")
old_particles = set(self.base_particles)
new_particles = {particle_symbol(key) for key in inputs.keys()}
missing_particles = old_particles - new_particles
if missing_particles:
raise AtomicError(
"Can only reset ionic fractions with a dict if "
"the new base particles are a superset of the "
"prior base particles. To change ionic fractions "
"for one base particle, use item assignment on the "
"IonizationStates instance instead.")
if isinstance(inputs, dict):
original_keys = inputs.keys()
ionfrac_types = {type(inputs[key]) for key in original_keys}
inputs_have_quantities = u.Quantity in ionfrac_types
if inputs_have_quantities and len(ionfrac_types) != 1:
raise TypeError(
"Ionic fraction information may only be inputted "
"as a Quantity object if all ionic fractions are "
"Quantity arrays with units of inverse volume.")
try:
particles = {key: Particle(key) for key in original_keys}
except (InvalidParticleError, TypeError) as exc:
raise AtomicError(
"Unable to create IonizationStates instance "
"because not all particles are valid.") from exc
# The particles whose ionization states are to be recorded
# should be elements or isotopes but not ions or neutrals.
for key in particles.keys():
is_element = particles[key].is_category('element')
has_charge_info = particles[key].is_category(
any_of=["charged", "uncharged"])
if not is_element or has_charge_info:
raise AtomicError(
f"{key} is not an element or isotope without "
f"charge information.")
# We are sorting the elements/isotopes by atomic number and
# mass number since we will often want to plot and analyze
# things and this is the most sensible order.
sorted_keys = sorted(original_keys,
key=lambda k: (
particles[k].atomic_number,
particles[k].mass_number
if particles[k].isotope else 0,
))
_elements_and_isotopes = []
_particle_instances = []
new_ionic_fractions = {}
if inputs_have_quantities:
n_elems = {}
for key in sorted_keys:
new_key = particles[key].particle
_particle_instances.append(particles[key])
if new_key in _elements_and_isotopes:
raise AtomicError(
"Repeated particles in IonizationStates.")
nstates_input = len(inputs[key])
nstates = particles[key].atomic_number + 1
if nstates != nstates_input:
raise AtomicError(
f"The ionic fractions array for {key} must "
f"have a length of {nstates}.")
_elements_and_isotopes.append(new_key)
if inputs_have_quantities:
try:
number_densities = inputs[key].to(u.m**-3)
n_elem = np.sum(number_densities)
new_ionic_fractions[new_key] = np.array(
number_densities / n_elem)
n_elems[key] = n_elem
except u.UnitConversionError as exc:
raise AtomicError(
"Units are not inverse volume.") from exc
elif isinstance(inputs[key],
np.ndarray) and inputs[key].dtype.kind == 'f':
new_ionic_fractions[particles[key].particle] = inputs[key]
else:
try:
new_ionic_fractions[particles[key].particle] = \
np.array(inputs[key], dtype=np.float)
except ValueError as exc:
raise AtomicError(
f"Inappropriate ionic fractions for {key}."
) from exc
for key in _elements_and_isotopes:
fractions = new_ionic_fractions[key]
if not np.all(np.isnan(fractions)):
if np.min(fractions) < 0 or np.max(fractions) > 1:
raise AtomicError(
f"Ionic fractions for {key} are not between 0 and 1."
)
if not np.isclose(
np.sum(fractions), 1, atol=self.tol, rtol=0):
raise AtomicError(
f"Ionic fractions for {key} are not normalized to 1."
)
# When the inputs provide the densities, the abundances must
# not have been provided because that would be redundant
# or contradictory information. The number density scaling
# factor might or might not have been provided. Have the
# number density scaling factor default to the total number
# of neutrals and ions across all elements and isotopes, if
# it was not provided. Then go ahead and calculate the
# abundances based on that. However, we need to be careful
# that the abundances are not overwritten during the
# instantiation of the class.
if inputs_have_quantities:
if np.isnan(self.n):
new_n = 0 * u.m**-3
for key in _elements_and_isotopes:
new_n += n_elems[key]
self.n = new_n
new_abundances = {}
for key in _elements_and_isotopes:
new_abundances[key] = np.float(n_elems[key] / self.n)
self._pars['abundances'] = new_abundances
elif isinstance(inputs, (list, tuple)):
try:
_particle_instances = [
Particle(particle) for particle in inputs
]
except (InvalidParticleError, TypeError) as exc:
raise AtomicError(
"Invalid inputs to IonizationStates.") from exc
_particle_instances.sort(key=lambda p: (
p.atomic_number, p.mass_number if p.isotope else 0))
_elements_and_isotopes = [
particle.particle for particle in _particle_instances
]
new_ionic_fractions = {
particle.particle: np.full(particle.atomic_number + 1,
fill_value=np.nan,
dtype=np.float64)
for particle in _particle_instances
}
else:
raise TypeError("Incorrect inputs to set ionic_fractions.")
for i in range(1, len(_particle_instances)):
if _particle_instances[
i - 1].element == _particle_instances[i].element:
if not _particle_instances[
i - 1].isotope and _particle_instances[i].isotope:
raise AtomicError(
"Cannot have an element and isotopes of that element.")
self._particle_instances = _particle_instances
self._base_particles = _elements_and_isotopes
self._ionic_fractions = new_ionic_fractions
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[:]
