From 17414b51db257566076cdda3bb56ca594e771a70 Mon Sep 17 00:00:00 2001 From: qezlou Date: Fri, 15 Nov 2024 22:48:25 -0600 Subject: [PATCH] small edits for: 1- Being able to pass exact nbins intead of velocity width 2- Explicitly turning the self-shiedling off by setting `turn_off_selfshied=True`. This is done by passing gamma (the Lorentzian width of the line) to 0. --- fake_spectra/griddedspectra.py | 37 ++++++++++++++++++++++++++++++---- fake_spectra/spectra.py | 25 ++++++++++++++++------- 2 files changed, 51 insertions(+), 11 deletions(-) diff --git a/fake_spectra/griddedspectra.py b/fake_spectra/griddedspectra.py index 3db5280..05d3edb 100644 --- a/fake_spectra/griddedspectra.py +++ b/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): diff --git a/fake_spectra/spectra.py b/fake_spectra/spectra.py index 2a07c70..1e50743 100644 --- a/fake_spectra/spectra.py +++ b/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