Plasma number densities using isotopes (#1978)

* Added isotope_mass as a plasma property similar to atomic_mass

* Added isotope number density calculation

* Added simple (currently failing) test

* Working tests

* Docstrings for new methods

* black

* Better variable name

* Fix

tardis/plasma/properties/atomic.py

@@ -6,6 +6,7 @@
 from scipy.special import expn
 from scipy.interpolate import PchipInterpolator
 from collections import Counter as counter
+import radioactivedecay as rd
 from tardis import constants as const
 
 from tardis.plasma.properties.base import (
@@ -32,6 +33,7 @@
     "LinesLowerLevelIndex",
     "LinesUpperLevelIndex",
     "AtomicMass",
+    "IsotopeMass",
     "IonizationData",
     "ZetaData",
     "NLTEData",
@@ -537,6 +539,52 @@ def calculate(self, atomic_data, selected_atoms):
             return atomic_data.atom_data.loc[selected_atoms].mass
 
 
+class IsotopeMass(ProcessingPlasmaProperty):
+    """
+    Attributes
+    ----------
+    isotope_mass : pandas.Series
+        Masses of the isotopes used. Indexed by isotope name e.g. 'Ni56'.
+    """
+
+    outputs = ("isotope_mass",)
+
+    def calculate(self, isotope_abundance):
+        """
+        Determine mass of each isotope.
+
+        Parameters
+        ----------
+        isotope_abundance : pandas.DataFrame
+            Fractional abundance of isotopes.
+
+        Returns
+        -------
+        pandas.DataFrame
+            Masses of the isotopes used. Indexed by isotope name e.g. 'Ni56'.
+        """
+        if getattr(self, self.outputs[0]) is not None:
+            return (getattr(self, self.outputs[0]),)
+        else:
+            if isotope_abundance.empty:
+                return None
+            isotope_mass_dict = {}
+            for Z, A in isotope_abundance.index:
+                element_name = rd.utils.Z_to_elem(Z)
+                isotope_name = element_name + str(A)
+
+                isotope_mass_dict[(Z, A)] = rd.Nuclide(isotope_name).atomic_mass
+
+            isotope_mass_df = pd.DataFrame.from_dict(
+                isotope_mass_dict, orient="index", columns=["mass"]
+            )
+            isotope_mass_df.index = pd.MultiIndex.from_tuples(
+                isotope_mass_df.index
+            )
+            isotope_mass_df.index.names = ["atomic_number", "mass_number"]
+            return isotope_mass_df / const.N_A
+
+
 class IonizationData(BaseAtomicDataProperty):
     """
     Attributes

tardis/plasma/properties/general.py

@@ -13,6 +13,7 @@
     "BetaRadiation",
     "GElectron",
     "NumberDensity",
+    "IsotopeNumberDensity",
     "SelectedAtoms",
     "ElectronTemperature",
     "BetaElectron",
@@ -95,6 +96,47 @@ def calculate(atomic_mass, abundance, density):
         return number_densities.div(atomic_mass.loc[abundance.index], axis=0)
 
 
+class IsotopeNumberDensity(ProcessingPlasmaProperty):
+    """
+    Calculate the atom number density based on isotope mass.
+
+    Attributes
+    ----------
+    isotope_number_density : Pandas DataFrame, dtype float
+                     Indexed by atomic number, columns corresponding to zones
+    """
+
+    outputs = ("isotope_number_density",)
+    latex_name = ("N_{i}",)
+
+    @staticmethod
+    def calculate(isotope_mass, isotope_abundance, density):
+        """
+        Calculate the atom number density based on isotope mass.
+
+        Parameters
+        ----------
+        isotope_mass : pandas.DataFrame
+            Masses of isotopes.
+        isotope_abundance : pandas.DataFrame
+            Fractional abundance of isotopes.
+        density : pandas.DataFrame
+            Density of each shell.
+
+        Returns
+        -------
+        pandas.DataFrame
+            Indexed by atomic number, columns corresponding to zones.
+        """
+        number_densities = isotope_abundance * density
+        isotope_number_density_array = (
+            number_densities.to_numpy() / isotope_mass.to_numpy()
+        )
+        return pd.DataFrame(
+            isotope_number_density_array, index=isotope_abundance.index
+        )
+
+
 class SelectedAtoms(ProcessingPlasmaProperty):
     """
     Attributes

tardis/plasma/properties/plasma_input.py

@@ -5,6 +5,7 @@
     "DilutionFactor",
     "AtomicData",
     "Abundance",
+    "IsotopeAbundance",
     "Density",
     "TimeExplosion",
     "JBlueEstimator",
@@ -62,6 +63,17 @@ class Abundance(Input):
     outputs = ("abundance",)
 
 
+class IsotopeAbundance(Input):
+    """
+    Attributes
+    ----------
+    isotope_abundance : Numpy array, dtype float
+        Fractional abundance of isotopes
+    """
+
+    outputs = ("isotope_abundance",)
+
+
 class Density(ArrayInput):
     """
     Attributes

tardis/plasma/properties/property_collections.py

@@ -150,3 +150,6 @@ class PlasmaPropertyCollection(list):
         TwoPhotonFrequencySampler,
     ]
 )
+isotope_properties = PlasmaPropertyCollection(
+    [IsotopeAbundance, IsotopeMass, IsotopeNumberDensity]
+)

tardis/plasma/standard_plasmas.py

@@ -27,6 +27,7 @@
     continuum_interaction_inputs,
     adiabatic_cooling_properties,
     two_photon_properties,
+    isotope_properties,
 )
 from tardis.plasma.exceptions import PlasmaConfigError
 
@@ -248,6 +249,10 @@ def assemble_plasma(config, model, atom_data=None):
                 electron_densities=electron_densities
             )
 
+    if not model.raw_isotope_abundance.empty:
+        plasma_modules += isotope_properties
+        kwargs.update(isotope_abundance=model.raw_isotope_abundance)
+
     kwargs["helium_treatment"] = config.plasma.helium_treatment
 
     plasma = BasePlasma(

tardis/plasma/tests/test_tardis_model_density_config.py

@@ -31,6 +31,15 @@ def test_electron_densities(raw_plasma):
     assert_almost_equal(raw_plasma.electron_densities[3], 2.6e14)
 
 
+def test_isotope_number_densities(raw_plasma):
+    assert_almost_equal(
+        raw_plasma.isotope_number_density.loc[(28, 56), 0], 9688803936.317898
+    )
+    assert_almost_equal(
+        raw_plasma.isotope_number_density.loc[(28, 58), 1], 13097656958.746628
+    )
+
+
 def test_t_rad(raw_plasma):
     assert_almost_equal(raw_plasma.t_rad[5], 76459.592)
     assert_almost_equal(raw_plasma.t_rad[3], 76399.042)