Merge pull request #377 from aoifeboyle/plasma/static_methods

Plasma static methods

tardis/plasma/properties/ion_population.py

@@ -10,9 +10,9 @@
 logger = logging.getLogger(__name__)
 
 __all__ = ['PhiSahaNebular', 'PhiSahaLTE', 'RadiationFieldCorrection',
-           'IonNumberDensity', 'PhiGeneral']
+           'IonNumberDensity']
 
-class PhiGeneral(ProcessingPlasmaProperty):
+class PhiSahaLTE(ProcessingPlasmaProperty):
     """
     Outputs:
     general_phi : Pandas DataFrame
@@ -21,18 +21,17 @@ class PhiGeneral(ProcessingPlasmaProperty):
         on PhiSahaLTE, but the code cannot deal with the inclusion of two
         properties that generate a property called 'phi'.
     """
-    outputs = ('general_phi',)
-    latex_name = ('\\Phi_{\\textrm{LTE}}',)
+    outputs = ('phi',)
+    latex_name = ('\\Phi',)
     latex_formula = ('\\dfrac{2Z_{i,j+1}}{Z_{i,j}}\\Big(\
                      \\dfrac{2\\pi m_{e}/\\beta_{\\textrm{rad}}}{h^2}\
                      \\Big)^{3/2}e^{\\dfrac{-\\chi_{i,j}}{kT_{\
                      \\textrm{rad}}}}',)
 
-    def calculate(self, g_electron, beta_rad, partition_function,
-        ionization_data):
+    @staticmethod
+    def calculate(g_electron, beta_rad, partition_function, ionization_data):
         def calculate_phis(group):
             return group[1:] / group[:-1].values
-
         phis = partition_function.groupby(level='atomic_number').apply(
             calculate_phis)
         phis = pd.DataFrame(phis.values, index=phis.index.droplevel(0))
@@ -54,7 +53,17 @@ class PhiSahaNebular(ProcessingPlasmaProperty):
                      \\dfrac{T_{\\textrm{electron}}}{T_{\\textrm{rad}}}\
                      \\right)^{1/2}',)
 
-    def calculate(self, general_phi, t_rad, w, zeta_data, t_electrons, delta):
+    def calculate(self, t_rad, w, zeta_data, t_electrons, delta,
+            g_electron, beta_rad, partition_function, ionization_data):
+        phi_lte = PhiSahaLTE.calculate(g_electron, beta_rad,
+            partition_function, ionization_data)
+        zeta = self.get_zeta_values(zeta_data, phi_lte, t_rad)
+        phis = phi_lte * w * ((zeta * delta) + w * (1 - zeta)) * \
+               (t_electrons/t_rad) ** .5
+        return phis
+
+    @staticmethod
+    def get_zeta_values(zeta_data, general_phi, t_rad):
         try:
             zeta = interpolate.interp1d(zeta_data.columns.values, zeta_data.ix[
                 general_phi.index].values)(t_rad)
@@ -65,23 +74,7 @@ def calculate(self, general_phi, t_rad, w, zeta_data, t_electrons, delta):
                              '- requested {2}'.format(
                 zeta_data.columns.values.min(), zeta_data.columns.values.max(),
                 t_rad))
-
-        phis = general_phi * w * ((zeta * delta) + w * (1 - zeta)) * \
-               (t_electrons/t_rad) ** .5
-        return phis
-
-class PhiSahaLTE(ProcessingPlasmaProperty):
-    """
-    Outputs:
-    phi_saha_lte: Pandas DataFrame
-        The ionization equilibrium as calculated using the Saha equation.
-    """
-    outputs = ('phi',)
-    latex_name = ('\\Phi',)
-    latex_formula = ('\\Phi_{\\textrm{LTE}}',)
-
-    def calculate(self, general_phi):
-        return general_phi
+        return zeta
 
 class RadiationFieldCorrection(ProcessingPlasmaProperty):
     """

tardis/plasma/properties/partition_function.py

@@ -22,7 +22,8 @@ class LevelBoltzmannFactorLTE(ProcessingPlasmaProperty):
     latex_formula = ('g_{i,j,k}e^{\\dfrac{-\\epsilon_{i,j,k}}{k_{\
         \\textrm{B}}T_{\\textrm{rad}}}}',)
 
-    def calculate(self, excitation_energy, g, beta_rad, levels):
+    @staticmethod
+    def calculate(excitation_energy, g, beta_rad, levels):
         exponential = np.exp(np.outer(excitation_energy.values, -beta_rad))
         level_boltzmann_factor_array = (g.values[np.newaxis].T *
                                         exponential)
@@ -44,13 +45,8 @@ class LevelBoltzmannFactorDiluteLTE(ProcessingPlasmaProperty):
 
     def calculate(self, levels, g, excitation_energy, beta_rad, w,
         metastability):
-        exponential = np.exp(np.outer(excitation_energy.values, -beta_rad))
-        level_boltzmann_factor_array = (g.values[np.newaxis].T *
-                                        exponential)
-        level_boltzmann_factor = pd.DataFrame(level_boltzmann_factor_array,
-                                              index=levels,
-                                              columns=np.arange(len(beta_rad)),
-                                              dtype=np.float64)
+        level_boltzmann_factor = LevelBoltzmannFactorLTE.calculate(
+            excitation_energy, g, beta_rad, levels)
         level_boltzmann_factor[~metastability] *= w
         return level_boltzmann_factor
 

tardis/plasma/properties/property_collections.py

@@ -5,7 +5,7 @@
     LinesUpperLevelIndex, TauSobolev, TRadiative, AtomicData, Abundance,
     Density, TimeExplosion, BetaSobolev, JBlues,
     TransitionProbabilities, StimulatedEmissionFactor, SelectedAtoms,
-    PhiGeneral, PhiSahaNebular, LevelBoltzmannFactorDiluteLTE, DilutionFactor,
+    PhiSahaNebular, LevelBoltzmannFactorDiluteLTE, DilutionFactor,
     ZetaData, ElectronTemperature, LinkTRadTElectron, BetaElectron,
     RadiationFieldCorrection, RadiationFieldCorrectionInput,
     LevelBoltzmannFactorNoNLTE, LevelBoltzmannFactorNLTE, NLTEData,
@@ -23,7 +23,7 @@ class PlasmaPropertyCollection(list):
     Levels, Lines, AtomicMass, PartitionFunction,
     GElectron, IonizationData, NumberDensity, LinesLowerLevelIndex,
     LinesUpperLevelIndex, TauSobolev, LevelNumberDensity, IonNumberDensity,
-    StimulatedEmissionFactor, SelectedAtoms, PhiGeneral, ElectronTemperature])
+    StimulatedEmissionFactor, SelectedAtoms, ElectronTemperature])
 lte_ionization_properties = PlasmaPropertyCollection([PhiSahaLTE])
 lte_excitation_properties = PlasmaPropertyCollection([LevelBoltzmannFactorLTE])
 macro_atom_properties = PlasmaPropertyCollection([BetaSobolev,

tardis/plasma/tests/conftest.py

@@ -7,18 +7,7 @@
 
 import tardis
 from tardis.atomic import AtomData
-from tardis.plasma.properties.ion_population import (PhiGeneral,
-    PhiSahaLTE, IonNumberDensity, PhiSahaNebular, RadiationFieldCorrection)
-from tardis.plasma.properties.general import (BetaRadiation, GElectron,
-    NumberDensity, ElectronTemperature, BetaElectron, SelectedAtoms)
-from tardis.plasma.properties.partition_function import (
-    LevelBoltzmannFactorLTE, PartitionFunction, LevelBoltzmannFactorDiluteLTE)
-from tardis.plasma.properties.atomic import (Levels, Lines, AtomicMass,
-    IonizationData, LinesUpperLevelIndex, LinesLowerLevelIndex, ZetaData,
-    Chi0)
-from tardis.plasma.properties.level_population import LevelNumberDensity
-from tardis.plasma.properties.radiative_properties import (TauSobolev,
-    StimulatedEmissionFactor, BetaSobolev, TransitionProbabilities)
+from tardis.plasma.properties import *
 
 # INPUTS
 
@@ -194,22 +183,17 @@ def partition_function(level_boltzmann_factor_lte):
 # ION POPULATION PROPERTIES
 
 @pytest.fixture
-def general_phi(g_electron, beta_rad, partition_function, ionization_data):
-    phi_general_module = PhiGeneral(None)
-    return phi_general_module.calculate(g_electron, beta_rad,
+def phi_saha_lte(g_electron, beta_rad, partition_function, ionization_data):
+    phi_saha_lte_module = PhiSahaLTE(None)
+    return phi_saha_lte_module.calculate(g_electron, beta_rad,
                                         partition_function, ionization_data)
 
 @pytest.fixture
-def phi_saha_lte(general_phi):
-    phi_module = PhiSahaLTE(None)
-    return phi_module.calculate(general_phi)
-
-@pytest.fixture
-def phi_saha_nebular(w, ionization_data, beta_rad, t_electrons, t_rad,
-        beta_electron, delta, general_phi, zeta_data):
+def phi_saha_nebular(t_rad, w, zeta_data, t_electrons, delta,
+            g_electron, beta_rad, partition_function, ionization_data):
     phi_saha_nebular_module = PhiSahaNebular(None)
-    return phi_saha_nebular_module.calculate(general_phi, t_rad, w, zeta_data,
-                                             t_electrons, delta)
+    return phi_saha_nebular_module.calculate(t_rad, w, zeta_data, t_electrons,
+        delta, g_electron, beta_rad, partition_function, ionization_data)
 
 @pytest.fixture
 def ion_number_density(phi_saha_lte, partition_function, number_density):