From 930286ea72b24105bab6a52f7a9f0cc87e426f83 Mon Sep 17 00:00:00 2001
From: Abhijeet Anand <abhijeetanand2011@gmail.com>
Date: Tue, 29 Aug 2023 10:37:58 -0700
Subject: [PATCH] cleaned up scripts

---
 py/redrock/archetypes.py |  2 +-
 py/redrock/zfind.py      |  2 +-
 py/redrock/zscan.py      | 54 +++++++++++++++++++++++++++++-----------
 3 files changed, 42 insertions(+), 16 deletions(-)

diff --git a/py/redrock/archetypes.py b/py/redrock/archetypes.py
index 9b5ea4fe..99fce260 100644
--- a/py/redrock/archetypes.py
+++ b/py/redrock/archetypes.py
@@ -134,7 +134,7 @@ def nearest_neighbour_model(self, spectra,weights,flux,wflux,dwave,z, legendre,
             #print(z, zzchi2, zzcoeff, fsstype)
             return zzchi2, zzcoeff, self._rrtype+':::%s'%(fsstype)
 
-
+    
     def get_best_archetype(self,spectra,weights,flux,wflux,dwave,z,legendre, per_camera, n_nearest):
 
         """Get the best archetype for the given redshift and spectype.
diff --git a/py/redrock/zfind.py b/py/redrock/zfind.py
index 26dc90ab..c7a73feb 100644
--- a/py/redrock/zfind.py
+++ b/py/redrock/zfind.py
@@ -535,7 +535,7 @@ def zfind(targets, templates, mp_procs=1, nminima=3, archetypes=None, priors=Non
             allzfit.append(astropy.table.Table(tzfit))
 
         allzfit = astropy.table.vstack(allzfit)
-        #print(allzfit['targetid', 'z', 'zwarn', 'chi2', 'deltachi2', 'spectype', 'subtype'])
+        #print(allzfit['targetid', 'z', 'zwarn', 'chi2', 'deltachi2','spectype', 'subtype'])
         # Cosmetic: move TARGETID to be first column as primary key
         try:
             allzfit.columns.move_to_end('targetid', last=False)
diff --git a/py/redrock/zscan.py b/py/redrock/zscan.py
index 5933ca73..d1ed2150 100644
--- a/py/redrock/zscan.py
+++ b/py/redrock/zscan.py
@@ -237,6 +237,42 @@ def calc_zchi2_one(spectra, weights, flux, wflux, tdata):
 
     return zchi2, zcoeff
 
+
+def Tb_for_archetype(spectra, tdata, nbasis, n_nbh, nleg):
+    
+    """
+    Parameters
+    ---------------------
+
+     spectra (object): target spectra object
+     tdata (dict): template data for model fit
+     nleg (int): number of Legendre polynomials
+     nbasis (int): number of templates
+     n_nbh (int): number of nearest best archetypes
+     nleg (int): number of Legendre polynomials
+
+    Returns
+    --------------------
+    Matrix of dot product
+
+    """ 
+    
+    Tb = list()
+    
+    for i,s in enumerate(spectra):
+        key = s.wavehash
+        tdata2 = np.zeros((tdata[key].shape[0],nbasis))
+        for k in range(n_nbh):
+            tdata2[:,k]   = tdata[key][:,k] # these are nearest archetype
+        if nleg>0:
+            for k in range(n_nbh, n_nbh+nleg):
+                tdata2[:,k+nleg*i] = tdata[key][:,k] # Legendre polynomials terms
+        Tb.append(s.Rcsr.dot(tdata2))
+
+    Tb = np.vstack(Tb)
+
+    return Tb
+
 def per_camera_coeff_with_least_square(spectra, tdata, nleg, method=None, n_nbh=None):
 
   """
@@ -263,27 +299,15 @@ def per_camera_coeff_with_least_square(spectra, tdata, nleg, method=None, n_nbh=
   
   ncam = 3 # number of cameras in DESI: b, r, z
 
-  Tb = list()
   flux = np.concatenate([s.flux for s in spectra])
   weights = np.concatenate([s.ivar for s in spectra])
   wflux = flux*weights
 
   nbasis = n_nbh+nleg*ncam # n_nbh : for actual physical archetype(s), nleg: number of legendre polynomials, ncamera: number of cameras
 
-  # linear templates in each camera (only the Legendre terms will vary per camera, the lead archetype(s) will remain same for entire spectra)
+  #linear templates in each camera (only the Legendre terms will vary per camera, the lead archetype(s) will remain same for entire spectra)
 
-  for i,s in enumerate(spectra):
-    key = s.wavehash
-    tdata2 = np.zeros((tdata[key].shape[0],nbasis))
-    for k in range(n_nbh):
-        tdata2[:,k]   = tdata[key][:,k] # these are nearest archetype
-    if nleg>0:
-        for k in range(n_nbh, n_nbh+nleg):
-            tdata2[:,k+nleg*i] = tdata[key][:,k] # Legendre polynomials terms
-    Tb.append(s.Rcsr.dot(tdata2))
-
-  Tb = np.vstack(Tb)
-  
+  Tb = Tb_for_archetype(spectra, tdata, nbasis, n_nbh, nleg)  
   M = Tb.T.dot(np.multiply(weights[:,None], Tb))
   y = Tb.T.dot(wflux)
   ret_zcoeff= {'alpha':[], 'b':[], 'r':[], 'z':[]}
@@ -609,7 +633,9 @@ def calc_batch_dot_product_3d3d_gpu(a, b, transpose_a=False, fullprecision=True)
 ###    templates are looped over on the CPU.  The operations performed
 ###    are very obviously analagous though and should be highly
 ###    maintainable.  The main difference is the extra loop on the CPU version
+
 def calc_zchi2_batch(spectra, tdata, weights, flux, wflux, nz, nbasis, solve_matrices_algorithm="PCA", use_gpu=False, fullprecision=True):
+    
     """Calculate a batch of chi2.
     For many redshifts and a set of spectral data, compute the chi2 for
     template data that is already on the correct grid.