def test_getitem_element(self, test):
"""Test that __get_item__ returns an IonizationState instance"""
instance = self.instances[test]
for key in instance.elements:
try:
expected = instance.ionic_fractions[key]
except Exception as exc:
raise AtomicError(
f"Unable to get ionic_fractions for '{key}' in "
f"test='{test}'.") from exc
try:
actual = instance[key].ionic_fractions
except Exception as exc:
raise AtomicError(f"Unable to get item {key} in test={test}.")
try:
if all(np.isnan(expected)):
test_passed = True
else:
test_passed = np.allclose(expected, actual)
except Exception:
raise TypeError(
f"For test='{test}' and key='{key}', cannot "
f"compare expected ionic fractions of {expected} "
f"with the resulting ionic fractions of {actual}."
) from None
if not test_passed:
raise AtomicError(
f"For test='{test}' and key='{key}', the expected "
f"ionic fractions of {expected} are not all equal "
f"to the resulting ionic fractions of {actual}.")
def __eq__(self, other):
"""
Return `True` if the ionic fractions for two `IonizationState`
instances are approximately equal to within the minimum `tol`
specified by either, and `False` otherwise.
Raises
------
AtomicError
If `other` is not an `~plasmapy.atomic.IonizationState`
instance, or if `other` corresponds to a different element.
Examples
--------
>>> IonizationState('H', [1, 0], tol=1e-6) == IonizationState('H', [1, 1e-6], tol=1e-6)
True
>>> IonizationState('H', [1, 0], tol=1e-8) == IonizationState('H', [1, 1e-6], tol=1e-5)
False
"""
if not isinstance(other, IonizationState):
raise AtomicError(
"Instances of the IonizationState class may only be "
"compared with other IonizationState instances.")
if self.element != other.element:
raise AtomicError(
"Only ionization states of the same element may be compared.")
# Use the tightest of the two absolute tolerances
min_tol = np.min([self.tol, other.tol])
return np.allclose(self.ionic_fractions,
other.ionic_fractions,
atol=min_tol)
def __init__(self,
particle: Particle,
ionic_fractions=None,
*,
T_e: u.K = np.nan * u.K,
kappa: Real = np.inf,
n_elem: u.m**-3 = np.nan * u.m**-3,
tol: Union[float, int] = 1e-15):
"""Initialize an `~plasmapy.atomic.IonizationState` instance."""
self._particle_instance = particle
try:
self.tol = tol
self.T_e = T_e
self.kappa = kappa
if not np.isnan(n_elem) and isinstance(ionic_fractions, u.Quantity) and \
ionic_fractions.si.unit == u.m ** -3:
raise AtomicError(
"Cannot simultaneously provide number density "
"through both n_elem and ionic_fractions.")
self.n_elem = n_elem
self.ionic_fractions = ionic_fractions
if ionic_fractions is None and not np.isnan(self.T_e):
warnings.warn(
"Collisional ionization equilibration has not yet "
"been implemented in IonizationState; cannot set "
"ionic fractions.")
except Exception as exc:
raise AtomicError(f"Unable to create IonizationState instance for "
f"{particle.particle}.") from exc
def __setitem__(self, key, value):
if isinstance(value, dict):
raise NotImplementedError("Dictionary assignment not implemented.")
else:
try:
particle = particle_symbol(key)
if particle not in self.elements:
raise AtomicError(
f"{key} is not one of the particles kept track of "
f"by this IonizationStates instance.")
new_fractions = np.array(value, dtype=np.float64)
if new_fractions.min() < 0 or new_fractions.max() > 1:
raise ValueError(
"Ionic fractions must be between 0 and 1.")
if not np.isclose(np.sum(new_fractions), 1):
raise ValueError("Ionic fractions are not normalized.")
if len(new_fractions) != atomic_number(particle) + 1:
raise ValueError(
f"Incorrect size of ionic fraction array for {key}.")
self._ionic_fractions[particle][:] = new_fractions[:]
except Exception as exc:
raise AtomicError(
f"Cannot set item for this IonizationStates "
f"instance for key = {repr(key)} and value = "
f"{repr(value)}")
def test_iteration(self, test_name: str):
"""Test that IonizationState instances iterate impeccably."""
try:
states = [state for state in self.instances[test_name]]
except Exception:
raise AtomicError(f"Unable to perform iteration for {test_name}.")
try:
integer_charges = [state.integer_charge for state in states]
ionic_fractions = np.array(
[state.ionic_fraction for state in states])
ionic_symbols = [state.ionic_symbol for state in states]
except Exception:
raise AtomicError("An attribute may be misnamed or missing.")
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(integer_charges == expected_charges):
errors.append(
f"The resulting integer charges are {integer_charges}, "
f"which are not equal to the expected integer charges, "
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 integer_charges
]
expected_symbols = [
particle.ionic_symbol for particle in expected_particles
]
if not 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)
raise AtomicError(errmsg)
def T_e(self, value: u.K):
"""Set the electron temperature."""
try:
value = value.to(u.K, equivalencies=u.temperature_energy())
except (AttributeError, u.UnitsError, u.UnitConversionError):
raise AtomicError("Invalid temperature.") from None
else:
if value < 0 * u.K:
raise AtomicError("T_e cannot be negative.")
self._T_e = value
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 n_H(self, n):
if n is None:
self._pars['n_H'] = n
else:
try:
self._pars['n_H'] = n.to(u.m**-3)
except u.UnitConversionError:
raise AtomicError("Units cannot be converted to u.m**-3.")
except Exception:
raise AtomicError(
f"{n} is not a valid number density.") from None
def number_densities(self, value: u.m**-3):
"""Set the number densities for each state."""
if np.any(value.value < 0):
raise AtomicError("Number densities cannot be negative.")
if len(value) != self.atomic_number + 1:
raise AtomicError(f"Incorrect number of charge states for "
f"{self.base_particle}")
value = value.to(u.m**-3)
self._n_elem = value.sum()
self._ionic_fractions = value / self._n_elem
def T_e(self, electron_temperature: u.K):
"""Set the electron temperature."""
try:
temperature = electron_temperature.to(
u.K, equivalencies=u.temperature_energy())
except (AttributeError, u.UnitsError):
raise AtomicError(
f"{electron_temperature} is not a valid temperature."
) from None
if temperature < 0 * u.K:
raise AtomicError("The electron temperature cannot be negative.")
self._pars['T_e'] = temperature
def n(self, n: u.m**-3):
"""Set the number density scaling factor."""
try:
n = n.to(u.m**-3)
except u.UnitConversionError as exc:
raise AtomicError(
"Units cannot be converted to u.m ** -3.") from exc
except Exception as exc:
raise AtomicError(f"{n} is not a valid number density.") from exc
if n < 0 * u.m**-3:
raise AtomicError("Number density cannot be negative.")
self._pars['n'] = n.to(u.m**-3)
def n_elem(self, value):
"""The number density of atoms plus ions of this species."""
if value is None:
self._n_elem = None
else:
if '_n_e' in dir(self) and self._n_e is not None:
raise AtomicError(
"Only one of n_e and n_elem may be set for a "
"single element, quasineutral plasma.")
try:
self._n_elem = value.to(u.m**-3)
except (AttributeError, u.UnitConversionError):
raise AtomicError(_number_density_errmsg) from None
def T_e(self, electron_temperature):
if electron_temperature is None:
self._pars['T_e'] = None
else:
try:
temp = electron_temperature.to(
u.K, equivalencies=u.temperature_energy())
except (AttributeError, u.UnitsError):
raise AtomicError("Invalid electron temperature.")
else:
if temp < 0 * u.K:
raise AtomicError(
"The electron temperature cannot be negative.")
self._pars['T_e'] = temp
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'] = None
elif not isinstance(abundances_dict, dict):
raise TypeError(
f"The abundances argument {abundances_dict} must be a dict with elements "
"or isotopes as keys and ")
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(
"The key {repr(old_key)} in the abundances "
"dictionary is not a valid element or isotope.")
new_elements = new_keys_dict.keys()
old_elements_set = set(self.elements)
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
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 number_densities(self):
"""Return the number densities for each state."""
if self._n_e is not None or self._n_elem is not None:
return (self.n_elem * self.ionic_fractions).to(u.m**-3)
else:
raise AtomicError(
"Insufficient information to return number densities.")
def antiparticle(self):
"""
Return the corresponding antiparticle, or raise an
`~plasmapy.utils.AtomicError` if the particle is not an
elementary particle.
This attribute may be accessed by using the unary operator ``~``
acting on a `~plasma.atomic.Particle` instance.
Examples
--------
>>> electron = Particle('e-')
>>> electron.antiparticle
Particle("e+")
>>> antineutron = Particle('antineutron')
>>> ~antineutron
Particle("n")
"""
if self.particle in _antiparticles.keys():
return Particle(_antiparticles[self.particle])
else:
raise AtomicError(
"The unary operator can only be used for elementary "
"particles and antiparticles.")
def n_e(self, value):
"""
Return the electron density assuming a single-species
plasma.
"""
if value is None:
self._n_e = None
return
elif self._n_elem is not None:
raise AtomicError("Only one of n_e and n_elem may be set for a "
"single element, quasineutral plasma.")
try:
self._n_e = value.to(u.m**-3)
self._n_elem = None
except (AttributeError, u.UnitConversionError):
raise AtomicError(_number_density_errmsg)
def __getitem__(self, value) -> State:
"""Return the ionic fraction(s)."""
if isinstance(value, slice):
raise TypeError("IonizationState instances cannot be sliced.")
if isinstance(value,
(int, np.integer)) and 0 <= value <= self.atomic_number:
result = State(value, self.ionic_fractions[value],
self.ionic_symbols[value])
else:
if not isinstance(value, Particle):
try:
value = Particle(value)
except InvalidParticleError as exc:
raise InvalidParticleError(
f"{value} is not a valid integer charge or "
f"particle.") from exc
same_element = value.element == self.element
same_isotope = value.isotope == self.isotope
has_charge_info = value.is_category(
any_of=["charged", "uncharged"])
if same_element and same_isotope and has_charge_info:
Z = value.integer_charge
result = State(Z, self.ionic_fractions[Z],
self.ionic_symbols[Z])
else:
if not same_element or not same_isotope:
raise AtomicError("Inconsistent element or isotope.")
elif not has_charge_info:
raise ChargeError("No integer charge provided.")
return result
def __init__(self,
particle: Particle,
ionic_fractions=None,
*,
T_e=None,
n_e=None,
n_elem=None,
tol: Union[float, int] = 1e-15):
"""Initialize a `~plasmapy.atomic.IonizationState` instance."""
self._particle = particle
try:
self.tol = tol
self.T_e = T_e
self.n_elem = n_elem
self.n_e = n_e
self.ionic_fractions = ionic_fractions
# This functionality has not yet been implemented:
# if self._ionic_fractions is None and self.T_e is not None:
# self.equilibrate()
except Exception as exc:
raise AtomicError(f"Unable to create IonizationState instance for "
f"{particle.particle}.") from exc
def process_particles_list(
unformatted_particles_list: List[Union[str, Particle]]) \
-> List[Particle]:
"""
Take an unformatted list of particles and puts each
particle into standard form, while allowing an integer and
asterisk immediately preceding a particle to act as a
multiplier. A string argument will be treated as a list
containing that string as its sole item.
"""
if isinstance(unformatted_particles_list, str):
unformatted_particles_list = [unformatted_particles_list]
if not isinstance(unformatted_particles_list, (list, tuple)):
raise TypeError("The input to process_particles_list should be a "
"string, list, or tuple.")
particles = []
for original_item in unformatted_particles_list:
try:
item = original_item.strip()
if item.count('*') == 1 and item[0].isdigit():
multiplier_str, item = item.split('*')
multiplier = int(multiplier_str)
else:
multiplier = 1
try:
particle = Particle(item)
except (InvalidParticleError) as exc:
raise AtomicError(errmsg) from exc
if particle.element and not particle.isotope:
raise AtomicError(errmsg)
[particles.append(particle) for i in range(multiplier)]
except Exception:
raise AtomicError(
f"{original_item} is not a valid reactant or "
"product in a nuclear reaction.") from None
return particles
def ionic_fractions(self, fractions):
"""
Set the ionic fractions, while checking that the new values are
valid and normalized to one.
"""
if fractions is None:
self._ionic_fractions = np.full(self.atomic_number + 1,
np.nan,
dtype=np.float64)
else:
try:
if not isinstance(fractions, np.ndarray) or 'float' not in str(
fractions.dtype):
fractions = np.array(fractions, dtype=np.float)
except Exception as exc:
raise AtomicError(
f"Unable to set ionic fractions of {self.element} "
f"to {fractions}.") from exc
if np.min(fractions) < 0:
raise AtomicError("Ionic fractions cannot be negative.")
if isinstance(fractions, u.Quantity):
if self._n_e is not None or self._n_elem is not None:
raise AtomicError(
"The ionization state may be set using number "
"densities for each ion only if neither of the "
"electron density and element density has "
"already been set.")
self.number_densities = fractions
else:
if not np.any(np.isnan(fractions)):
total = np.sum(fractions)
if not np.isclose(total, 1, atol=self.tol, rtol=0):
raise AtomicError(
f"The sum of the ionic fractions of {self.element} "
f"equals {total}, which is not approximately one.")
if not len(fractions) == self.atomic_number + 1:
raise AtomicError(
f"len(fractions) equals {len(fractions)}, but "
f"should equal {self.atomic_number + 1} which "
f"is the atomic number of {self.element} + 1.")
self._ionic_fractions = fractions
def test_ionic_fractions(self, test):
errmsg = ""
elements_actual = self.instances[test].elements
inputs = tests[test]["inputs"]
if isinstance(inputs, dict):
input_keys = tests[test]["inputs"].keys()
for element, input_key in zip(elements_actual, input_keys):
expected = np.array(tests[test]["inputs"][input_key])
if isinstance(expected, u.Quantity):
expected = np.array(expected.value /
np.sum(expected.value))
#if not isinstance(expected, np.ndarray)
actual = self.instances[test].ionic_fractions[element]
if not np.allclose(actual, expected):
errmsg += (
f"\n\nThere is a discrepancy in ionic fractions for "
f"({test}, {element}, {input_key})\n"
f" expected = {expected}\n"
f" actual = {actual}")
if not isinstance(actual, np.ndarray) or isinstance(
actual, u.Quantity):
raise AtomicError(
f"\n\nNot a numpy.ndarray: ({test}, {element})")
else:
elements_expected = {
particle_symbol(element)
for element in inputs
}
assert set(self.instances[test].elements) == elements_expected
for element in elements_expected:
assert all(
np.isnan(self.instances[test].ionic_fractions[element]))
if errmsg:
raise AtomicError(errmsg)
def n_H(self):
"""
The number density of hydrogen neutrals and atoms of all
isotopes, if defined.
"""
if 'H' not in self.elements or self._pars['n_H'] is None:
raise AtomicError("The number density of hydrogen is not ")
return self._pars['n_H']
def log_abundances(self):
if self._pars['abundances'] is not None:
log_abundances_dict = {}
for key in self.abundances.keys():
log_abundances_dict[key] = np.log10(self.abundances[key])
return log_abundances_dict
else:
raise AtomicError("No abundances are available.")
def __init__(self,
inputs: Union[Dict[str, np.ndarray], List, Tuple],
*,
T_e: u.K = np.nan * u.K,
equilibrate: Optional[bool] = None,
abundances: Optional[Dict[str, Real]] = None,
log_abundances: Optional[Dict[str, Real]] = None,
n: u.m**-3 = np.nan * u.m**-3,
tol: Real = 1e-15,
kappa: Real = np.inf):
abundances_provided = abundances is not None or log_abundances is not None
set_abundances = True
if isinstance(inputs, dict):
all_quantities = np.all(
[isinstance(fracs, u.Quantity) for fracs in inputs.values()])
if all_quantities:
right_units = np.all(
[fracs[0].si.unit == u.m**-3 for fracs in inputs.values()])
if not right_units:
raise AtomicError(
"Units must be inverse volume for number densities.")
if abundances_provided:
raise AtomicError(
"Abundances cannot be provided if inputs "
"provides number density information.")
set_abundances = False
try:
self._pars = collections.defaultdict(lambda: None)
self.T_e = T_e
self.n = n
self.tol = tol
self.ionic_fractions = inputs
if set_abundances:
self.abundances = abundances
self.log_abundances = log_abundances
self.kappa = kappa
except Exception as exc:
raise AtomicError(
"Unable to create IonizationStates instance.") from exc
if equilibrate:
self.equilibrate() # for now, this raises a NotImplementedError
def __eq__(self, other):
if not isinstance(other, IonizationStates):
raise TypeError(
"IonizationStates instance can only be compared with "
"other IonizationStates instances.")
if self.base_particles != other.base_particles:
raise AtomicError(
"Two IonizationStates instances can be compared only "
"if the base particles are the same.")
min_tol = np.min([self.tol, other.tol])
# Check any of a whole bunch of equality measures, recalling
# that np.nan == np.nan is False.
for attribute in ['T_e', 'n_e', 'kappa']:
this = eval(f"self.{attribute}")
that = eval(f"other.{attribute}")
# TODO: Maybe create a function in utils called same_enough
# TODO: that would take care of all of these disparate
# TODO: equality measures.
this_equals_that = np.any([
this == that,
this is that,
np.isnan(this) and np.isnan(that),
np.isinf(this) and np.isinf(that),
u.quantity.allclose(this, that, rtol=min_tol),
])
if not this_equals_that:
return False
for attribute in ['ionic_fractions', 'number_densities']:
this_dict = eval(f"self.{attribute}")
that_dict = eval(f"other.{attribute}")
for particle in self.base_particles:
this = this_dict[particle]
that = that_dict[particle]
this_equals_that = np.any([
this is that,
np.all(np.isnan(this)) and np.all(np.isnan(that)),
u.quantity.allclose(this, that, rtol=min_tol),
])
if not this_equals_that:
return False
return True
def log_abundances(self, value):
if value is not None:
try:
new_abundances_input = {}
for key in value.keys():
new_abundances_input[key] = 10**value[key]
self.abundances = new_abundances_input
except Exception as exc:
raise AtomicError("Invalid log_abundances.")
def ionic_fractions(self, fractions):
"""
Set the ionic fractions, while checking that the new values are
valid and normalized to one.
"""
if fractions is None or np.all(np.isnan(fractions)):
self._ionic_fractions = np.full(self.atomic_number + 1,
np.nan,
dtype=np.float64)
return
try:
if np.min(fractions) < 0:
raise AtomicError("Cannot have negative ionic fractions.")
if len(fractions) != self.atomic_number + 1:
raise AtomicError("The length of ionic_fractions must be "
f"{self.atomic_number + 1}.")
if isinstance(fractions, u.Quantity):
fractions = fractions.to(u.m**-3)
self.n_elem = np.sum(fractions)
self._ionic_fractions = np.array(fractions / self.n_elem)
else:
fractions = np.array(fractions, dtype=np.float64)
sum_of_fractions = np.sum(fractions)
all_nans = np.all(np.isnan(fractions))
if not all_nans:
if np.any(fractions < 0) or np.any(fractions > 1):
raise AtomicError(
"Ionic fractions must be between 0 and 1.")
if not np.isclose(
sum_of_fractions, 1, rtol=0, atol=self.tol):
raise AtomicError("Ionic fractions must sum to one.")
self._ionic_fractions = fractions
except Exception as exc:
raise AtomicError(
f"Unable to set ionic fractions of {self.element} "
f"to {fractions}.") from exc
def T_e(self, value):
if value is None:
self._T_e = None
else:
try:
value = value.to(u.K, equivalencies=u.temperature_energy())
except (AttributeError, u.UnitsError):
raise AtomicError("Invalid temperature.") from None
self._T_e = value
def log_abundances(self, value: Optional[Dict[str, Real]]):
"""
Set the base 10 logarithm of the relative abundances.
"""
if value is not None:
try:
new_abundances_input = {}
for key in value.keys():
new_abundances_input[key] = 10**value[key]
self.abundances = new_abundances_input
except Exception:
raise AtomicError("Invalid log_abundances.") from None