def abundances(self, abundances_dict: Optional[Dict[ParticleLike, Real]]): """ 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 = {} for old_key in old_keys: new_keys_dict[particle_symbol(old_key)] = old_key except Exception: raise ParticleError( 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 ParticleError( 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 ParticleError(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 {particle}." ) elif not 0 <= int_charge <= atomic_number(particle): raise ChargeError( f"{int_charge} is not a valid charge for {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 ionic_fractions(self, inputs: Union[Dict, List, Tuple]): """ Set the ionic fractions. Notes ----- The ionic fractions are initialized during instantiation of `~plasmapy.particles.IonizationStateCollection`. 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 ------ `~plasmapy.particles.exceptions.ParticleError` 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 IonizationStateCollection instance as a whole. An # example of the problem is `s = IonizationStateCollection(["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 ParticleError( "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 " "IonizationStateCollection 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 ParticleError( "Unable to create IonizationStateCollection 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 ParticleError( 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. def _sort_entries_by_atomic_and_mass_numbers(k): return ( particles[k].atomic_number, particles[k].mass_number if particles[k].isotope else 0, ) sorted_keys = sorted( original_keys, key=_sort_entries_by_atomic_and_mass_numbers ) _elements_and_isotopes = [] _particle_instances = [] new_ionic_fractions = {} if inputs_have_quantities: n_elems = {} for key in sorted_keys: new_key = particles[key].symbol _particle_instances.append(particles[key]) if new_key in _elements_and_isotopes: raise ParticleError( "Repeated particles in IonizationStateCollection." ) nstates_input = len(inputs[key]) nstates = particles[key].atomic_number + 1 if nstates != nstates_input: raise ParticleError( 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 ParticleError("Units are not inverse volume.") from exc elif ( isinstance(inputs[key], np.ndarray) and inputs[key].dtype.kind == "f" ): new_ionic_fractions[particles[key].symbol] = inputs[key] else: try: new_ionic_fractions[particles[key].symbol] = np.array( inputs[key], dtype=np.float ) except ValueError as exc: raise ParticleError( 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 ParticleError( 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 ParticleError( 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.n0): new_n = 0 * u.m ** -3 for key in _elements_and_isotopes: new_n += n_elems[key] self.n0 = new_n new_abundances = {} for key in _elements_and_isotopes: new_abundances[key] = np.float(n_elems[key] / self.n0) 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 ParticleError( "Invalid inputs to IonizationStateCollection." ) from exc _particle_instances.sort(key=_atomic_number_and_mass_number) _elements_and_isotopes = [ particle.symbol for particle in _particle_instances ] new_ionic_fractions = { particle.symbol: 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 ParticleError( "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 IonizationStateCollection instance for " f"key = {repr(key)} and value = {repr(value)}" ) try: particle = particle_symbol(key) self.ionic_fractions[key] except (ParticleError, 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.n0) if abundance_is_undefined: if n_is_defined: self._pars["abundances"][particle] = new_n_elem / self.n0 elif all_abundances_are_nan: self.n0 = new_n_elem self._pars["abundances"][particle] = 1 else: raise ParticleError( 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 normalized to one." ) self._ionic_fractions[particle][:] = new_fractions[:]