General bugfixing

docs/plasma_doc/plasma_plots/lte_ionization_balance.py

@@ -2,7 +2,7 @@
 from tardis import atomic, plasma, model_radial_oned
 
 atom_data = atomic.AtomData.from_hdf5()
-number_density = model_radial_oned.calculate_atom_number_densities(atom_data, {'Si':1}, density=1e-13)
+number_density = model_radial_oned.calculate_atom_number_densities(atom_data, {'Si': 1}, density=1e-13)
 atom_data.prepare_atom_data(set(number_density.index.values.astype(int)), max_ion_number=3)
 lte_plasma = plasma.LTEPlasma(number_density, atom_data, max_ion_number=3)
 siI = []
@@ -13,7 +13,7 @@
 for t_rad in t_rads:
     lte_plasma.update_radiationfield(t_rad)
     si_number_density = number_density.get_value(14, 'number_density')
-    ion_density = lte_plasma.ion_number_density / si_number_density
+    ion_density = lte_plasma.ion_populations / si_number_density
     siI.append(ion_density.get_value((14, 0)))
     siII.append(ion_density.get_value((14, 1)))
     siIII.append(ion_density.get_value((14, 2)))

docs/plasma_doc/plasma_plots/nebular_ionization_balance_w.5.py

@@ -6,7 +6,7 @@
 atom_fname = os.path.expanduser('~/.tardis/kurucz_atom.h5')
 
 atom_data = atomic.AtomData.from_hdf5(atom_fname, use_zeta_data=True)
-number_density = model_radial_oned.calculate_atom_number_densities(atom_data, {'Si':1}, density=1e-13)
+number_density = model_radial_oned.calculate_atom_number_densities(atom_data, {'Si': 1}, density=1e-13)
 atom_data.prepare_atom_data(set(number_density.index.values.astype(int)), max_ion_number=3)
 nebular_plasma = plasma.NebularPlasma(number_density, atom_data, max_ion_number=3)
 siI = []
@@ -17,7 +17,7 @@
 for t_rad in t_rads:
     nebular_plasma.update_radiationfield(t_rad, w=.5)
     si_number_density = number_density.get_value(14, 'number_density')
-    ion_density = nebular_plasma.ion_number_density / si_number_density
+    ion_density = nebular_plasma.ion_populations / si_number_density
     siI.append(ion_density.get_value((14, 0)))
     siII.append(ion_density.get_value((14, 1)))
     siIII.append(ion_density.get_value((14, 2)))

docs/plasma_doc/plasma_plots/nebular_ionization_balance_w1.py

@@ -6,7 +6,7 @@
 atom_fname = os.path.expanduser('~/.tardis/kurucz_atom.h5')
 
 atom_data = atomic.AtomData.from_hdf5(atom_fname, use_zeta_data=True)
-number_density = model_radial_oned.calculate_atom_number_densities(atom_data, {'Si':1}, density=1e-13)
+number_density = model_radial_oned.calculate_atom_number_densities(atom_data, {'Si': 1}, density=1e-13)
 atom_data.prepare_atom_data(set(number_density.index.values.astype(int)), max_ion_number=3)
 nebular_plasma = plasma.NebularPlasma(number_density, atom_data, max_ion_number=3)
 siI = []
@@ -17,7 +17,7 @@
 for t_rad in t_rads:
     nebular_plasma.update_radiationfield(t_rad, w=1.)
     si_number_density = number_density.get_value(14, 'number_density')
-    ion_density = nebular_plasma.ion_number_density / si_number_density
+    ion_density = nebular_plasma.ion_populations / si_number_density
     siI.append(ion_density.get_value((14, 0)))
     siII.append(ion_density.get_value((14, 1)))
     siIII.append(ion_density.get_value((14, 2)))

tardis/atomic.py

@@ -7,6 +7,8 @@
 import os
 import h5py
 
+import pdb
+
 from astropy import table, units, constants
 
 from collections import OrderedDict
@@ -525,52 +527,85 @@ def prepare_atom_data(self, selected_atomic_numbers, line_interaction_type='scat
                 self.macro_atom_data['destination_level_idx'] = (np.ones(len(self.macro_atom_data)) * -1).astype(
                     np.int64)
 
-        #Setting NLTE species
-        self.set_nlte_mask(nlte_species)
-
-
-    def set_nlte_mask(self, nlte_species):
-
-        logger.debug('Setting NLTE Species Mask for %s' % nlte_species)
-        self.nlte_mask = np.zeros(self.levels.shape[0]).astype(bool)
-
-        for species in nlte_species:
-            current_mask = (self.levels.index.get_level_values(0) == species[0]) & \
-                           (self.levels.index.get_level_values(1) == species[1])
-
-            self.nlte_mask |= current_mask
-
-    def prepare_nlte_indices(self, nlte_species):
-        pass
-
-    def get_collision_coefficients(self, atomic_number, ion_number, level_number_lower, level_number_upper, t_electron):
-        if self.has_collision_data:
-            try:
-                C_lus = self.collision_data.ix[
-                            (atomic_number, ion_number, level_number_lower, level_number_upper)].values[1:]
-                C_lu = np.interp(t_electron, self.collision_data_temperatures, C_lus)
-                C_ul = C_lu * self.collision_data.ix[
-                    (atomic_number, ion_number, level_number_lower, level_number_upper)].values[0]
-
-            except pd.core.indexing.IndexingError:
-                C_lu = 0
-                C_ul = 0
-                logger.debug('Could not find collision data for atom=%d ion=%d lvl_lower=%d lvl_upper=%d',
-                             atomic_number, ion_number, level_number_lower, level_number_upper)
-            else:
-                logger.debug('Found collision data for atom=%d ion=%d lvl_lower=%d lvl_upper=%d',
-                             atomic_number, ion_number, level_number_lower, level_number_upper)
-
-            return C_lu, C_ul
+        self.nlte_data = NLTEData(self, nlte_species)
 
-        else:
-            return 0., 0.
 
     def __repr__(self):
         return "<Atomic Data UUID=%s MD5=%s Lines=%d Levels=%d>" % \
                (self.uuid1, self.md5, self.lines_data.atomic_number.count(), self.levels_data.energy.count())
 
 
+class NLTEData(object):
+    def __init__(self, atom_data, nlte_species):
+        logger.info('Preparing the NLTE data')
+        self.atom_data = atom_data
+        self.lines = atom_data.lines.reset_index()
+        self.nlte_species = nlte_species
+        self._init_indices()
+        self._create_nlte_mask()
+        if atom_data.has_collision_data:
+            self._create_collision_coefficient_matrix()
+
+
+    def _init_indices(self):
+        self.lines_idx = {}
+        self.lines_level_number_lower = {}
+        self.lines_level_number_upper = {}
+        self.A_uls = {}
+        self.B_uls = {}
+        self.B_lus = {}
+
+        for species in self.nlte_species:
+            lines_idx = np.where((self.lines.atomic_number == species[0]) &
+                                 (self.lines.ion_number == species[1]))
+            self.lines_idx[species] = lines_idx
+            self.lines_level_number_lower[species] = self.lines.level_number_lower.values[lines_idx].astype(int)
+            self.lines_level_number_upper[species] = self.lines.level_number_upper.values[lines_idx].astype(int)
+
+            self.A_uls[species] = self.atom_data.lines.A_ul.values[lines_idx]
+            self.B_uls[species] = self.atom_data.lines.B_ul.values[lines_idx]
+            self.B_lus[species] = self.atom_data.lines.B_lu.values[lines_idx]
+
+    def _create_nlte_mask(self):
+        self.nlte_mask = np.zeros(self.atom_data.levels.energy.count()).astype(bool)
+
+        for species in self.nlte_species:
+            current_mask = (self.atom_data.levels.index.get_level_values(0) == species[0]) & \
+                           (self.atom_data.levels.index.get_level_values(1) == species[1])
+
+            self.nlte_mask |= current_mask
 
 
+    def _create_collision_coefficient_matrix(self):
+        self.C_ul_interpolator = {}
+        self.delta_E_matrices = {}
+        self.g_ratio_matrices = {}
+        collision_group = self.atom_data.collision_data.groupby(level=['atomic_number', 'ion_number'])
+        for species in self.nlte_species:
+            no_of_levels = self.atom_data.levels.ix[species].energy.count()
+            C_ul_matrix = np.zeros((no_of_levels, no_of_levels, len(self.atom_data.collision_data_temperatures)))
+            delta_E_matrix = np.zeros((no_of_levels, no_of_levels))
+            g_ratio_matrix = np.zeros((no_of_levels, no_of_levels))
+
+            for (atomic_number, ion_number, level_number_lower, level_number_upper), line in \
+                collision_group.get_group(species).iterrows():
+                C_ul_matrix[level_number_lower, level_number_upper, :] = line.values[2:]
+                delta_E_matrix[level_number_lower, level_number_upper] = line['delta_e']
+                #TODO TARDISATOMIC fix change the g_ratio to be the otherway round - I flip them now here.
+                g_ratio_matrix[level_number_lower, level_number_upper] = line['g_ratio']
+            self.C_ul_interpolator[species] = interpolate.interp1d(self.atom_data.collision_data_temperatures,
+                                                                   C_ul_matrix)
+            self.delta_E_matrices[species] = delta_E_matrix
+
+            self.g_ratio_matrices[species] = g_ratio_matrix
+
+
+    def get_collision_matrix(self, species, t_electron):
+        c_ul_matrix = self.C_ul_interpolator[species](t_electron)
+
+        c_ul_matrix[np.isnan(c_ul_matrix)] = 0.0
+        #TODO in tardisatomic the g_ratio is the other way round - here I'll flip it in prepare_collision matrix
+        c_lu_matrix = c_ul_matrix * np.exp(-self.delta_E_matrices[species] / t_electron) * self.g_ratio_matrices[
+            species]
+        return c_ul_matrix + c_lu_matrix.transpose()
 

tardis/config_reader.py

@@ -519,7 +519,13 @@ def from_yaml(cls, fname, args=None):
             logger.warn('"w_epsilon" not specified in plasma section - setting it to 1e-10')
             config_dict['w_epsilon'] = 1e-10
 
-            ##### spectrum section ######
+        ##### NLTE Section #####
+
+        config_dict['nlte_options'] = yaml_dict.pop('nlte', {})
+
+
+
+        ##### spectrum section ######
         spectrum_section = yaml_dict.pop('spectrum')
         spectrum_start = parse2quantity(spectrum_section['start']).to('angstrom', units.spectral())
         spectrum_end = parse2quantity(spectrum_section['end']).to('angstrom', units.spectral())

tardis/data/synpp_default.yaml

@@ -0,0 +1,37 @@
+#-
+#- Need to document...
+#-
+---
+output :
+    min_wl      : 100.0        # min. wavelength in AA
+    max_wl      : 10000.0       # max. wavelength in AA
+    wl_step     : 5.0           # wavelength spacing in AA
+grid :
+    bin_width   : 0.3           # opacity bin size in kkm/s
+    v_size      : 20           # size of line-forming region grid
+    v_outer_max : 30.0          # fastest ejecta velocity in kkm/s
+opacity :
+    line_dir    : /Users/wkerzend/software/synapps-linedb/es-data/lines     # path to atomic line data
+    ref_file    : /Users/wkerzend/software/synapps-linedb/es-data/refs.dat  # path to ref. line data
+    form        : exp           # parameterization (only exp for now)
+    v_ref       : 10.0          # reference velocity for parameterization
+    log_tau_min : -5.0          # opacity threshold
+source :
+    mu_size     : 10            # number of angles for source integration
+spectrum :
+    p_size      : 60            # number of phot. impact parameters for spectrum
+    flatten     : No            # divide out continuum or not
+setups :
+    -   a0      :  1.0          # constant term
+        a1      :  0.0          # linear warp term
+        a2      :  0.0          # quadratic warp term
+        v_phot  :  10.0         # velocity at photosphere (kkm/s)
+        v_outer :  30.0         # outer velocity of line forming region (kkm/s)
+        t_phot  :  10.0         # blackbody photosphere temperature (kK)
+        ions    :  [ 1400, 1401, 1402, 1403 ]   # ions (100*Z+I, I=0 is neutral)
+        active  :  [  Yes,  Yes,  Yes,  Yes ]   # actually use the ion or not
+        log_tau :  [-99, -99 , -99, -99]   # ref. line opacity at v_ref
+        v_min   :  [ 10.0, 10.0, 10.0, 10.0 ]   # lower cutoff (kkm/s)
+        v_max   :  [ 30.0, 30.0, 30.0, 30.0 ]   # upper cutoff (kkm/s)
+        aux     :  [ 1e300 ,1e300, 1e300, 1e300 ]   # e-folding for exp form
+        temp    :  [ 10.0, 10.0, 10.0, 10.0 ]   # Boltzmann exc. temp. (kK)