From bc2be7a6e512c256031d35c2c27559ae2cf8c8a9 Mon Sep 17 00:00:00 2001 From: a_anand Date: Sat, 25 Mar 2023 18:05:16 -0700 Subject: [PATCH] added full_galaxy_dir as argument --- py/redrock/archetypes.py | 6 ++-- py/redrock/external/desi.py | 14 +++++++-- py/redrock/fitz.py | 7 +++-- py/redrock/nearest_neighbours.py | 51 ++++++++++++++++++++++++++++---- py/redrock/zfind.py | 9 +++--- 5 files changed, 70 insertions(+), 17 deletions(-) diff --git a/py/redrock/archetypes.py b/py/redrock/archetypes.py index 90c88874..11ec8db5 100644 --- a/py/redrock/archetypes.py +++ b/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()} diff --git a/py/redrock/external/desi.py b/py/redrock/external/desi.py index 5595f6f0..3d063eba 100644 --- a/py/redrock/external/desi.py +++ b/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) diff --git a/py/redrock/fitz.py b/py/redrock/fitz.py index f8b694e6..d33e04ea 100644 --- a/py/redrock/fitz.py +++ b/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, diff --git a/py/redrock/nearest_neighbours.py b/py/redrock/nearest_neighbours.py index 99fc5ef3..1f8a04c4 100644 --- a/py/redrock/nearest_neighbours.py +++ b/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 + + diff --git a/py/redrock/zfind.py b/py/redrock/zfind.py index 48333722..5831693c 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, 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.