added full_galaxy_dir as argument

py/redrock/archetypes.py

@@ -49,7 +49,6 @@ def __init__(self, filename):
         self.wave = np.asarray(hdr['CRVAL1'] + hdr['CDELT1']*np.arange(self.flux.shape[1]))
         self.filename = filename
         ## Loading some global data once as it will be used many times
-        self.archetype_galaxies = params_for_all_galaxies() ## Dictionary for ELG, LRG, BGS
 
         if hdr['LOGLAM']:
             self.wave = 10**self.wave
@@ -89,7 +88,7 @@ def eval(self, subtype, dwave, coeff, wave, z):
         return flux
 
 
-    def get_best_archetype(self,spectra,weights,flux,wflux,dwave,z,legendre, nearest_nbh, n_nbh):
+    def get_best_archetype(self,spectra,weights,flux,wflux,dwave,z,legendre, nearest_nbh, n_nbh, archetype_data_dict):
         """Get the best archetype for the given redshift and spectype.
 
         Args:
@@ -102,6 +101,7 @@ def get_best_archetype(self,spectra,weights,flux,wflux,dwave,z,legendre, nearest
             legendre (dic): legendre polynomial
             nearest_nbh (bool): nearest-neighbour approach on archetypes
             n_nbh (int): number of nearest neighbours taken into account (if nearest_nbh is True)
+            archetype_data_dict (dict): dictionary containing the template archetype galaxy data (physical properties)
 
         Returns:
             chi2 (float): chi2 of best archetype
@@ -129,7 +129,7 @@ def get_best_archetype(self,spectra,weights,flux,wflux,dwave,z,legendre, nearest
             subtype = self._full_type[iBest].split(':::')[1].split('_')[0] #redrock best subtype
             if spectype=='GALAXY':
                 rrarchs_params = return_galaxy_archetype_properties(self.filename) 
-                new_arch, gal_inds = return_N_nearest_archetypes_from_synthetic_spectra(arch_id=iBest, archetype_data=rrarchs_params, gal_data=self.archetype_galaxies[subtype], n_nbh=n_nbh)
+                new_arch, gal_inds = return_N_nearest_archetypes_from_synthetic_spectra(arch_id=iBest, archetype_data=rrarchs_params, gal_data=archetype_data_dict[subtype], n_nbh=n_nbh)
                 new_arch = new_arch.astype('float64') # to let it work with trapz_rebin and Numba
                 for i in range(n_nbh):
                     binned = {hs:trapz_rebin(self.wave*(1.+z), new_arch[i], wave) for hs, wave in dwave.items()}

py/redrock/external/desi.py

@@ -41,6 +41,9 @@
 
 from ..archetypes import All_archetypes
 
+from ..nearest_neighbours import find_template_galaxy_archetypes
+
+from ..nearest_neighbours import params_for_all_galaxies
 
 def write_zbest(outfile, zbest, fibermap, exp_fibermap, tsnr2,
         template_version, archetype_version,
@@ -563,6 +566,10 @@ def rrdesi(options=None, comm=None):
     parser.add_argument("-n_nbh", type=int, default=9,
         required=False, help="if nearest_nbh True, N-nearest neighbours taken into account (default is 9)")
 
+    parser.add_argument("--full_galaxy_dir", type=str, default=None,
+        required=False,
+        help="template archetype galaxy directory for nearest_nbh approach")
+
     parser.add_argument("-d", "--details", type=str, default=None,
         required=False, help="output file for full redrock fit details")
 
@@ -701,12 +708,13 @@ def rrdesi(options=None, comm=None):
                     comm.Abort()
                 else:
                     sys.exit(1)
+            tempfilenames = find_template_galaxy_archetypes(args.full_galaxy_dir)
             print('\n=================\n')    
-            print('REDSHIFT: Archetype with galaxy properties is provided\n')
+            print('REDSHIFT: Archetype with galaxy properties is provided and full_galaxy_dir with archetype data is also found\n')
             print('REDSHIFT: Nearest neighbour approach is provided, so will apply the N-nearest neighbour approach on Redrock\n')
             print('REDSHIFT: %d nearest neighbours will be used...\n'%(args.n_nbh))
             print('===================\n')
-
+            archetype_data_dict = params_for_all_galaxies(tempfilenames)
 
     targetids = None
     if args.targetids is not None:
@@ -830,7 +838,7 @@ def rrdesi(options=None, comm=None):
 
         scandata, zfit = zfind(targets, dtemplates, mpprocs,
             nminima=args.nminima, archetypes=args.archetypes, 
-            nearest_nbh=args.nearest_nbh, n_nbh=args.n_nbh,
+            nearest_nbh=args.nearest_nbh, n_nbh=args.n_nbh, archetype_data_dict=archetype_data_dict,
             priors=args.priors, chi2_scan=args.chi2_scan, use_gpu=use_gpu)
 
         stop = elapsed(start, "Computing redshifts", comm=comm)

py/redrock/fitz.py

@@ -107,7 +107,7 @@ def minfit(x, y):
     return (x0, xerr, y0, zwarn)
 
 
-def fitz(zchi2, redshifts, spectra, template, nminima=3, archetype=None, nearest_nbh=False, n_nbh=9, use_gpu=False):
+def fitz(zchi2, redshifts, spectra, template, nminima=3, archetype=None, nearest_nbh=False,archetype_data_dict=None, n_nbh=9, use_gpu=False):
     """Refines redshift measurement around up to nminima minima.
 
     TODO:
@@ -121,6 +121,9 @@ def fitz(zchi2, redshifts, spectra, template, nminima=3, archetype=None, nearest
         template (Template): the template for this fit.
         nminima (int): the number of minima to consider.
         use_gpu (bool): use GPU or not
+        nearest_nbh (bool), for applying nearest neighbour approach (default False)
+        n_nbh (int); if above true, then number of nearest neighbour to take into account
+        archetype_data_dict (dict): dictionary containing the template archetype galaxy data
 
     Returns:
         Table: the fit parameters for the minima.
@@ -245,7 +248,7 @@ def fitz(zchi2, redshifts, spectra, template, nminima=3, archetype=None, nearest
                 chi2=chi2min, zz=zz, zzchi2=zzchi2,
                 coeff=coeff))
         else:
-            chi2min, coeff, fulltype = archetype.get_best_archetype(spectra,weights,flux,wflux,dwave,zbest,legendre, nearest_nbh, n_nbh)
+            chi2min, coeff, fulltype = archetype.get_best_archetype(spectra,weights,flux,wflux,dwave,zbest,legendre, nearest_nbh, n_nbh, archetype_data_dict)
 
             results.append(dict(z=zbest, zerr=zerr, zwarn=zwarn,
                 chi2=chi2min, zz=zz, zzchi2=zzchi2,

py/redrock/nearest_neighbours.py

@@ -16,14 +16,17 @@ def read_fits_data(filename, nhdu):
     return data
 
 
-def synthetic_galaxy_data_read(gals):
+def synthetic_galaxy_data_read(filename):
 
     """ Returns fluxes and physical parameters (dictionary) 
         for synthetic galaxy spectra generated with desisim
         by Abhijeet Anand, LBL, email: AbhijeetAnand@lbl.gov
+
+        Input:
+            filename (str): filename containing template galaxy archetype data
     """
     all_gal_data = {}
-    file_gals = '/global/cfs/cdirs/desi/users/abhijeet/synthetic_spectra/%s/%s_spectral_flux_rest_frame.fits'%(gals, gals)
+    file_gals = filenme
     temp_gal_data = read_fits_data(file_gals, nhdu=2)
     all_gal_data['FLUX'] = fits.open(file_gals)[1].data
     for key in temp_gal_data.dtype.names:
@@ -85,15 +88,53 @@ def return_galaxy_archetype_properties(archetypes):
         params[key] = data[key]
     return params
 
-def params_for_all_galaxies():
+def params_for_all_galaxies(lstfilenames):
     """
         Returns dictionaries for all galaxies (ELG, LRG, BGS)
         access each dictionary as fin_dict[subtype]
+        Input:
+            lstfilenames (list of filenames): filename containing template galaxy archetype data
+            Naming convention used: {gals}_..*.fits, where gals = ELG, LRG or BGS
     """
     synthetic_galaxies = {}
-    for galaxies in ['ELG', 'LRG', 'BGS']:
-        synthetic_galaxies[galaxies] = synthetic_galaxy_data_read(galaxies)
+    for filename in lstfilenames:
+        galaxies = filename[0:3] # will be as ELG, LRG or BGS
+        synthetic_galaxies[galaxies] = synthetic_galaxy_data_read(filename)
     return synthetic_galaxies
 
 
+def find_template_galaxy_archetypes(full_galaxy_dir=None):
+    """Return list of  \*_spectral_flux_rest_frame.fits, template galaxy data for ELG, LRG, BGS
+
+    Search directories in this order, returning results from first one found:
+        - full_galaxy_dir
+        - $RR_ARCHDATA_DIR
+
+    Args:
+       full_galaxy_dir (str): optional directory containing the template galaxy data for ELG, LRG, BGS
+
+    Returns:
+        list: a list of template galaxy data files.
+
+    """
+    if full_galaxy_dir is None:
+        if 'RR_ARCHDATA_DIR' in os.environ:
+            full_galaxy_dir = os.environ['RR_ARCHDATA_DIR']
+        else:
+            thisdir = os.path.dirname(__file__)
+            archdir = os.path.join(os.path.abspath(thisdir), 'full_galaxy_dir')
+            if os.path.exists(archdir):
+                full_galaxy_dir = archdir
+            else:
+                raise IOError("ERROR: can't find full_galaxy_dir, $RR_ARCHDATA_DIR")
+                sys.exit(1)
+        lstfilename = sorted(glob(os.path.join(full_galaxy_dir, '*.fits')))
+    else:
+        full_galaxy_dir_expand = os.path.expandvars(full_galaxy_dir)
+        lstfilename = glob(os.path.join(afull_galaxy_dir_expand, '*.fits'))
+        lstfilename = sorted([ f.replace(full_galaxy_dir_expand,full_galaxy_dir) for f in lstfilename])
+
+    return lstfilename
+
+
 

py/redrock/zfind.py

@@ -78,7 +78,7 @@ def sort_dict_by_cols(d, colnames, sort_first_column_first = True):
         d[k] = d[k][idx]
     return
 
-def _mp_fitz(chi2, target_data, t, nminima, qout, archetype, nearest_nbh, n_nbh, use_gpu):
+def _mp_fitz(chi2, target_data, t, nminima, qout, archetype, nearest_nbh, n_nbh, archetype_data_dict, use_gpu):
     """Wrapper for multiprocessing version of fitz.
     """
     try:
@@ -88,7 +88,7 @@ def _mp_fitz(chi2, target_data, t, nminima, qout, archetype, nearest_nbh, n_nbh,
         results = list()
         for i, tg in enumerate(target_data):
             zfit = fitz(chi2[i], t.template.redshifts, tg.spectra,
-                t.template, nminima=nminima, archetype=archetype, nearest_nbh=nearest_nbh, n_nbh=n_nbh, use_gpu=use_gpu)
+                t.template, nminima=nminima, archetype=archetype, nearest_nbh=nearest_nbh, n_nbh=n_nbh, archetype_data_dict=archetype_data_dict, use_gpu=use_gpu)
             results.append( (tg.id, zfit) )
         qout.put(results)
     except:
@@ -208,7 +208,7 @@ def sort_zfit_dict(zfit):
     sort_dict_by_cols(zfit, ('__badfit__', 'chi2'), sort_first_column_first=True)
     zfit.pop('__badfit__')
 
-def zfind(targets, templates, mp_procs=1, nminima=3, archetypes=None, nearest_nbh=False, n_nbh=9, priors=None, chi2_scan=None, use_gpu=False):
+def zfind(targets, templates, mp_procs=1, nminima=3, archetypes=None, nearest_nbh=False, n_nbh=9, archetype_data_dict=archetype_data_dict, priors=None, chi2_scan=None, use_gpu=False):
     """Compute all redshift fits for the local set of targets and collect.
 
     Given targets and templates distributed across a set of MPI processes,
@@ -232,6 +232,7 @@ def zfind(targets, templates, mp_procs=1, nminima=3, archetypes=None, nearest_nb
             to use for final fitz choice of best chi2 vs. z minimum.
         nearest_nbh (bool), for applying nearest neighbour approach (default False)
         n_nbh (int); if above true, then number of nearest neighbour to take into account
+        archetype_data_dict (dict): dictionary containing the template archetype galaxy data
         priors (str, optional): file containing redshift priors
         chi2_scan (str, optional): file containing already computed chi2 scan
         use_gpu (bool, optional): use gpu for calc_zchi2
@@ -333,7 +334,7 @@ def zfind(targets, templates, mp_procs=1, nminima=3, archetypes=None, nearest_nb
                 zfit = fitz(results[tg.id][ft]['zchi2'] \
                     + results[tg.id][ft]['penalty'],
                     t.template.redshifts, tg.spectra,
-                    t.template, nminima=nminima,archetype=archetype, nearest_nbh=nearest_nbh, n_nbh=n_nbh, use_gpu=use_gpu)
+                    t.template, nminima=nminima,archetype=archetype, nearest_nbh=nearest_nbh, n_nbh=n_nbh, archetype_data_dict=archetype_data_dict, use_gpu=use_gpu)
                 results[tg.id][ft]['zfit'] = zfit
         else:
             # Multiprocessing case.