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))