Merge 4e317d6 into f91489e

tardis/io/config_reader.py

@@ -880,10 +880,37 @@ def from_config_dict(cls, raw_dict, atom_data=None, test_parser=False):
 
 
         #^^^^^^^ NLTE subsection of Plasma end
-
+
+		# HACHINGER SECTION **************************************************************
+
+        equilibrium_with_continuum_config_dict={}
+        equilibrium_with_continuum_species = []
+        if 'equilibrium_with_continuum'in plasma_section:
+            equilibrium_with_continuum_section = plasma_section['equilibrium_with_continuum']
+            if 'species' in equilibrium_with_continuum_section:
+                equilibrium_with_continuum_species_list = equilibrium_with_continuum_section.pop('species')
+                for species_string in equilibrium_with_continuum_species_list:
+                    equilibrium_with_continuum_species.append(species_string_to_tuple(species_string))
+
+                equilibrium_with_continuum_config_dict['species'] = equilibrium_with_continuum_species
+                equilibrium_with_continuum_config_dict['species_string'] = equilibrium_with_continuum_species_list
+                equilibrium_with_continuum_config_dict.update(equilibrium_with_continuum_section)
+
+        if not equilibrium_with_continuum_config_dict:
+            equilibrium_with_continuum_config_dict['species'] = []
+
+        plasma_config_dict['equilibrium_with_continuum'] = equilibrium_with_continuum_config_dict
+
+        # HE IONISATION FORCING **********************************************************
+
+        helium_forcing = plasma_section.get('helium_forcing')
+
+        helium_forcing_config_dict = {}
+    	helium_forcing_config_dict = helium_forcing       
+        plasma_config_dict['helium_forcing'] = helium_forcing_config_dict
+
         config_dict['plasma'] = plasma_config_dict
-
-
+
         #^^^^^^^^^^^^^^ End of Plasma Section
 
         ##### Monte Carlo Section

tardis/model.py

@@ -128,7 +128,9 @@ def __init__(self, tardis_config):
                                                          nlte_config=tardis_config.plasma.nlte,
                                                          delta_treatment=tardis_config.plasma.delta_treatment,
                                                          ionization_mode=tardis_config.plasma.ionization,
-                                                         excitation_mode=tardis_config.plasma.excitation)
+                                                         excitation_mode=tardis_config.plasma.excitation,
+                                                         equilibrium_with_continuum_config=tardis_config.plasma.equilibrium_with_continuum,
+                                                         helium_forcing_config=tardis_config.plasma.helium_forcing)
 
 
 
@@ -321,7 +323,7 @@ def update_radiationfield(self, log_sampling=5):
 
 
     def simulate(self, update_radiation_field=True, enable_virtual=False, initialize_j_blues=False,
-                 initialize_nlte=False):
+                 initialize_nlte=False, initialize_equilibrium_with_continuum=False, initialize_helium_forcing=False):
         """
         Run a simulation
         """

tardis/plasma_array.py

@@ -11,6 +11,7 @@
 from tardis import macro_atom, io
 from tardis.io.util import parse_abundance_dict_to_dataframe
 
+from copy import deepcopy
 
 logger = logging.getLogger(__name__)
 
@@ -94,7 +95,9 @@ class BasePlasmaArray(object):
     @classmethod
     def from_abundance(cls, abundance_dict, density, atom_data, time_explosion,
                        nlte_config=None, ionization_mode='lte',
-                       excitation_mode='lte'):
+                       excitation_mode='lte', 
+                       equilibrium_with_continuum_config = None, 
+                       helium_forcing_config=None):
         """
         Initializing the abundances from the a dictionary like {'Si':0.5, 'Fe':0.5} and a density.
         All other parameters are the same as the normal initializer
@@ -133,19 +136,23 @@ def from_abundance(cls, abundance_dict, density, atom_data, time_explosion,
 
         return cls(number_densities, atom_data, time_explosion.to('s').value,
                    nlte_config=nlte_config, ionization_mode=ionization_mode,
-                   excitation_mode=excitation_mode)
+                   excitation_mode=excitation_mode, 
+                   equilibrium_with_continuum_config = equilibrium_with_continuum_config,
+                   helium_forcing_config = helium_forcing_config)
 
     @classmethod
     def from_hdf5(cls, hdf5store):
         raise NotImplementedError()
 
 
-    def __init__(self, number_densities, atom_data, time_explosion, delta_treatment=None, nlte_config=None,
-                 ionization_mode='lte', excitation_mode='lte'):
+    def __init__(self, number_densities, atom_data, time_explosion, delta_treatment=None, nlte_config=None, equilibrium_with_continuum_config=None,
+                 helium_forcing_config=None, ionization_mode='lte', excitation_mode='lte'):
         self.number_densities = number_densities
         self.atom_data = atom_data
         self.time_explosion = time_explosion
         self.nlte_config = nlte_config
+        self.equilibrium_with_continuum_config = equilibrium_with_continuum_config
+        self.helium_forcing_config = helium_forcing_config
         self.delta_treatment = delta_treatment
         self.electron_densities = self.number_densities.sum(axis=0)
 
@@ -236,6 +243,9 @@ def update_radiationfield(self, t_rads, ws, j_blues=None, t_electrons=None, n_e_
 
         while True:
             self.calculate_ion_populations(phis)
+            if self.equilibrium_with_continuum_config is not None and self.equilibrium_with_continuum_config.species:
+            	stored = {}
+            	self.equilibrium_with_continuum_approximation(self.atom_data, stored)
             ion_numbers = self.ion_populations.index.get_level_values(1).values
             ion_numbers = ion_numbers.reshape((ion_numbers.shape[0], 1))
             new_electron_densities = (self.ion_populations.values * ion_numbers).sum(axis=0)
@@ -253,6 +263,13 @@ def update_radiationfield(self, t_rads, ws, j_blues=None, t_electrons=None, n_e_
             self.electron_densities = 0.5 * (new_electron_densities + self.electron_densities)
 
         self.calculate_level_populations(initialize_nlte=initialize_nlte, excitation_mode=self.excitation_mode)
+
+        # This sets the level populations of the relevant elements back to what was determined by the equilibrium_with_continuum_approximation function.
+        if self.equilibrium_with_continuum_config is not None and self.equilibrium_with_continuum_config.species:
+        	for value in range (0, len(self.equilibrium_with_continuum_config.species)):
+        		equilibrium_with_continuum_element = self.equilibrium_with_continuum_config.species[value][0]
+        		self.level_populations.ix[equilibrium_with_continuum_element].update(stored[str(equilibrium_with_continuum_element)])
+
         self.tau_sobolevs = self.calculate_tau_sobolev()
 
         if self.nlte_config is not None and self.nlte_config.species:
@@ -506,6 +523,7 @@ def calculate_level_populations(self, initialize_nlte=False, excitation_mode='lt
 
         This function updates the 'number_density' column on the levels table (or adds it if non-existing)
         """
+
         Z = self.partition_functions.ix[self.atom_data.levels.index.droplevel(2)].values
 
         ion_number_density = self.ion_populations.ix[self.atom_data.levels.index.droplevel(2)].values
@@ -523,7 +541,6 @@ def calculate_level_populations(self, initialize_nlte=False, excitation_mode='lt
         else:
             self.level_populations.update(level_populations[~self.atom_data.nlte_data.nlte_levels_mask])
 
-
     def calculate_nlte_level_populations(self):
         """
         Calculating the NLTE level populations for specific ions
@@ -732,15 +749,79 @@ def to_hdf5(self, hdf5_store, path, mode='full'):
         else:
             raise NotImplementedError('Currently only mode="full" is supported.')
 
-
-
-
-
-
-
-
-
-
-
-
-
+    def equilibrium_with_continuum_approximation(self, atom_data, stored):
+        for species in self.equilibrium_with_continuum_config.species:
+            ion_species = (species[0], species[1]+1) # The singly ionised species
+            total = self.ion_populations.ix[species[0]].sum() # Stores total number of atoms for later normalisation
+
+            # Sets the singly ionised ground state population to 1
+            self.level_populations.ix[ion_species].ix[0][np.isnan(self.level_populations.ix[ion_species].ix[0])] = 1.0
+
+            # Neutral Populations (New Way):
+
+            neutral = self.level_population_proportionalities.ix[species] * \
+            self.electron_densities * np.exp(self.atom_data.ionization_data.ionization_energy.ix[ion_species[0]].ix[ion_species[1]] * \
+            self.beta_rads) * (1 / (2 * atom_data.levels.g.ix[ion_species].ix[0] * self.g_electrons * self.ws))
+
+            self.level_populations.ix[species].update(neutral) # This line raising warning, but it still works.
+
+            # Forces helium ionised population upwards if selected            
+            if species==(2,0) and self.helium_forcing_config==True:
+            	factor = pd.Series(index = self.electron_densities.index)
+
+            	for number in factor.index:
+            		# Adjusts fit for number of shells
+            		corrected_number = ((number+1)*(38/len(self.electron_densities.index)))
+            		# Function fitted from SH's results
+            		factor[number] = ((0.00003*((corrected_number)^(3))) - (0.0009*((corrected_number)^(2))) + (0.0106*corrected_number) + 0.0551)
+            		# Forcing only applied if it increases ionised population (to prevent decreasing it at early times)
+            		if ((self.level_populations.ix[ion_species].ix[0][number]*factor[number])<(self.level_populations.ix[species].ix[0][number])):
+            			self.level_populations.ix[species].ix[0][number] = (self.level_populations.ix[ion_species].ix[0][number]*factor[number])
+
+            # Further Ionised Populations (Old Way):
+            for value in range(1, species[0]): #Loop through all higher ions
+
+            	# Excited populations
+            	excited = self.level_population_proportionalities.ix[species[0],value].mul(self.ws * \
+            	self.level_populations.ix[species[0],value].ix[0] * \
+            	(self.atom_data.levels.g.ix[species[0],value].ix[0]**(-1)))
+
+                self.level_populations.ix[species[0], value].update(excited)
+
+                # Sets the singly ionised ground state population back to 1
+                if value==1:
+                	reset = self.level_populations.ix[ion_species].ix[0].mul(1 / self.ws)
+                	self.level_populations.ix[ion_species].ix[0].update(reset)
+
+                # Ionised ground state population
+
+                delta = self.delta_treatment
+                zeta_data = self.atom_data.zeta_data
+
+                zeta = interpolate.interp1d(zeta_data.columns.values, zeta_data.ix[species[0], value+1].values)(self.t_rads)
+
+                ground = 2 * ((self.electron_densities)**(-1)) * self.g_electrons * (self.t_electrons / self.t_rads)**0.5 * \
+                np.exp(-1 * self.atom_data.ionization_data.ionization_energy.ix[species[0], value+1] * self.beta_rads) * \
+                self.ws * (zeta + self.ws * (1 - zeta))
+
+                statistical_weights = ground.mul(self.atom_data.levels.g.ix[species[0], 
+                value+1].ix[0] * ((self.atom_data.levels.g.ix[species[0], value].ix[0])**(-1)))
+
+                self.level_populations.ix[species[0], value+1].ix[0].update(statistical_weights)
+
+            # Normalisation:
+
+            test_total = self.level_populations.ix[species[0]].sum()
+            normalisation = self.level_populations.ix[species[0]] * (total / test_total)
+            self.level_populations.ix[species[0]].update(normalisation)
+
+            # Update ion populations:
+
+            for x in self.ion_populations.ix[species[0]].index:
+            	self.ion_populations.ix[species[0], x] = self.level_populations.ix[species[0], x].sum()
+
+        	element = '%s' %(species[0])
+        	# Stores the level populations so that they can be applied again after level populations calculated in main code
+        	stored[element] = deepcopy(self.level_populations.ix[species[0]])
+
+        return atom_data, stored