def test_iteration(self, test_name: str): """Test that `IonizationState` instances iterate impeccably.""" states = list(self.instances[test_name]) charge_numbers = [state.charge_number for state in states] ionic_fractions = np.array([state.ionic_fraction for state in states]) ionic_symbols = [state.ionic_symbol for state in states] 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(charge_numbers == expected_charges): errors.append( f"The resulting charge numbers are {charge_numbers}, " f"which are not equal to the expected charge numbers, " 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 charge_numbers ] expected_symbols = [particle.ionic_symbol for particle in expected_particles] if 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) pytest.fail(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_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_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) ]), ("The ionic fractions for IonizationState are not defaulting " "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_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 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))
def sort_key(k): return atomic_number(k), mass_number(k) if Particle( k).isotope else 0
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[:]