diff --git a/py/redrock/archetypes.py b/py/redrock/archetypes.py index e25c64ef..3150753b 100644 --- a/py/redrock/archetypes.py +++ b/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,9 +319,9 @@ 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) @@ -328,7 +329,7 @@ def get_best_archetype(self,target,weights,flux,wflux,dwave,z, per_camera, n_nea (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) diff --git a/py/redrock/external/desi.py b/py/redrock/external/desi.py index 9f5e1aa8..cb7bc919 100644 --- a/py/redrock/external/desi.py +++ b/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) diff --git a/py/redrock/fitz.py b/py/redrock/fitz.py index 6e0f66fd..4b6ba289 100644 --- a/py/redrock/fitz.py +++ b/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) diff --git a/py/redrock/zfind.py b/py/redrock/zfind.py index 09f3e9dc..04dc462f 100644 --- a/py/redrock/zfind.py +++ b/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,8 +588,9 @@ def zfind(targets, templates, mp_procs=1, nminima=3, archetypes=None, priors=Non del allresults[tid]['meta'] # Standardize column sizes - if allzfit['subtype'].dtype != '