def test_particle_input_simple(func, args, kwargs, symbol): """ Test that simple functions decorated by particle_input correctly return the correct Particle object. """ try: expected = Particle(symbol) except Exception as e: raise ParticleError( f"Cannot create Particle class from symbol {symbol}") from e try: result = func(*args, **kwargs) except Exception as e: raise ParticleError( f"An exception was raised while trying to execute " f"{func} with args = {args} and kwargs = {kwargs}.") from e assert result == expected, ( f"The result {repr(result)} does not equal the expected value of " f"{repr(expected)}.\n\n" f"func = {func}\n" f"args = {args}\n" f"kwargs = {kwargs}\nsymbol = {symbol}\n" f"{result._attributes}\n" f"{expected._attributes}\n")
def __init__( self, particle: Particle, ionic_fractions=None, *, T_e: u.K = np.nan * u.K, T_i: u.K = None, kappa: Real = np.inf, n_elem: u.m**-3 = np.nan * u.m**-3, tol: Union[float, int] = 1e-15, ): """ Initialize an `~plasmapy.particles.ionization_state.IonizationState` instance. """ self._number_of_particles = particle.atomic_number + 1 if particle.is_ion or particle.is_category(require=("uncharged", "element")): if ionic_fractions is None: ionic_fractions = np.zeros(self._number_of_particles) ionic_fractions[particle.charge_number] = 1.0 particle = Particle( particle.isotope if particle.isotope else particle.element) else: raise ParticleError( "The ionic fractions must not be specified when " "the input particle to IonizationState is an ion.") self._particle = particle try: self.tol = tol self.T_e = T_e self.T_i = T_i self.kappa = kappa if (not np.isnan(n_elem) and isinstance(ionic_fractions, u.Quantity) and ionic_fractions.si.unit == u.m**-3): raise ParticleError( "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 ParticleError( f"Unable to create IonizationState object for {particle.symbol}." ) from exc
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 ParticleError("Invalid temperature.") from None else: if value < 0 * u.K: raise ParticleError("T_e cannot be negative.") self._T_e = value
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 number_densities(self, value: u.m**-3): """Set the number densities for each state.""" if np.any(value.value < 0): raise ParticleError("Number densities cannot be negative.") if len(value) != self.atomic_number + 1: raise ParticleError( f"Incorrect number of charge states for {self.base_particle}") value = value.to(u.m**-3) self._n_elem = value.sum() self._ionic_fractions = value / self._n_elem
def n0(self, n: u.m ** -3): """Set the number density scaling factor.""" try: n = n.to(u.m ** -3) except u.UnitConversionError as exc: raise ParticleError("Units cannot be converted to u.m ** -3.") from exc except Exception as exc: raise ParticleError(f"{n} is not a valid number density.") from exc if n < 0 * u.m ** -3: raise ParticleError("Number density cannot be negative.") self._pars["n"] = n.to(u.m ** -3)
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 ParticleError( f"{electron_temperature} is not a valid temperature." ) from None if temperature < 0 * u.K: raise ParticleError("The electron temperature cannot be negative.") self._pars["T_e"] = temperature
def __init__( self, inputs: Union[Dict[str, np.ndarray], List, Tuple], *, T_e: u.K = np.nan * u.K, abundances: Optional[Dict[str, Real]] = None, log_abundances: Optional[Dict[str, Real]] = None, n0: 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 ParticleError( "Units must be inverse volume for number densities." ) if abundances_provided: raise ParticleError( "Abundances cannot be provided if inputs " "provides number density information." ) set_abundances = False try: self._pars = dict() self.T_e = T_e self.n0 = n0 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 ParticleError( "Unable to create IonizationStateCollection object." ) 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 ParticleError(errmsg) from exc if particle.element and not particle.isotope: raise ParticleError(errmsg) [particles.append(particle) for i in range(multiplier)] except Exception: raise ParticleError( f"{original_item} is not a valid reactant or " "product in a nuclear reaction.") from None return particles
def test_particle_input_classes(): instance = Test_particle_input() symbol = "muon" expected = Particle(symbol) try: result_noparens = instance.method_noparens(symbol) except Exception as e: raise ParticleError("Problem with method_noparens") from e try: result_parens = instance.method_parens(symbol) except Exception as e: raise ParticleError("Problem with method_parens") from e assert result_parens == result_noparens == expected
def __init__(self, ion: Particle, ionic_fraction=None, number_density=None): try: self.ion = ion self.ionic_fraction = ionic_fraction self.number_density = number_density except Exception as exc: raise ParticleError("Unable to create IonicLevel object") from exc
def function_to_test_annotations(particles: Union[Tuple, List], resulting_particles): """ Test that a function with an argument annotated with (Particle, Particle, ...) or [Particle] returns a tuple of expected Particle instances. Arguments ========= particles: tuple or list A collection containing many items, each of which may be a valid representation of a particle or a `~plasmapy.particles.Particle` instance """ expected = [ particle if isinstance(particle, Particle) else Particle(particle) for particle in particles ] # Check that the returned values are Particle instances because # running: # Particle('p+') == 'p+' # will return True because of how Particle.__eq__ is set up. returned_particle_instances = all( isinstance(p, Particle) for p in resulting_particles) returned_correct_instances = all(expected[i] == resulting_particles[i] for i in range(len(particles))) if not returned_particle_instances: raise ParticleError( f"A function decorated by particle_input did not return " f"a collection of Particle instances for input of " f"{repr(particles)}, and instead returned" f"{repr(resulting_particles)}.") if not returned_correct_instances: raise ParticleError( f"A function decorated by particle_input did not return " f"{repr(expected)} as expected, and instead returned " f"{repr(resulting_particles)}.")
def __eq__(self, other): if not isinstance(other, IonizationStateCollection): raise TypeError( "IonizationStateCollection instance can only be compared with " "other IonizationStateCollection instances.") if self.base_particles != other.base_particles: raise ParticleError( "Two IonizationStateCollection 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 = getattr(self, attribute) that = getattr(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 = getattr(self, attribute) that_dict = getattr(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 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 ParticleError("Cannot have negative ionic fractions.") if len(fractions) != self.atomic_number + 1: raise ParticleError("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 ParticleError( "Ionic fractions must be between 0 and 1.") if not np.isclose( sum_of_fractions, 1, rtol=0, atol=self.tol): raise ParticleError("Ionic fractions must sum to one.") self._ionic_fractions = fractions except Exception as exc: raise ParticleError( f"Unable to set ionic fractions of {self.element} to {fractions}." ) from exc
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 ParticleError("Invalid log_abundances.") from None
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 = { atom: 10**log_abundance for atom, log_abundance in value.items() } self.abundances = new_abundances_input except ParticleError: raise ParticleError("Invalid log_abundances.") from None
def __init__(self, ion: Particle, ionic_fraction=None, number_density=None, T_i=None): try: self.ion = ion self.ionic_fraction = ionic_fraction self.number_density = number_density self.T_i = T_i except (ValueError, TypeError) as exc: raise ParticleError("Unable to create IonicLevel object") from exc
def test_list_annotation(particles: Union[Tuple, List]): try: resulting_particles = function_with_list_annotation(particles, "ignore", x="ignore") except Exception as exc2: raise ParticleError( f"Unable to evaluate a function decorated by particle_input" f" with an annotation of [Particle] for inputs of" f" {repr(particles)}.") from exc2 function_to_test_annotations(particles, resulting_particles)
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_Particle_class(arg, kwargs, expected_dict): """ Test that `~plasmapy.particles.Particle` objects for different subatomic particles, elements, isotopes, and ions return the expected properties. Provide a detailed error message that lists all of the inconsistencies with the expected results. """ call = call_string(Particle, arg, kwargs) errmsg = "" try: particle = Particle(arg, **kwargs) except Exception as exc: raise ParticleError(f"Problem creating {call}") from exc for key in expected_dict.keys(): expected = expected_dict[key] if inspect.isclass(expected) and issubclass(expected, Exception): # Exceptions are expected to be raised when accessing certain # attributes for some particles. For example, accessing a # neutrino's mass should raise a MissingParticleDataError since # only upper limits of neutrino masses are presently available. # If expected_dict[key] is an exception, then check to make # sure that this exception is raised. try: with pytest.raises(expected): exec(f"particle.{key}") except pytest.fail.Exception: errmsg += f"\n{call}[{key}] does not raise {expected}." except Exception: errmsg += (f"\n{call}[{key}] does not raise {expected} but " f"raises a different exception.") else: try: result = eval(f"particle.{key}") assert result == expected or u.isclose(result, expected) except AssertionError: errmsg += (f"\n{call}.{key} returns {result} instead " f"of the expected value of {expected}.") except Exception: errmsg += f"\n{call}.{key} raises an unexpected exception." if len(errmsg) > 0: raise Exception(f"Problems with {call}:{errmsg}")
def __getitem__(self, value) -> IonicLevel: """Return information for a single ionization level.""" if isinstance(value, slice): return [ IonicLevel( ion=Particle(self.base_particle, Z=val), ionic_fraction=self.ionic_fractions[val], number_density=self.number_densities[val], T_i=self.T_i[val], ) for val in range(0, self._number_of_particles)[value] ] if isinstance(value, Integral) and 0 <= value <= self.atomic_number: result = IonicLevel( ion=Particle(self.base_particle, Z=value), ionic_fraction=self.ionic_fractions[value], number_density=self.number_densities[value], T_i=self.T_i[value], ) else: if not isinstance(value, Particle): try: value = Particle(value) except InvalidParticleError as exc: raise InvalidParticleError( f"{value} is not a valid charge number or 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.charge_number result = IonicLevel( ion=Particle(self.base_particle, Z=Z), ionic_fraction=self.ionic_fractions[Z], number_density=self.number_densities[Z], T_i=self.T_i[Z], ) else: if not same_element or not same_isotope: raise ParticleError("Inconsistent element or isotope.") elif not has_charge_info: raise ChargeError("No charge number provided.") return result
def T_i(self, value: u.K): """Set the ion temperature.""" if value is None: self._T_i = np.repeat(self._T_e, self._number_of_particles) return if value.size == 1: self._T_i = np.repeat(value, self._number_of_particles) elif value.size == self._number_of_particles: self._T_i = value else: error_str = ( "T_i must be set with either one common temperature" f" for all ions, or a set of {self._number_of_particles} of them. " ) if value.size == 5 and self._number_of_particles != 5: error_str += " For {self.base_particle}, five is right out." raise ParticleError(error_str)
def Z_most_abundant(self) -> List[Integral]: """ A `list` of the charge numbers with the highest ionic fractions. Examples -------- >>> He = IonizationState('He', [0.2, 0.5, 0.3]) >>> He.Z_most_abundant [1] >>> Li = IonizationState('Li', [0.4, 0.4, 0.2, 0.0]) >>> Li.Z_most_abundant [0, 1] """ if np.any(np.isnan(self.ionic_fractions)): raise ParticleError( f"Cannot find most abundant ion of {self.base_particle} " f"because the ionic fractions have not been defined.") return np.flatnonzero( self.ionic_fractions == self.ionic_fractions.max()).tolist()
def __getitem__(self, value) -> IonicFraction: """Return information for a single ionization level.""" if isinstance(value, slice): raise TypeError("IonizationState instances cannot be sliced.") if isinstance(value, Integral) and 0 <= value <= self.atomic_number: result = IonicFraction( ion=Particle(self.base_particle, Z=value), ionic_fraction=self.ionic_fractions[value], number_density=self.number_densities[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 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 = IonicFraction( ion=Particle(self.base_particle, Z=Z), ionic_fraction=self.ionic_fractions[Z], number_density=self.number_densities[Z], ) else: if not same_element or not same_isotope: raise ParticleError("Inconsistent element or isotope.") elif not has_charge_info: raise ChargeError("No integer charge provided.") return result
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[:]
def get_particle(argname, params, already_particle, funcname): argval, Z, mass_numb = params """ Convert the argument to a `~plasmapy.particles.particle_class.Particle` object if it is not already one. """ if not already_particle: if not isinstance(argval, (numbers.Integral, str, tuple, list)): raise TypeError( f"The argument {argname} to {funcname} must be " f"a string, an integer or a tuple or list of them " f"corresponding to an atomic number, or a " f"Particle object.") try: particle = Particle(argval, Z=Z, mass_numb=mass_numb) except InvalidParticleError as e: raise InvalidParticleError( _particle_errmsg(argname, argval, Z, mass_numb, funcname)) from e # We will need to do the same error checks whether or not the # argument is already an instance of the Particle class. if already_particle: particle = argval # If the name of the argument annotated with Particle in the # decorated function is element, isotope, or ion; then this # decorator should raise the appropriate exception when the # particle ends up not being an element, isotope, or ion. cat_table = [ ("element", particle.element, InvalidElementError), ("isotope", particle.isotope, InvalidIsotopeError), ("ion", particle.ionic_symbol, InvalidIonError), ] for category_name, category_symbol, CategoryError in cat_table: if argname == category_name and not category_symbol: raise CategoryError( f"The argument {argname} = {repr(argval)} to " f"{funcname} does not correspond to a valid " f"{argname}.") # Some functions require that particles be charged, or # at least that particles have charge information. _charge_number = particle._attributes["charge number"] must_be_charged = "charged" in require must_have_charge_info = set(any_of) == {"charged", "uncharged"} uncharged = _charge_number == 0 lacks_charge_info = _charge_number is None if must_be_charged and (uncharged or must_have_charge_info): raise ChargeError( f"A charged particle is required for {funcname}.") if must_have_charge_info and lacks_charge_info: raise ChargeError( f"Charge information is required for {funcname}.") # Some functions require particles that belong to more complex # classification schemes. Again, be sure to provide a # maximally useful error message. if not particle.is_category( require=require, exclude=exclude, any_of=any_of): raise ParticleError( _category_errmsg(particle, require, exclude, any_of, funcname)) return particle
def wrapper(*args, **kwargs): annotations = wrapped_function.__annotations__ bound_args = wrapped_signature.bind(*args, **kwargs) default_arguments = bound_args.signature.parameters arguments = bound_args.arguments argnames = bound_args.signature.parameters.keys() # Handle optional-only arguments in function declaration for default_arg in default_arguments: # The argument is not contained in `arguments` if the # user does not explicitly pass an optional argument. # In such cases, manually add it to `arguments` with # the default value of parameter. if default_arg not in arguments: arguments[default_arg] = default_arguments[ default_arg].default funcname = wrapped_function.__name__ args_to_become_particles = [] for argname in annotations.keys(): if isinstance(annotations[argname], tuple): if argname == "return": continue annotated_argnames = annotations[argname] expected_params = len(annotated_argnames) received_params = len(arguments[argname]) if expected_params != received_params: raise ValueError( f"Number of parameters allowed in the tuple " f"({expected_params} parameters) are " f"not equal to number of parameters passed in " f"the tuple ({received_params} parameters).") elif isinstance(annotations[argname], list): annotated_argnames = annotations[argname] expected_params = len(annotated_argnames) if expected_params > 1: raise TypeError( "Put in [Particle] as the annotation to " "accept arbitrary number of Particle arguments.") else: annotated_argnames = (annotations[argname], ) for annotated_argname in annotated_argnames: is_particle = (annotated_argname is Particle or annotated_argname is Optional[Particle]) if is_particle and argname != "return": args_to_become_particles.append(argname) if not args_to_become_particles: raise ParticleError( f"None of the arguments or keywords to {funcname} " f"have been annotated with Particle, as required " f"by the @particle_input decorator.") elif len(args_to_become_particles) > 1: if "Z" in argnames or "mass_numb" in argnames: raise ParticleError( f"The arguments Z and mass_numb in {funcname} are not " f"allowed when more than one argument or keyword is " f"annotated with Particle in functions decorated " f"with @particle_input.") for x in args_to_become_particles: if (annotations[x] is Particle and isinstance(arguments[x], (tuple, list)) and len(arguments[x]) > 1): raise TypeError( f"You cannot pass a tuple or list containing " f"Particles when only single Particle was " f"expected, instead found {arguments[x]}. If you " f"intend to pass more than 1 Particle instance, " f"use a tuple or a list type. " f"That is use (Particle, Particle, ...) or " f"[Particle] in function declaration.") # If the number of arguments and keywords annotated with # Particle is exactly one, then the Z and mass_numb keywords # can be used without potential for ambiguity. Z = arguments.get("Z", None) mass_numb = arguments.get("mass_numb", None) # Go through the argument names and check whether or not they are # annotated with Particle. If they aren't, include the name and # value of the argument as an item in the new keyword arguments # dictionary unchanged. If they are annotated with Particle, then # either convert the representation of a Particle to a Particle if # it is not already a Particle and then do error checks. new_kwargs = {} for argname in argnames: raw_argval = arguments[argname] if isinstance(raw_argval, (tuple, list)): # Input argument value is a tuple or list # of corresponding particles or atomic values. argval_tuple = raw_argval particles = [] else: # Otherwise convert it to tuple anyway so it can work # with loops too. argval_tuple = (raw_argval, ) for pos, argval in enumerate(argval_tuple): should_be_particle = argname in args_to_become_particles # If the argument is not annotated with Particle, then we just # pass it through to the new keywords without doing anything. if not should_be_particle: new_kwargs[argname] = raw_argval continue # Occasionally there will be functions where it will be # useful to allow None as an argument. # In case annotations[argname] is a collection (which looks # like (Particle, Optional[Particle], ...) or [Particle]) if isinstance(annotations[argname], tuple): optional_particle = (annotations[argname][pos] is Optional[Particle]) elif isinstance(annotations[argname], list): optional_particle = annotations[argname] == [ Optional[Particle] ] else: # Otherwise annotations[argname] must be a Particle itself optional_particle = annotations[argname] is Optional[ Particle] if (optional_particle or none_shall_pass) and argval is None: particle = None else: params = (argval, Z, mass_numb) already_particle = isinstance(argval, Particle) particle = get_particle(argname, params, already_particle, funcname) if isinstance(raw_argval, (tuple, list)): # If passed argument is a tuple or list, keep # appending them. particles.append(particle) # Set appended values if current iteration is the # last iteration. if (pos + 1) == len(argval_tuple): new_kwargs[argname] = tuple(particles) del particles else: # Otherwise directly set values new_kwargs[argname] = particle return wrapped_function(**new_kwargs)
def T_e(self) -> u.K: """Return the electron temperature.""" if self._T_e is None: raise ParticleError("No electron temperature has been specified.") return self._T_e.to(u.K, equivalencies=u.temperature_energy())
def average_ion( self, *, include_neutrals: bool = True, use_rms_charge: bool = False, use_rms_mass: bool = False, ) -> CustomParticle: """ Return a |CustomParticle| representing the mean particle included across all ionization states. By default, this method will use the weighted mean to calculate the properties of the |CustomParticle|, where the weights for each ionic level is given by its ionic fraction multiplied by the abundance of the base element or isotope. If ``use_rms_charge`` or ``use_rms_mass`` is `True`, then this method will return the root mean square of the charge or mass, respectively. Parameters ---------- include_neutrals : `bool`, optional, keyword-only If `True`, include neutrals when calculating the mean values of the different particles. If `False`, exclude neutrals. Defaults to `True`. use_rms_charge : `bool`, optional, keyword-only If `True`, use the root mean square charge instead of the mean charge. Defaults to `False`. use_rms_mass : `bool`, optional, keyword-only If `True`, use the root mean square mass instead of the mean mass. Defaults to `False`. Raises ------ `~plasmapy.particles.exceptions.ParticleError` If the abundance of any of the elements or isotopes is not defined and the |IonizationStateCollection| instance includes more than one element or isotope. Returns ------- ~plasmapy.particles.particle_class.CustomParticle Examples -------- >>> states = IonizationStateCollection( ... {"H": [0.1, 0.9], "He": [0, 0.1, 0.9]}, ... abundances={"H": 1, "He": 0.1} ... ) >>> states.average_ion() CustomParticle(mass=2.12498...e-27 kg, charge=1.5876...e-19 C) >>> states.average_ion(include_neutrals=False, use_rms_charge=True, use_rms_mass=True) CustomParticle(mass=2.633...e-27 kg, charge=1.805...e-19 C) """ min_charge = 0 if include_neutrals else 1 all_particles = ParticleList() all_abundances = [] for base_particle in self.base_particles: ionization_state = self[base_particle] ionic_levels = ionization_state.to_list()[min_charge:] all_particles.extend(ionic_levels) base_particle_abundance = self.abundances[base_particle] if np.isnan(base_particle_abundance): if len(self) == 1: base_particle_abundance = 1 else: raise ParticleError( "Unable to provide an average particle without abundances." ) ionic_fractions = ionization_state.ionic_fractions[min_charge:] ionic_abundances = base_particle_abundance * ionic_fractions all_abundances.extend(ionic_abundances) return all_particles.average_particle( use_rms_charge=use_rms_charge, use_rms_mass=use_rms_mass, abundances=all_abundances, )