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.