From c5e1dce7c50857c60925c92c762ebb0da766270e Mon Sep 17 00:00:00 2001 From: Aoife Boyle Date: Mon, 17 Aug 2015 16:45:46 +0100 Subject: [PATCH] Simplified using static methods. --- tardis/plasma/properties/ion_population.py | 43 ++++++++----------- .../plasma/properties/partition_function.py | 12 ++---- .../plasma/properties/property_collections.py | 4 +- tardis/plasma/tests/conftest.py | 32 ++++---------- 4 files changed, 32 insertions(+), 59 deletions(-) diff --git a/tardis/plasma/properties/ion_population.py b/tardis/plasma/properties/ion_population.py index ac41be97a1a..ffd46a0d3a9 100644 --- a/tardis/plasma/properties/ion_population.py +++ b/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): """ diff --git a/tardis/plasma/properties/partition_function.py b/tardis/plasma/properties/partition_function.py index bef2600cb47..8ed7bb8e54c 100644 --- a/tardis/plasma/properties/partition_function.py +++ b/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 diff --git a/tardis/plasma/properties/property_collections.py b/tardis/plasma/properties/property_collections.py index cb7a8b769a9..01f24db09d8 100644 --- a/tardis/plasma/properties/property_collections.py +++ b/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, diff --git a/tardis/plasma/tests/conftest.py b/tardis/plasma/tests/conftest.py index 4aa8f703a71..a013b4ac3ae 100644 --- a/tardis/plasma/tests/conftest.py +++ b/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):