Merge pull request #2 from wkerzendorf/plasma/ab_cleanup

fix for chi_0 problem
tardis-sn · Sep 15, 2015 · 191f711 · 191f711
2 parents 84e7950 + aecec3d
commit 191f711
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 6 deletions.
diff --git a/tardis/plasma/properties/ion_population.py b/tardis/plasma/properties/ion_population.py
@@ -97,31 +97,42 @@ def __init__(self, plasma_parent=None, departure_coefficient=None,
             self.delta_treatment = self.plasma_parent.delta_treatment
         except:
             self.delta_treatment = delta_treatment
+
         self.chi_0_species = chi_0_species
 
+
+    def _set_chi_0(self, ionization_data):
+        if self.chi_0_species == (20, 2):
+            self.chi_0 = 1.9020591570241798e-11
+        else:
+            self.chi_0 = ionization_data.ionization_energy.ix[self.chi_0_species]
+
+
     def calculate(self, w, ionization_data, beta_rad, t_electrons, t_rad,
         beta_electron):
+
+        if getattr(self, 'chi_0', None) is None:
+            self._set_chi_0(ionization_data)
+
         if self.delta_treatment is None:
-            chi_0 = ionization_data.ionization_energy.ix[self.chi_0_species[
-                0]].ix[self.chi_0_species[1]]
             if self.departure_coefficient is None:
                 departure_coefficient = 1. / w
             else:
                 departure_coefficient = self.departure_coefficient
             radiation_field_correction = -np.ones((len(ionization_data), len(
                 beta_rad)))
             less_than_chi_0 = (
-                ionization_data.ionization_energy < chi_0).values
+                ionization_data.ionization_energy < self.chi_0).values
             factor_a = (t_electrons / (departure_coefficient * w * t_rad))
             radiation_field_correction[~less_than_chi_0] = factor_a * \
                 np.exp(np.outer(ionization_data.ionization_energy.values[
                 ~less_than_chi_0], beta_rad - beta_electron))
             radiation_field_correction[less_than_chi_0] = 1 - np.exp(np.outer(
                 ionization_data.ionization_energy.values[less_than_chi_0],
-                beta_rad) - beta_rad * chi_0)
+                beta_rad) - beta_rad * self.chi_0)
             radiation_field_correction[less_than_chi_0] += factor_a * np.exp(
                 np.outer(ionization_data.ionization_energy.values[
-                less_than_chi_0],beta_rad) - chi_0 * beta_electron)
+                less_than_chi_0],beta_rad) - self.chi_0 * beta_electron)
         else:
             radiation_field_correction = np.ones((len(ionization_data),
                 len(beta_rad))) * self.plasma_parent.delta_treatment

diff --git a/tardis/plasma/tests/conftest.py b/tardis/plasma/tests/conftest.py
@@ -211,7 +211,7 @@ def electron_densities(phi_saha_lte, partition_function, number_density):
 @pytest.fixture
 def delta(w, ionization_data, beta_rad, t_electrons, t_rad, beta_electron,
           levels):
-    delta_module = RadiationFieldCorrection(chi_0_species=(2,2))
+    delta_module = RadiationFieldCorrection(chi_0_species=(2,1))
     return delta_module.calculate(w, ionization_data, beta_rad, t_electrons,
         t_rad, beta_electron)