args added for default archetype values, removed hard coded ncamera

py/redrock/archetypes.py

@@ -161,7 +161,7 @@ def eval(self, subtype, dwave, coeff, wave, z):
 
         return flux
 
-    def nearest_neighbour_model(self, target,weights,flux,wflux,dwave,z, n_nearest, zzchi2, trans, per_camera, dedges=None, binned=None, use_gpu=False, prior=None):
+    def nearest_neighbour_model(self, target,weights,flux,wflux,dwave,z, n_nearest, zzchi2, trans, per_camera, dedges=None, binned=None, use_gpu=False, prior=None, ncam=None):
 
         """
         
@@ -181,6 +181,7 @@ def nearest_neighbour_model(self, target,weights,flux,wflux,dwave,z, n_nearest,
         binned (dict): already computed dictionary of rebinned fluxes
         use_gpu (bool): use GPU or not
         prior (2d array): prior matrix on coefficients (1/sig**2)
+        ncam (int): Number of camera for given Instrument
 
         
         Returns:
@@ -222,7 +223,7 @@ def nearest_neighbour_model(self, target,weights,flux,wflux,dwave,z, n_nearest,
             nbasis = tdata[hs].shape[2]
         if per_camera:
             #Use CPU mode since small tdata
-            (zzchi2, zzcoeff) = per_camera_coeff_with_least_square_batch(spectra, tdata, weights, flux, wflux, nleg, n_nearest, method='bvls', n_nbh=n_nearest, prior=prior, use_gpu=False)
+            (zzchi2, zzcoeff) = per_camera_coeff_with_least_square_batch(spectra, tdata, weights, flux, wflux, nleg, n_nearest, method='bvls', n_nbh=n_nearest, prior=prior, use_gpu=False, ncam=ncam)
         else:
             #Use CPU mode for calc_zchi2 since small tdata
             (zzchi2, zzcoeff) = calc_zchi2_batch(spectra, tdata, weights, flux, wflux, 1, nbasis, use_gpu=False)
@@ -276,7 +277,7 @@ def get_best_archetype(self,target,weights,flux,wflux,dwave,z, per_camera, n_nea
             dedges = None
 
         if per_camera:
-            ncam=3 # b, r, z cameras
+            ncam=len(list(dwave.keys())) # for e.g., DESI has three cameras namely b, r, z
         else:
             ncam = 1 # entire spectra
 
@@ -318,17 +319,17 @@ def get_best_archetype(self,target,weights,flux,wflux,dwave,z, per_camera, n_nea
             nbasis = tdata[hs].shape[2]
         if per_camera:
             if (use_gpu):
-                (zzchi2, zzcoeff) = per_camera_coeff_with_least_square_batch(target, tdata, gpuweights, gpuflux, gpuwflux, nleg, self._narch, method=solve_method, n_nbh=1, prior=prior, use_gpu=use_gpu)
+                (zzchi2, zzcoeff) = per_camera_coeff_with_least_square_batch(target, tdata, gpuweights, gpuflux, gpuwflux, nleg, self._narch, method=solve_method, n_nbh=1, prior=prior, use_gpu=use_gpu, ncam=ncam)
             else:
-                (zzchi2, zzcoeff) = per_camera_coeff_with_least_square_batch(target, tdata, weights, flux, wflux, nleg, self._narch, method=solve_method, n_nbh=1, prior=prior, use_gpu=use_gpu)
+                (zzchi2, zzcoeff) = per_camera_coeff_with_least_square_batch(target, tdata, weights, flux, wflux, nleg, self._narch, method=solve_method, n_nbh=1, prior=prior, use_gpu=use_gpu, ncam=ncam)
         else:
             if (use_gpu):
                 (zzchi2, zzcoeff) = calc_zchi2_batch(spectra, tdata, gpuweights, gpuflux, gpuwflux, self._narch, nbasis, use_gpu=use_gpu)
             else:
                 (zzchi2, zzcoeff) = calc_zchi2_batch(spectra, tdata, weights, flux, wflux, self._narch, nbasis, use_gpu=use_gpu)
 
         if n_nearest is not None:
-            best_chi2, best_coeff, best_fulltype = self.nearest_neighbour_model(target,weights,flux,wflux,dwave,z, n_nearest, zzchi2, trans, per_camera, dedges=dedges, binned=binned, use_gpu=use_gpu, prior=prior)
+            best_chi2, best_coeff, best_fulltype = self.nearest_neighbour_model(target,weights,flux,wflux,dwave,z, n_nearest, zzchi2, trans, per_camera, dedges=dedges, binned=binned, use_gpu=use_gpu, prior=prior, ncam=ncam)
             return best_chi2, best_coeff, best_fulltype
         else:
             iBest = np.argmin(zzchi2)

py/redrock/external/desi.py

@@ -557,18 +557,21 @@ def rrdesi(options=None, comm=None):
         required=False,
         help="archetype file or directory for final redshift comparison")
 
-    parser.add_argument("-deg_legendre", type=int, default=3,
-        required=False, help="if archetypes are provided legendre polynomials upto deg_legendre-1 will be used (default is 3)")
+    parser.add_argument("--archetypes-legendre", default=False, action="store_true",
+        required=False, help="Use this flag with archetypes if want to run archetype percamera approach with default values")
+
+    parser.add_argument("--archetype-legendre-degree", type=int, default=2,
+        required=False, help="if archetypes are provided legendre polynomials upto deg_legendre-1 will be used (default is 2)")
 
-    parser.add_argument("-n_nearest", type=int, default=None,
+    parser.add_argument("--archetype-nnearest", type=int, default=None,
         required=False, help="number of nearest archetypes (in chi2 space) to be used in archetype modeling, must be greater than 1 (default is None)")
 
-    parser.add_argument("-nz", type=int, default=15,
+    parser.add_argument("-zminfit_npoints", type=int, default=15,
         required=False, help="number of finer redshift to be used around best fit redshifts (default is 15)")
 
-    parser.add_argument("--per-camera", default=False, action="store_true",
+    parser.add_argument("--archetype-legendre-percamera", default=True, action="store_true",
         required=False, help="If True, in archetype mode the fitting will be done for each camera/band")
-
+    
     parser.add_argument("-d", "--details", type=str, default=None,
         required=False, help="output file for full redrock fit details")
 
@@ -593,8 +596,8 @@ def rrdesi(options=None, comm=None):
     parser.add_argument("--nminima", type=int, default=3,
         required=False, help="the number of redshift minima to search")
 
-    parser.add_argument("--prior_sigma", type=float, default=None,
-        required=False, help="sigma to add as prior in solving linear equation, 1/sig**2 will be added")
+    parser.add_argument("--archetype-legendre-prior", type=float, default=0.1,
+        required=False, help="sigma to add as prior in solving linear equation, 1/sig**2 will be added, default is 0.1")
 
     parser.add_argument("--allspec", default=False, action="store_true",
         required=False, help="use individual spectra instead of coadd")
@@ -713,15 +716,7 @@ def rrdesi(options=None, comm=None):
 
                 if os.path.isfile(args.archetypes):
                     print('Archetype is a file and it exists and readable\n')
-                    print('Archetype will only be applied to that spectype\n')
-                if args.deg_legendre>0:
-                    print('legendre polynomials of degrees %s will be added to Archetypes\n'%([i for i in range(args.deg_legendre)]))
-                else:
-                    print('No Legendre polynomial will be added to archetypes...\n')
-                if args.n_nearest:
-                    print('n_nearest argument is provided so %d nearest neighbour (in chi2 space) to the best archetype will be used for further modeling...\n'%(args.n_nearest-1))
-                if args.prior_sigma:
-                    print('prior_sigma = %s has been provided, so a prior will be added while solving for the coefficients\n'%(args.prior_sigma))
+                    print('Archetype will only be applied to SPECTYPE %s\n'%(os.path.basename(args.archetypes).split('-')[1].split('.')[0].upper()))
             else:
                 print("ERROR: can't find archetypes_dir or it is unreadable\n")
                 sys.stdout.flush()
@@ -836,7 +831,37 @@ def rrdesi(options=None, comm=None):
 
         # Get the dictionary of wavelength grids
         dwave = targets.wavegrids()
+
+        ncamera = len(list(dwave.keys())) # number of cameras for given instrument
+        if args.archetypes_legendre:
+            print('\n--archetypes-legendre argument is provided so using default values..\n')
+            if ncamera>1: 
+                archetype_legendre_percamera = True
+                print('Number of cameras = %d, percamera fitting will be done\n'%(ncamera))
+            else:
+                archetype_legendre_percamera = False
+                print('No per camera fitting will be done..\n')
+            if args.archetype_legendre_degree>0:
+                print('legendre polynomials of degrees %s will be added to Archetypes\n'%([i for i in range(args.archetype_legendre_degree)]))
+            else:
+                print('No Legendre polynomial will be added to archetypes...\n')
+            archetype_legendre_prior = 0.1
+            print('archetype_legendre_prior = %s has been provided, so a prior will be added while solving for the coefficients\n'%(args.archetype_legendre_prior))
+        else:
+            archetype_legendre_prior = None
+            print('\nNo prior sigma will be added while solving for Archetypes coefficients\n')
+            if args.archetype_legendre_degree>0:
+                print('legendre polynomials of degrees %s will be added to Archetypes\n'%([i for i in range(args.archetype_legendre_degree)]))
+            else:
+                print('No Legendre polynomial will be added to archetypes...\n')
 
+            if ncamera>1:
+                archetype_legendre_percamera = True
+                print('Number of cameras = %d, percamera fitting will be done\n'%(ncamera))
+            else:
+                archetype_legendre_percamera = False
+                print('Number of cameras = %d, No per camera fitting will be done..\n'%(ncamera))
+
         stop = elapsed(start, "Read and distribution of {} targets"\
             .format(len(targets.all_target_ids)), comm=comm)
 
@@ -852,7 +877,7 @@ def rrdesi(options=None, comm=None):
 
         scandata, zfit = zfind(targets, dtemplates, mpprocs,
             nminima=args.nminima, archetypes=args.archetypes,
-            priors=args.priors, chi2_scan=args.chi2_scan, use_gpu=use_gpu, nz=args.nz, per_camera=args.per_camera, deg_legendre=args.deg_legendre, n_nearest=args.n_nearest, prior_sigma=args.prior_sigma)
+            priors=args.priors, chi2_scan=args.chi2_scan, use_gpu=use_gpu, zminfit_npoints=args.zminfit_npoints, per_camera=archetype_legendre_percamera, deg_legendre=args.archetype_legendre_degree, n_nearest=args.archetype_nnearest, prior_sigma=archetype_legendre_prior, ncamera=ncamera)
 
         stop = elapsed(start, "Computing redshifts", comm=comm)
 

py/redrock/fitz.py

@@ -116,16 +116,27 @@ def legendre_calculate(nleg, dwave):
 
     return legendre
 
-def prior_on_coeffs(n_nbh, deg_legendre, sigma):
+def prior_on_coeffs(n_nbh, deg_legendre, sigma, ncamera):
+
+    """
+    Args:
+        n_nbh (int): number of dominant archetypes
+        deg_legendre (int): number of Legendre polynomials 
+        sigma (int): prior sigma to be used for archetype fitting
+        ncamera (int): number of cameras for given instrument
+    Returns:
+        2d array to be added while solving for archetype fitting
 
-    nbasis = n_nbh+deg_legendre*3 # 3 desi cameras
+    """
+
+    nbasis = n_nbh+deg_legendre*ncamera # 3 desi cameras
     prior = np.zeros((nbasis, nbasis), dtype='float64');np.fill_diagonal(prior, 1/(sigma**2))
     for i in range(n_nbh):
         prior[i][i]=0. ## Do not add prior to the archetypes, added only to the Legendre polynomials
     return prior
 
 
-def fitz(zchi2, redshifts, target, template, nminima=3, archetype=None, use_gpu=False, deg_legendre=None, nz=15, per_camera=False, n_nearest=None, prior_sigma=None):
+def fitz(zchi2, redshifts, target, template, nminima=3, archetype=None, use_gpu=False, deg_legendre=None, zminfit_npoints=15, per_camera=False, n_nearest=None, prior_sigma=None):
     """Refines redshift measurement around up to nminima minima.
 
     TODO:
@@ -140,7 +151,7 @@ def fitz(zchi2, redshifts, target, template, nminima=3, archetype=None, use_gpu=
         nminima (int): the number of minima to consider.
         use_gpu (bool): use GPU or not
         deg_legendre (int): in archetype mode polynomials upto deg_legendre-1 will be used
-        nz (int): number of finer redshift pixels to search for final redshift - default 15
+        zminfit_npoints (int): number of finer redshift pixels to search for final redshift - default 15
         per_camera: (bool): True if fitting needs to be done in each camera for archetype mode
         n_nearest: (int): number of nearest neighbours to be used in chi2 space (including best archetype) 
         prior_sigma (float): prior to add in the final solution matrix: added as 1/(prior_sigma**2) only for per-camera mode
@@ -177,8 +188,10 @@ def fitz(zchi2, redshifts, target, template, nminima=3, archetype=None, use_gpu=
 
     results = list()
     #Moved default nz to arg list
-    if (nz is None):
+    if (zminfit_npoints is None):
         nz = 15
+    else:
+        nz = zminfit_npoints
 
     for imin in find_minima(zchi2):
         if len(results) == nminima:
@@ -302,10 +315,14 @@ def fitz(zchi2, redshifts, target, template, nminima=3, archetype=None, use_gpu=
                 coeff=coeff))
         else:
             if prior_sigma is not None:
+                if per_camera:
+                    ncamera = len(list(dwave.keys())) # number of cameras, for e.g. DESI has three cameras
+                else:
+                    ncamera = 1
                 if n_nearest is None:
-                    prior = prior_on_coeffs(1, deg_legendre, prior_sigma)
+                    prior = prior_on_coeffs(1, deg_legendre, prior_sigma, ncamera)
                 else:
-                    prior = prior_on_coeffs(n_nearest, deg_legendre, prior_sigma)
+                    prior = prior_on_coeffs(n_nearest, deg_legendre, prior_sigma, ncamera)
             else:
                 prior=None
             chi2min, coeff, fulltype = archetype.get_best_archetype(target,weights,flux,wflux,dwave,zbest, per_camera, n_nearest, trans=trans, use_gpu=use_gpu, prior=prior)

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, use_gpu, deg_legendre, nz, per_camera, n_nearest, prior_sigma):
+def _mp_fitz(chi2, target_data, t, nminima, qout, archetype, use_gpu, deg_legendre, zminfit_npoints, per_camera, n_nearest, prior_sigma):
     """Wrapper for multiprocessing version of fitz.
     """
     try:
@@ -88,7 +88,7 @@ def _mp_fitz(chi2, target_data, t, nminima, qout, archetype, use_gpu, deg_legend
         results = list()
         for i, tg in enumerate(target_data):
             zfit = fitz(chi2[i], t.template.redshifts, tg,
-                t.template, nminima=nminima, archetype=archetype, use_gpu=use_gpu, nz=nz, per_camera=per_camera, deg_legendre=deg_legendre, n_nearest=n_nearest, prior_sigma=prior_sigma)
+                t.template, nminima=nminima, archetype=archetype, use_gpu=use_gpu, zminfit_npoints=zminfit_npoints, per_camera=per_camera, deg_legendre=deg_legendre, n_nearest=n_nearest, prior_sigma=prior_sigma)
             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, priors=None, chi2_scan=None, use_gpu=False, nz=15, per_camera=None, deg_legendre=None, n_nearest=None, prior_sigma=None):
+def zfind(targets, templates, mp_procs=1, nminima=3, archetypes=None, priors=None, chi2_scan=None, use_gpu=False, zminfit_npoints=15, per_camera=None, deg_legendre=None, n_nearest=None, prior_sigma=None, ncamera=None):
     """Compute all redshift fits for the local set of targets and collect.
 
     Given targets and templates distributed across a set of MPI processes,
@@ -234,10 +234,11 @@ def zfind(targets, templates, mp_procs=1, nminima=3, archetypes=None, priors=Non
         chi2_scan (str, optional): file containing already computed chi2 scan
         use_gpu (bool, optional): use gpu for calc_zchi2
         deg_legendre (int): in archetype mode polynomials upto deg_legendre-1 will be used
-        nz (int): number of finer redshift pixels to search for final redshift
+        zminfit_npoints (int): number of finer redshift pixels to search for final redshift
         per_camera: (bool): True if fitting needs to be done in each camera for archetype mode
         n_nearest (int): number of nearest neighbours to be used in chi2 space (including best archetype)
         prior_sigma (float): prior to add in the final solution matrix: added as 1/(prior_sigma**2) only for per-camera mode
+        zminfit_npoints (int): number of minimum redshift to be fit for final redshift estimation
 
     Returns:
         tuple: (allresults, allzfit), where "allresults" is a dictionary of the
@@ -349,7 +350,7 @@ def zfind(targets, templates, mp_procs=1, nminima=3, archetypes=None, priors=Non
                 zfit = fitz(results[tg.id][ft]['zchi2'] \
                     + results[tg.id][ft]['penalty'],
                     t.template.redshifts, tg,
-                    t.template, nminima=nminima,archetype=archetype, use_gpu=use_gpu, deg_legendre=deg_legendre, nz=nz, per_camera=per_camera, n_nearest=n_nearest, prior_sigma=prior_sigma)
+                    t.template, nminima=nminima,archetype=archetype, use_gpu=use_gpu, deg_legendre=deg_legendre, zminfit_npoints=zminfit_npoints, per_camera=per_camera, n_nearest=n_nearest, prior_sigma=prior_sigma)
                 results[tg.id][ft]['zfit'] = zfit
         else:
             # Multiprocessing case.
@@ -373,7 +374,7 @@ def zfind(targets, templates, mp_procs=1, nminima=3, archetypes=None, priors=Non
                     eff_chi2[i,:] = results[tg.id][ft]['zchi2'] \
                         + results[tg.id][ft]['penalty']
                 p = mp.Process(target=_mp_fitz, args=(eff_chi2,
-                    target_data, t, nminima, qout, archetype, use_gpu, deg_legendre, nz, per_camera, n_nearest, prior_sigma))
+                    target_data, t, nminima, qout, archetype, use_gpu, deg_legendre, zminfit_npoints, per_camera, n_nearest, prior_sigma))
                 procs.append(p)
                 p.start()
 
@@ -587,20 +588,20 @@ def zfind(targets, templates, mp_procs=1, nminima=3, archetypes=None, priors=Non
             del allresults[tid]['meta']
 
         # Standardize column sizes
-        if allzfit['subtype'].dtype != '<U32':
-            allzfit.replace_column('subtype', allzfit['subtype'].astype('<U32'))
+        if allzfit['subtype'].dtype != '<U20':
+            allzfit.replace_column('subtype', allzfit['subtype'].astype('<U20'))
+            #TODO: Think about standardizing the array dtype in case nearest neighbours is used
 
         if allzfit['spectype'].dtype != '<U6':
             allzfit.replace_column('spectype',allzfit['spectype'].astype('<U6'))
 
         if not per_camera:
             maxcoeff = np.max([t.template.nbasis for t in templates])
         else:
-            ncam = 3 # three cameras of DESI:: b, r, z
             if n_nearest is not None:
-                maxcoeff = max(np.max([t.template.nbasis for t in templates]), ncam*(deg_legendre)+n_nearest)
+                maxcoeff = max(np.max([t.template.nbasis for t in templates]), ncamera*(deg_legendre)+n_nearest)
             else:
-                maxcoeff = max(np.max([t.template.nbasis for t in templates]), ncam*(deg_legendre)+1)
+                maxcoeff = max(np.max([t.template.nbasis for t in templates]), ncamera*(deg_legendre)+1)
         ntarg, ncoeff = allzfit['coeff'].shape
         if ncoeff != maxcoeff:
             coeff = np.zeros((ntarg, maxcoeff), dtype=allzfit['coeff'].dtype)