Merge branch 'fastspec-pca-templates' of https://github.com/desihub/r…

…edrock into fastspec-pca-templates

Attempting merge after git reported conflict due to branch being
modified by me and someone else

py/redrock/external/desi.py

@@ -575,6 +575,15 @@ def rrdesi(options=None, comm=None):
     parser.add_argument("--chi2-scan", type=str, default=None,
         required=False, help="Load the chi2-scan from the input file")
 
+    parser.add_argument("--zscan-galaxy", type=str, default='-0.005,1.7,3e-4',
+        required=False, help="Redshift scan parameters for galaxies: zmin,zmax,dz")
+
+    parser.add_argument("--zscan-qso", type=str, default='0.05,6.0,5e-4',
+        required=False, help="Redshift scan parameters for QSO: zmin,zmax,dz")
+
+    parser.add_argument("--zscan-star", type=str, default='-0.002,0.00201,4e-5',
+        required=False, help="Redshift scan parameters for stars: zmin,zmax,dz")
+
     parser.add_argument("-n", "--ntargets", type=int,
         required=False, help="the number of targets to process in each file")
 
@@ -813,7 +822,12 @@ def rrdesi(options=None, comm=None):
         # Read the template data
         # Pass both use_gpu (this proc) and args.gpu (if any proc is using GPU)
         dtemplates = load_dist_templates(dwave, templates=args.templates,
-            comm=comm, mp_procs=mpprocs, redistribute=redistribute_templates, use_gpu=use_gpu, gpu_mode=args.gpu)
+                                         zscan_galaxy=args.zscan_galaxy,
+                                         zscan_qso=args.zscan_qso,
+                                         zscan_star=args.zscan_star,
+                                         comm=comm, mp_procs=mpprocs,
+                                         redistribute=redistribute_templates,
+                                         use_gpu=use_gpu, gpu_mode=args.gpu)
 
         # Compute the redshifts, including both the coarse scan and the
         # refinement.  This function only returns data on the rank 0 process.

py/redrock/templates.py

@@ -31,7 +31,8 @@ class Template(object):
 
     """
     def __init__(self, filename=None, spectype=None, redshifts=None,
-        wave=None, flux=None, subtype=None):
+                 wave=None, flux=None, subtype=None,
+                 zscan_galaxy=None, zscan_qso=None, zscan_star=None):
 
         if filename is not None:
             fx = None
@@ -73,15 +74,23 @@ def __init__(self, filename=None, spectype=None, redshifts=None,
             self._rrtype = hdr['RRTYPE'].strip().upper()
             if old_style_templates:
                 if self._rrtype == 'GALAXY':
-                    # redshifts = 10**np.arange(np.log10(1+0.005),
-                    # np.log10(1+2.0), 1.5e-4) - 1
-                    self._redshifts = 10**np.arange(np.log10(1-0.005),
-                        np.log10(1+7.0), 3e-4) - 1
+                    if zscan_galaxy is not None:
+                        zmin, zmax, dz = zscan_galaxy.split(',')
+                        self._redshifts = 10**np.arange(np.log10(1+float(zmin)), np.log10(1+float(zmax)), float(dz)) - 1
+                    else:
+                        self._redshifts = 10**np.arange(np.log10(1-0.005), np.log10(1+1.7), 3e-4) - 1
                 elif self._rrtype == 'STAR':
-                    self._redshifts = np.arange(-0.002, 0.00201, 4e-5)
+                    if zscan_star is not None:
+                        zmin, zmax, dz = zscan_star.split(',')
+                        self._redshifts = np.arange(float(zmin), float(zmax), float(dz))
+                    else:
+                        self._redshifts = np.arange(-0.002, 0.00201, 4e-5)
                 elif self._rrtype == 'QSO':
-                    self._redshifts = 10**np.arange(np.log10(1+0.05),
-                        np.log10(1+6.0), 5e-4) - 1
+                    if zscan_qso is not None:
+                        zmin, zmax, dz = zscan_qso.split(',')
+                        self._redshifts = 10**np.arange(np.log10(1+float(zmin)), np.log10(1+float(zmax)), float(dz)) - 1
+                    else:
+                        self._redshifts = 10**np.arange(np.log10(1+0.05), np.log10(1+6.0), 5e-4) - 1
                 else:
                     raise ValueError("Unknown redshift range to use for "
                         "template type {}".format(self._rrtype))
@@ -446,7 +455,9 @@ def cycle(self):
         return True
 
 
-def load_dist_templates(dwave, templates=None, comm=None, mp_procs=1, redistribute=False, use_gpu=False, gpu_mode=False):
+def load_dist_templates(dwave, templates=None, comm=None, mp_procs=1,
+                        zscan_galaxy=None, zscan_qso=None, zscan_star=None,
+                        redistribute=False, use_gpu=False, gpu_mode=False):
     """Read and distribute templates from disk.
 
     This reads one or more template files from disk and distributes them among
@@ -509,7 +520,8 @@ def load_dist_templates(dwave, templates=None, comm=None, mp_procs=1, redistribu
     template_data = list()
     if (comm is None) or (comm.rank == 0):
         for t in template_files:
-            template_data.append(Template(filename=t))
+            template_data.append(Template(filename=t, zscan_galaxy=zscan_galaxy,
+                                          zscan_star=zscan_star, zscan_qso=zscan_qso))
 
     if comm is not None:
         template_data = comm.bcast(template_data, root=0)

py/redrock/utils.py

@@ -15,8 +15,11 @@
 
 from . import constants
 
+<<<<<<< HEAD
 #IGM = Inoue14()
 
+=======
+>>>>>>> f1eaa9b7bad3c5a397a4ed53bf5a215a01a9f787
 class Inoue14(object):
     def __init__(self, scale_tau=1.):
         """
@@ -605,7 +608,14 @@ def distribute_work(nproc, ids, weights=None, capacities=None):
     return dist
 
 
+<<<<<<< HEAD
 def transmission_IGM_old(zObj, lObs, use_gpu=False):
+=======
+
+
+
+def transmission_Lyman(zObj, lObs, use_gpu=False):
+>>>>>>> f1eaa9b7bad3c5a397a4ed53bf5a215a01a9f787
     """Calculate the transmitted flux fraction from the Lyman series
     This returns the transmitted flux fraction:
     1 -> everything is transmitted (medium is transparent)
@@ -629,14 +639,15 @@ def transmission_IGM_old(zObj, lObs, use_gpu=False):
         scalar input; nz x nlambda in case of array input)
 
     """
-    if (use_gpu):
+    if use_gpu:
         import cupy as cp
         tile = cp.tile
         asarray = cp.asarray
     else:
         tile = np.tile
         asarray = np.asarray
 
+<<<<<<< HEAD
     #Lyman_series = constants.Lyman_series
     min_wave = 0
     if (np.isscalar(zObj)):
@@ -671,15 +682,48 @@ def transmission_IGM_old(zObj, lObs, use_gpu=False):
             T[iz, :] = IGM.full_IGM_old(myz, lRF[iz, :]) # .get() hack
     if use_gpu:
         T = asarray(T) # hack
+=======
+    ##Lyman_series = constants.Lyman_series
+    #min_wave = 0
+    #if np.isscalar(zObj):
+    #    #zObj is a float
+    #    lRF = lObs/(1.+zObj)
+    #else:
+    #    if (len(zObj) == 0):
+    #        #Empty z array
+    #        return np.ones((0, len(lObs)), dtype=np.float64)
+    #    #This is an array of float
+    #    min_wave = lObs.min()/(1+zObj.max())
+    #    #if (lObs.min()/(1+zObj.max()) > Lyman_series['Lya']['line']):
+    #    #if (min_wave > Lyman_series['Lya']['line']):
+    #    #    #Return None if wavelength range doesn't overlap with Lyman series
+    #    #    #No need to perform any calculations in this case
+    #    #    return None
+    #    if (not use_gpu and type(zObj) != np.ndarray):
+    #        #Cupy array passed??
+    #        zObj = zObj.get()
+    #    lObs = tile(lObs, (zObj.size, 1))
+    #    lRF = lObs/(1.+asarray(zObj)[:,None])
+
+    T = tile(np.ones_like(lObs), (zObj.size, 1))
+
+    W = np.where(lObs.min() / (1. + zObj) < 1215.67)[0]
+    if len(W) > 0:
+        for iz in W:
+            T[iz, :] = asarray(IGM.full_IGM(zObj[iz], lObs))
+
+>>>>>>> f1eaa9b7bad3c5a397a4ed53bf5a215a01a9f787
     #for l in list(Lyman_series.keys()):
     #    if (min_wave > Lyman_series[l]['line']):
     #        continue
     #    w      = lRF<Lyman_series[l]['line']
     #    zpix   = lObs[w]/Lyman_series[l]['line']-1.
     #    tauEff = Lyman_series[l]['A']*(1.+zpix)**Lyman_series[l]['B']
     #    T[w]  *= np.exp(-tauEff)
-    if (np.isscalar(zObj) and use_gpu):
-        T = asarray(T)
+
+    #if (np.isscalar(zObj) and use_gpu):
+    #    T = asarray(T)
+
     return T
 
 def transmission_IGM(zObj, lObs, use_gpu=False):