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 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_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)]), (
f"The ionic fractions for IonizationState are not defaulting "
f"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_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 AtomicError(
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 __init__(
self,
initial: IonizationStates,
n_init: u.Quantity,
T_e_init: u.Quantity,
max_steps: int,
time_start: u.Quantity,
):
self._elements = list(initial.ionic_fractions.keys())
self._abundances = initial.abundances
self._max_steps = max_steps
self._nstates = {
elem: atomic_number(elem) + 1
for elem in self.elements
}
self._ionic_fractions = {
elem: np.full((max_steps + 1, self.nstates[elem]),
np.nan,
dtype=np.float64)
for elem in self.elements
}
self._number_densities = {
elem: np.full(
(max_steps + 1, self.nstates[elem]), np.nan, dtype=np.float64)
* u.cm**-3
for elem in self.elements
}
self._n_elem = {
elem: np.full(max_steps + 1, np.nan) * u.cm**-3
for elem in self.elements
}
self._n_e = np.full(max_steps + 1, np.nan) * u.cm**-3
self._T_e = np.full(max_steps + 1, np.nan) * u.K
self._time = np.full(max_steps + 1, np.nan) * u.s
self._index = 0
self._assign(
new_time=time_start,
new_ionfracs=initial.ionic_fractions,
new_n=n_init,
new_T_e=T_e_init,
)
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 __init__(self, element='H'):
"""Read in the """
self._element = element
self._temperature = None
#
# 1. Read ionization and recombination rates
#
data_dir = __path__[0] + '/data/ionizrecombrates/chianti_8.07/'
filename = data_dir + 'ionrecomb_rate.h5'
f = h5py.File(filename, 'r')
atomic_numb = atomic.atomic_number(element)
nstates = atomic_numb + 1
self._temperature_grid = f['te_gird'][:]
ntemp = len(self._temperature_grid)
c_ori = f['ioniz_rate'][:]
r_ori = f['recomb_rate'][:]
f.close()
#
# Ionization and recombination rate for the current element
#
c_rate = np.zeros((ntemp, nstates))
r_rate = np.zeros((ntemp, nstates))
for ite in range(ntemp):
for i in range(nstates-1):
c_rate[ite, i] = c_ori[i, atomic_numb-1, ite]
for i in range(1, nstates):
r_rate[ite, i] = r_ori[i-1, atomic_numb-1, ite]
#
# 2. Definet the grid size
#
self._ntemp = ntemp
self._atomic_numb = atomic_numb
self._nstates = nstates
#
# Compute eigenvalues and eigenvectors
#
self._ionization_rate = np.ndarray(shape=(ntemp, nstates),
dtype=np.float64)
self._recombination_rate = np.ndarray(shape=(ntemp, nstates),
dtype=np.float64)
self._equilibrium_states = np.ndarray(shape=(ntemp, nstates),
dtype=np.float64)
self._eigenvalues = np.ndarray(shape=(ntemp, nstates),
dtype=np.float64)
self._eigenvectors = np.ndarray(shape=(ntemp, nstates, nstates),
dtype=np.float64)
self._eigenvector_inverses = np.ndarray(
shape=(ntemp, nstates, nstates),
dtype=np.float64)
#
# Save ionization and recombination rates
#
self._ionization_rate = c_rate
self._recombination_rate = r_rate
#
# Define the coefficients matrix A. The first dimension is
# for elements, and the second number of equations.
#
neqs = nstates
A = np.ndarray(shape=(nstates, neqs), dtype=np.float64)
#
# Enter temperature loop over the whole temperature grid
#
for ite in range(ntemp):
# Ionization and recombination rate at Te(ite)
carr = c_rate[ite, :]
rarr = r_rate[ite, :]
# Equilibirum
eqi = self._function_eqi(carr, rarr, atomic_numb)
# Initialize A to zero
for ion in range(nstates):
for jon in range(nstates):
A[ion, jon] = 0.0
# Give coefficients
for ion in range(1, nstates-1):
A[ion, ion-1] = carr[ion-1]
A[ion, ion] = -(carr[ion]+rarr[ion])
A[ion, ion+1] = rarr[ion+1]
# The first row
A[0, 0] = -carr[0]
A[0, 1] = rarr[1]
# The last row
A[nstates-1, nstates-2] = carr[nstates-2]
A[nstates-1, nstates-1] = -rarr[nstates-1]
# Compute eigenvalues and eigenvectors using Scipy
la, v = LA.eig(A)
# Rerange the eigenvalues. Try a simple way in here.
idx = np.argsort(la)
la = la[idx]
v = v[:, idx]
# Compute inverse of eigenvectors
v_inverse = LA.inv(v)
# transpose the order to as same as the Fortran Version
v = v.transpose()
v_inverse = v_inverse.transpose()
# Save eigenvalues and eigenvectors into arrays
for j in range(nstates):
self._eigenvalues[ite, j] = la[j]
self._equilibrium_states[ite, j] = eqi[j]
for i in range(nstates):
self._eigenvectors[ite, i, j] = v[i, j]
self._eigenvector_inverses[ite, i, j] = v_inverse[i, j]
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[:]
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))
