NumberDensity as Input to plasma

tardis/model/base.py

@@ -161,6 +161,45 @@ def t_radiative(self, new_t_radiative):
                 "Trying to set t_radiative for different number of shells."
             )
 
+    @property
+    def elemental_number_density(self):
+        elemental_number_density = (
+            (
+                self.composition.elemental_mass_fraction
+                * self.composition.density
+            )
+            .divide(self.composition.element_masses, axis=0)
+            .dropna()
+        )
+        elemental_number_density = elemental_number_density.iloc[
+            :,
+            self.geometry.v_inner_boundary_index : self.geometry.v_outer_boundary_index,
+        ]
+        elemental_number_density.columns = range(
+            len(elemental_number_density.columns)
+        )
+        return elemental_number_density
+
+    @property
+    def isotopic_number_density(self):
+        isotopic_number_density = (
+            self.composition.isotopic_mass_fraction * self.composition.density
+        ).divide(
+            self.composition.isotope_masses.loc[
+                self.composition.isotopic_mass_fraction.index
+            ]
+            * u.u.to(u.g),
+            axis=0,
+        )
+        isotopic_number_density = isotopic_number_density.iloc[
+            :,
+            self.geometry.v_inner_boundary_index : self.geometry.v_outer_boundary_index,
+        ]
+        isotopic_number_density.columns = range(
+            len(isotopic_number_density.columns)
+        )
+        return isotopic_number_density
+
     @property
     def radius(self):
         return self.time_explosion * self.velocity

tardis/plasma/properties/atomic.py

@@ -31,8 +31,6 @@
     "Lines",
     "LinesLowerLevelIndex",
     "LinesUpperLevelIndex",
-    "AtomicMass",
-    "IsotopeMass",
     "IonizationData",
     "ZetaData",
     "NLTEData",
@@ -76,8 +74,6 @@ def _filter_atomic_property(self, levels, selected_atoms):
         return levels[levels.index.isin(selected_atoms, level="atomic_number")]
 
     def _set_index(self, levels):
-        # levels = levels.set_index(['atomic_number', 'ion_number',
-        #                          'level_number'])
         return (
             levels.index,
             levels["energy"],
@@ -541,69 +537,6 @@ def calculate(self, level_idxs2line_idx, level_idxs2continuum_idx):
         return level_idxs2transition_idx
 
 
-class AtomicMass(ProcessingPlasmaProperty):
-    """
-    Attributes
-    ----------
-    atomic_mass : pandas.Series
-        Atomic masses of the elements used. Indexed by atomic number.
-    """
-
-    outputs = ("atomic_mass",)
-
-    def calculate(self, atomic_data, selected_atoms):
-        if getattr(self, self.outputs[0]) is not None:
-            return (getattr(self, self.outputs[0]),)
-        else:
-            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

@@ -12,8 +12,6 @@
 __all__ = [
     "BetaRadiation",
     "GElectron",
-    "NumberDensity",
-    "IsotopeNumberDensity",
     "SelectedAtoms",
     "ElectronTemperature",
     "BetaElectron",
@@ -79,64 +77,6 @@ def calculate(self, beta_electron):
         return super(ThermalGElectron, self).calculate(beta_electron)
 
 
-class NumberDensity(ProcessingPlasmaProperty):
-    """
-    Attributes
-    ----------
-    number_density : Pandas DataFrame, dtype float
-                     Indexed by atomic number, columns corresponding to zones
-    """
-
-    outputs = ("number_density",)
-    latex_name = ("N_{i}",)
-
-    @staticmethod
-    def calculate(atomic_mass, abundance, density):
-        number_densities = 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

@@ -1,10 +1,14 @@
-from tardis.plasma.properties.base import Input, ArrayInput, DataFrameInput
+from tardis.plasma.properties.base import (
+    ArrayInput,
+    Input,
+)
 
 __all__ = [
     "TRadiative",
     "DilutionFactor",
     "AtomicData",
     "Abundance",
+    "NumberDensity",
     "IsotopeAbundance",
     "Density",
     "TimeExplosion",
@@ -160,3 +164,15 @@ class NLTEIonizationSpecies(Input):
 class NLTEExcitationSpecies(Input):
 
     outputs = ("nlte_excitation_species",)
+
+
+class NumberDensity(Input):
+    """
+    Attributes
+    ----------
+    number_density : Pandas DataFrame, dtype float
+                     Indexed by atomic number, columns corresponding to zones
+    """
+
+    outputs = ("number_density",)
+    latex_name = ("N_{i}",)

tardis/plasma/properties/property_collections.py

@@ -9,7 +9,7 @@ class PlasmaPropertyCollection(list):
     [
         TRadiative,
         Abundance,
-        Density,
+        NumberDensity,
         TimeExplosion,
         AtomicData,
         DilutionFactor,
@@ -25,11 +25,9 @@ class PlasmaPropertyCollection(list):
         BetaRadiation,
         Levels,
         Lines,
-        AtomicMass,
         PartitionFunction,
         GElectron,
         IonizationData,
-        NumberDensity,
         LinesLowerLevelIndex,
         LinesUpperLevelIndex,
         TauSobolev,
@@ -153,6 +151,3 @@ class PlasmaPropertyCollection(list):
         TwoPhotonFrequencySampler,
     ]
 )
-isotope_properties = PlasmaPropertyCollection(
-    [IsotopeAbundance, IsotopeMass, IsotopeNumberDensity]
-)