Merge cf392e9 into b722b18

py/desisim/scripts/quickquasars.py

@@ -5,6 +5,8 @@
 import time
 
 import numpy as np
+from scipy.constants import speed_of_light
+from scipy.stats import cauchy
 from astropy.table import Table,Column
 import astropy.io.fits as pyfits
 import multiprocessing
@@ -28,6 +30,7 @@
 from speclite import filters
 from desitarget.cuts import isQSO_colors
 
+c = speed_of_light/1000. #- km/s
 def parse(options=None):
     parser=argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                    description="""Fast simulation of QSO Lya spectra into the final DESI format (Spectra class) that can be directly used as
@@ -63,6 +66,7 @@ def parse(options=None):
 
     parser.add_argument('--sigma_kms_zfit',nargs='?',type=float,const=400,help="Adds a gaussian error to the quasar redshift, to simulate the redshift fitting step. E.g. --sigma_kms_zfit 200 will use a sigma value of 200 km/s. If a number is not specified the default is used.")
 
+    parser.add_argument('--shift_kms_los',nargs='?',type=float,const=100,help="Adds a shift to the quasar redshift as in arXiv:1708.02225. If a number is not specified the default (100km/s) is used.")
     parser.add_argument('--overwrite', action = "store_true" ,help="rerun if spectra exists (default is skip)")
     parser.add_argument('--target-selection', action = "store_true" ,help="apply QSO target selection cuts to the simulated quasars")
     parser.add_argument('--mags', action = "store_true", help="DEPRECATED; use --bbflux")
@@ -84,6 +88,15 @@ def parse(options=None):
     return args
 
 
+def mod_cauchy(loc,scale,size,cut):
+    samples=cauchy.rvs(loc=loc,scale=scale,size=3*size)
+    samples=samples[abs(samples)<cut]   
+    if len(samples)>=size:
+       samples=samples[:size] 
+    else:     
+        samples=mod_cauchy(loc,scale,size,cut)   ##Only added for the very unlikely case that there are not enough samples after the cut.
+    return samples
+
 def get_spectra_filename(args,nside,pixel):
     if args.outfile :
         return args.outfile
@@ -235,9 +248,17 @@ def simulate_one_healpix(ifilename,args,model,obsconditions,decam_and_wise_filte
     # might add metal transmission as well (from the HDU file).
     log.info("Read transmission file {}".format(ifilename))
     trans_wave, transmission, metadata, dla_info = read_lya_skewers(ifilename,read_dlas=(args.dla=='file'),add_metals=args.metals_from_file)
+
+###ADD dz_fog before generate the continua
+
+    Z_noFOG=np.copy(metadata['Z'])
+    log.info("Add FOG to redshift with sigma {} to quasar redshift".format(args.sigma_kms_fog)) 
+    dz_fog=(args.sigma_kms_fog/c)*(1.+metadata['Z'])*np.random.normal(0,1,len(metadata['Z']))    
+    metadata['Z']+=dz_fog
     ok = np.where(( metadata['Z'] >= args.zmin ) & (metadata['Z'] <= args.zmax ))[0]
     transmission = transmission[ok]
     metadata = metadata[:][ok]
+    Z_noFOG=Z_noFOG[ok]
 
     # option to make for BOSS+eBOSS
     if not eboss is None:
@@ -277,6 +298,7 @@ def simulate_one_healpix(ifilename,args,model,obsconditions,decam_and_wise_filte
             return
         transmission = transmission[selection]
         metadata = metadata[:][selection]
+        Z_noFOG=Z_noFOG[selection]
 
     nqso=transmission.shape[0]
     if args.downsampling is not None :
@@ -289,7 +311,9 @@ def simulate_one_healpix(ifilename,args,model,obsconditions,decam_and_wise_filte
             return
         transmission = transmission[indices]
         metadata = metadata[:][indices]
-        nqso = transmission.shape[0]
+        Z_noFOG=Z_noFOG[indices]
+        nqso = transmission.shape[0] 
+
 
     if args.nmax is not None :
         if args.nmax < nqso :
@@ -298,6 +322,7 @@ def simulate_one_healpix(ifilename,args,model,obsconditions,decam_and_wise_filte
             indices = (np.random.uniform(size=args.nmax)*nqso).astype(int)
             transmission = transmission[indices]
             metadata = metadata[:][indices]
+            Z_noFOG=Z_noFOG[indices]
             nqso = args.nmax
 
     # In previous versions of the London mocks we needed to enforce F=1 for
@@ -413,6 +438,7 @@ def simulate_one_healpix(ifilename,args,model,obsconditions,decam_and_wise_filte
         tmp_qso_flux = np.zeros([nqso, len(model.basewave)], dtype='f4')
         tmp_qso_wave = model.basewave
 
+
     for these, issouth in zip( (north, south), (False, True) ):
 
         # number of quasars in these
@@ -458,6 +484,7 @@ def simulate_one_healpix(ifilename,args,model,obsconditions,decam_and_wise_filte
         if args.balprob<=1. and args.balprob >0:
             log.info("Adding BALs with probability {}".format(args.balprob))
             # save current random state
+
             rnd_state = np.random.get_state()
             tmp_qso_flux,meta_bal=bal.insert_bals(tmp_qso_wave,tmp_qso_flux, metadata['Z'],
                                                   balprob=args.balprob,seed=seed)
@@ -532,6 +559,7 @@ def simulate_one_healpix(ifilename,args,model,obsconditions,decam_and_wise_filte
         metadata     = metadata[:][selection]
         meta         = meta[:][selection]
         qsometa      = qsometa[:][selection]
+        Z_noFOG      = Z_noFOG[selection]
 
         for band in bands :
             bbflux[band] = bbflux[band][selection]
@@ -577,17 +605,20 @@ def simulate_one_healpix(ifilename,args,model,obsconditions,decam_and_wise_filte
                 meta=specmeta,seed=seed,fibermap_columns=fibermap_columns,use_poisson=False) # use Poisson = False to get reproducible results.
 
 
-##Adedd to write the truth file, includen metadata for DLA's and BALs
-    log.info("Added FOG to redshift with sigma {} to zbest".format(args.sigma_kms_fog))
-    dz_fog=(args.sigma_kms_fog/299792.458)*(1.+metadata['Z'])*np.random.normal(0,1,nqso)
-    meta.rename_column('REDSHIFT','TRUEZ_noFOG')
+##Adedd to write the truth file, including metadata for DLA's and BALs      
+
     log.info('Writing a truth file  {}'.format(truth_filename))
+    meta.rename_column('REDSHIFT','TRUEZ')
+    meta.add_column(Column(Z_noFOG,name='TRUEZ_noFOG'))
+
     if 'Z_noRSD' in metadata.dtype.names:
         meta.add_column(Column(metadata['Z_noRSD'],name='TRUEZ_noRSD'))
     else:
         log.info('Z_noRSD field not present in transmission file. Z_noRSD not saved to truth file')
+    if args.shift_kms_los:
+       metadata['Z']+=(args.shift_kms_los/c*(1.+metadata['Z']))
+       log.info('Added a shift of {} km/s to the redshift'.format(str(args.shift_kms_los)))
 
-    meta.add_column(Column(metadata['Z']+dz_fog,name='TRUEZ'))
     hdu = pyfits.convenience.table_to_hdu(meta)
     hdu.header['EXTNAME'] = 'TRUTH'
     hduqso=pyfits.convenience.table_to_hdu(qsometa)
@@ -620,14 +651,14 @@ def simulate_one_healpix(ifilename,args,model,obsconditions,decam_and_wise_filte
             ('BRICKNAME', (str,8))]
         zbest = Table(np.zeros(nqso, dtype=columns))
         zbest["CHI2"][:]      = 0.
-        zbest["Z"]            = metadata['Z']+dz_fog
+        zbest["Z"]            = metadata['Z']
         zbest["ZERR"][:]      = 0.
 
         if args.sigma_kms_zfit:
            log.info("Added zfit error with sigma {} to zbest".format(args.sigma_kms_zfit))
-           sigma_zfit=(args.sigma_kms_zfit/299792.458)*(1.+metadata['Z'])
-           zbest["Z"]+=sigma_zfit*np.random.normal(0,1,nqso)
-           zbest["ZERR"]=sigma_zfit
+           dz_fit=mod_cauchy(loc=0,scale=args.sigma_kms_zfit,size=nqso,cut=3000)*(1.+metadata['Z'])/c
+           zbest["Z"]+=dz_fit
+        zbest["ZERR"][:]      = 0
         zbest["ZWARN"][:]     = 0
         zbest["SPECTYPE"][:]  = "QSO"
         zbest["SUBTYPE"][:]   = ""