Merge pull request #72 from qezlou/no-selfshielding

small edits for:

fake_spectra/griddedspectra.py

@@ -9,8 +9,36 @@
 class GriddedSpectra(spectra.Spectra):
     """Generate regular grid of spectra along a given axis."""
 
-    def __init__(self,num, base, nspec=200, MPI=None, res = None, savefile="gridded_spectra.hdf5", reload_file=True, axis=1, **kwargs):
-
+    def __init__(self,num, base, nspec=200, MPI=None, res = None, nbins=None, savefile="gridded_spectra.hdf5", reload_file=True, grid_axes=None, grid_cofm=None, axis=1, **kwargs):
+        """
+        Parameters
+        ----------
+        num : int, required
+            Snapshot number
+        base : str, required
+            Base directory of the snapshot, i.e. containing .hdf5 or PART_XXXX files
+        nspec : int, optional
+            Number of skewers per side of the grid in the transverse direction
+        MPI : object, optional
+            MPI communicator
+        res : float, required if `nbins` is not provided
+            Resolution of the skewers in km/s
+        nbins : int, required if `res` is not provided
+            Number of bins in the skewers
+        savefile : str, optional
+            Name of the file to save the spectra
+        reload_file : bool, optional
+            Try to load the spectra file if it exists
+        grid_axes : array of shape (nspec*nspec*nspec,3), optional
+            Axes of the grid. Onlly if not a unifrom grid across the box is desired.
+            if none, `get_axes_and_cofm` will be called to get the axes
+        grid_cofm : array of shape (nspec*nspec*nspec,3), optional
+            Centers of the grid. Onlly if not a unifrom grid across the box is desired
+            if none, `get_axes_and_cofm` will be called to get the centers
+        axis : int, optional
+            Axis along which to place the skewers. 
+            Only if `grid_axes` and `grid_cofm` are not provided.
+        """
         if reload_file is False:
             # if reading skewers from file, no need to read the snapshot
             grid_cofm=None
@@ -21,10 +49,11 @@ def __init__(self,num, base, nspec=200, MPI=None, res = None, savefile="gridded_
             self.box = f.get_header_attr("BoxSize")
             del f
             # get position of skewers in the grid
-            grid_axes, grid_cofm = self.get_axes_and_cofm(nspec,axis)
+            if grid_cofm is None:
+                grid_axes, grid_cofm = self.get_axes_and_cofm(nspec,axis)
 
         # call constructor of base class
-        spectra.Spectra.__init__(self,num,base,cofm=grid_cofm,axis=grid_axes,MPI=MPI, res=res,savefile=savefile,reload_file=reload_file,**kwargs)
+        spectra.Spectra.__init__(self,num,base,cofm=grid_cofm,axis=grid_axes,MPI=MPI, res=res, nbins=nbins, savefile=savefile,reload_file=reload_file,**kwargs)
 
 
     def get_axes_and_cofm(self,nspec,axis):

fake_spectra/spectra.py

@@ -82,8 +82,8 @@ class Spectra:
                      when converting comoving to velocity units.
                      If not proviced, will assume flat LCDM and compute it.
     """
-    def __init__(self, num, base, cofm, axis, MPI=None, res=1., cdir=None, savefile="spectra.hdf5",
-                 savedir=None, reload_file=False, spec_res = 0,load_halo=False, units=None, sf_neutral=True,
+    def __init__(self, num, base, cofm, axis, MPI=None, nbins=None, res=1., cdir=None, savefile="spectra.hdf5",
+                 savedir=None, reload_file=False, spec_res = 0,load_halo=False, units=None, sf_neutral=True, turn_off_selfshield=False,
                  quiet=False, load_snapshot=True, gasprop=None, gasprop_args=None, kernel=None, use_external_Hz=None):
 
         #Present for compatibility. Functionality moved to HaloAssignedSpectra
@@ -124,6 +124,8 @@ def __init__(self, num, base, cofm, axis, MPI=None, res=1., cdir=None, savefile=
         self.discarded=0
         self.npart=0
 
+        self.turn_off_selfshield = turn_off_selfshield
+
         self.spec_res = spec_res
         self.cdir = cdir
         #Minimum length of spectra within which to look at metal absorption (in km/s)
@@ -218,9 +220,14 @@ def __init__(self, num, base, cofm, axis, MPI=None, res=1., cdir=None, savefile=
         if reload_file:
             # if reloading from snapshot, pixel width must have been defined
             if res is None:
-                raise ValueError('pixel width (res) not provided')
-            # nbins must be an integer
-            self.nbins=int(self.vmax/res)
+                if nbins is not None:
+                   self.nbins = nbins
+                   res = self.vmax/(1.*nbins)
+                else:
+                   raise ValueError('pixel width (res) not provided')
+            if nbins is None:
+                # nbins must be an integer
+                self.nbins=int(self.vmax/res)
             # velocity bin size (kms^-1)
             self.dvbin = self.vmax / (1.*self.nbins)
         else:
@@ -659,7 +666,11 @@ def _get_elem_den(self, elem, ion, den, temp, ind, ind2):
     def _do_interpolation_work(self, pos, vel, elem_den, temp, hh, amumass, line, get_tau):
         """Run the interpolation on some pre-determined arrays, spat out by _read_particle_data"""
         #Factor of 10^-8 converts line width (lambda_X) from Angstrom to cm
-        return _Particle_Interpolate(get_tau*1, self.nbins, self.kernel_int, self.box, self.velfac, self.atime, line.lambda_X*1e-8, line.gamma_X, line.fosc_X, amumass, self.tautail, pos, vel, elem_den, temp, hh, self.axis, self.cofm)
+        if self.turn_off_selfshield:
+            gamma_X = 0
+        else:
+            gamma_X = line.gamma_X
+        return _Particle_Interpolate(get_tau*1, self.nbins, self.kernel_int, self.box, self.velfac, self.atime, line.lambda_X*1e-8, gamma_X, line.fosc_X, amumass, self.tautail, pos, vel, elem_den, temp, hh, self.axis, self.cofm)
 
     def particles_near_lines(self, pos, hh, axis=None, cofm=None):
         """Filter a particle list, returning an index list of those near sightlines"""
@@ -807,7 +818,7 @@ def compute_spectra(self, elem, ion, ll, get_tau):
         #Do remaining files
         for nn in xrange(1, nsegments):
             tresult = self._interpolate_single_file(nn, elem, ion, ll, get_tau)
-            print("Interpolation %.1f percent done" % (100*nn/nsegments), flush=True)
+            print(f"rank = {self.rank} | "+"Interpolation %.1f percent done" % (100*nn/nsegments), flush=True)
             #Add new file
             result += tresult
             del tresult