Skip to content

Commit

Permalink
Merge branch 'master' of https://github.com/igmhub/SaclayMocks
Browse files Browse the repository at this point in the history
  • Loading branch information
Thomas committed Jul 1, 2019
2 parents bd58279 + 9f803f6 commit 2ff1d74
Showing 1 changed file with 55 additions and 58 deletions.
113 changes: 55 additions & 58 deletions bin/draw_qso.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -222,7 +222,7 @@ def main():
# we then build the probability ptot = c*p12 + (1-c)*p23
# where c is a coefficient, store in etc/qso_lognormal_coef.txt
if not random_cond:
print("Computing exp(a(z)*g) for the 2 lognormal fields and interpolating...")
print("Computing exp(a(z)*g) for the 3 lognormal fields and interpolating...")
t3 = time()
z1 = constant.z_QSO_bias_1
z2 = constant.z_QSO_bias_2
Expand Down Expand Up @@ -366,11 +366,19 @@ def main():
if (not random_cond):
nnQSO=0
qsofits = FITS(out_file, 'rw', clobber=True) # output file
z_list = []
z_rsd_list = []
ra_list = []
dec_list = []
xx_list = []
yy_list = []
zz_list = []
if not random_cond:
names = ["Z_QSO_NO_RSD", "Z_QSO_RSD", "RA", "DEC", "HDU", "THING_ID", "PLATE", "MJD", "FIBERID", "PMF", "XX", "YY", "ZZ"] # , "XGRID", "YGRID", "ZGRID"]
else:
names = ["Z", "RA", "DEC", "HDU", "THING_ID", "PLATE", "MJD", "FIBERID", "PMF", "XX", "YY", "ZZ"] # , "XGRID", "YGRID", "ZGRID"]

t4 = time()
for mz in range(NZ):
XX = x_axis
YY = y_axis
Expand Down Expand Up @@ -425,9 +433,9 @@ def main():
redshift = args.zfix*np.ones_like(RR)
if redshift.min() > z_max +1.: continue
if not random_cond and rsd:
vpar = (XX.reshape(-1,1)*vx[:, :, mz]
+ YY*vy[:, :, mz]
+ ZZ*vz[:, :, mz]) / RR
vpar = (XXX*vx[:, :, mz]
+ YYY*vy[:, :, mz]
+ ZZZ*vz[:, :, mz]) / RR
msk = np.where(redshift < z_max + 1.) # don't go above z=5
if args.zfix is None:
RR_RSD = RR.copy()
Expand Down Expand Up @@ -468,16 +476,6 @@ def main():
ra = ra[iqso]
dec = dec[iqso]

# # Draw random xyz in the cell
# X = x_axis[iqso[0]] + np.random.uniform(-DX/2, DX/2, len(iqso[0]))
# Y = y_axis[iqso[1]] + np.random.uniform(-DY/2, DY/2, len(iqso[0]))
# Z = ZZ + np.random.uniform(-DZ/2, DZ/2, len(iqso[0]))
# zzz = z_of_R(np.sqrt(X*X+Y*Y+Z*Z)/h)

# ra, dec, _ = box.ComputeRaDecR2(XX, YY, ZZ, np.radians(ra0), np.radians(dec0)) # (NQS0_)
# ra = np.degrees(ra)
# dec = np.degrees(dec)

# Select desi footprint
if desi:
msk = util.desi_footprint(ra, dec)
Expand All @@ -489,49 +487,52 @@ def main():
ZZ = ZZ[msk]
zzz_RSD = zzz_RSD[msk]

un = np.ones(len(zzz)).astype(int)
# XGRID = np.array(iqso[0]) + i_slice*NX
# YGRID = np.array(iqso[1])
# ZGRID = un * mz
#if (len(XGRID)>0):
# print ( np.mean(np.log(ptot[XGRID,YGRID,ZGRID])) )
THING_ID = (chunk*1e9 + i_slice*1e6 + np.arange(nQSO, nQSO+len(zzz)) + 1).astype(int) # start at 1
HDU = un * i_slice # QSOhdu
plate = THING_ID
mjd = np.random.randint(51608, high=57521, size=len(zzz))
fiberid = np.random.randint(1,high=1001, size=len(zzz))
# pmf = np.array([plate[i].astype(str)+'-'+mjd[i].astype(str)+'-'+fiberid[i].astype(str) for i in range(len(zzz))])
pmf = np.empty((len(zzz)), dtype='|S21')
for i in range(len(zzz)):
l_tmp = plate[i].astype(str)+'-'+mjd[i].astype(str)+'-'+fiberid[i].astype(str)
pmf[i] = l_tmp
nQSO += len(zzz)
# print(mz)
# print(len(zzz), len(tanx), len(X))
#print (len(zzz),len(ra),len(dec))
if not random_cond:
array_list = [np.float32(zzz), np.float32(zzz_RSD), np.float32(ra), np.float32(dec), np.int32(HDU), THING_ID, plate, np.int32(mjd), np.int32(fiberid), pmf, np.float32(XX), np.float32(YY), np.float32(ZZ)] # , XGRID, YGRID, ZGRID]
else:
array_list = [np.float32(zzz), np.float32(ra), np.float32(dec), np.int32(HDU), THING_ID, plate, np.int32(mjd), np.int32(fiberid), pmf, np.float32(XX), np.float32(YY), np.float32(ZZ)] # , XGRID, YGRID, ZGRID]

# print("MZ : {}".format(mz))
# print(array_list)
ra_list.append(ra)
dec_list.append(dec)
z_list.append(zzz)
z_rsd_list.append(zzz_RSD)
xx_list.append(XX)
yy_list.append(YY)
zz_list.append(ZZ)

# end of loop on qso
t5 = time()
print("End of loop on QSO. Took {} s".format(t5 - t4))
if len(ra_list) > 0:
ra_list = np.concatenate(ra_list)
dec_list = np.concatenate(dec_list)
z_list = np.concatenate(z_list)
z_rsd_list = np.concatenate(z_rsd_list)
xx_list = np.concatenate(xx_list)
yy_list = np.concatenate(yy_list)
zz_list = np.concatenate(zz_list)
# write to fits file
un = np.ones(nQSO)
thing_id = (chunk*1e9 + i_slice*1e6 + np.arange(nQSO) + 1).astype(int) # start at 1
hdu = un * i_slice # QSOhdu
plate = thing_id
mjd = np.random.randint(51608, high=57521, size=nQSO)
fiberid = np.random.randint(1,high=1001, size=nQSO)
pmf = np.empty(nQSO, dtype='|S21')
for i in range(nQSO):
l_tmp = plate[i].astype(str)+'-'+mjd[i].astype(str)+'-'+fiberid[i].astype(str)
pmf[i] = l_tmp
if not random_cond:
array_list = [np.float32(z_list), np.float32(z_rsd_list), np.float32(ra_list), np.float32(dec_list), np.int32(hdu), thing_id, plate, np.int32(mjd), np.int32(fiberid), pmf, np.float32(xx_list), np.float32(yy_list), np.float32(zz_list)]
else:
array_list = [np.float32(z_list), np.float32(ra_list), np.float32(dec_list), np.int32(hdu), thing_id, plate, np.int32(mjd), np.int32(fiberid), pmf, np.float32(xx_list), np.float32(yy_list), np.float32(zz_list)]

if (len(qsofits) == 1):
qsofits.write(array_list, names=names)
#exit(0)
else:
qsofits[-1].append(array_list, names=names)
# this is not making a new HDU, this is appending to last hdu, i.e. hdu=1
qsofits.write(array_list, names=names)

if len(qsofits) == 1:
print("No QSO drawn, creating a null table")
if not random_cond:
null_table = [np.array([], dtype=np.float32), np.array([], dtype=np.float32), np.array([], dtype=np.float32), np.array([], dtype=np.float32), np.array([], dtype=np.int32), np.array([], dtype=np.int64), np.array([], dtype=np.int64), np.array([], dtype=np.int32), np.array([], dtype=np.int32), np.array([], dtype='|S21'), np.array([], dtype=np.float32), np.array([], dtype=np.float32), np.array([], dtype=np.float32)]
else:
null_table = [np.array([], dtype=np.float32), np.array([], dtype=np.float32), np.array([], dtype=np.float32), np.array([], dtype=np.int32), np.array([], dtype=np.int64), np.array([], dtype=np.int64), np.array([], dtype=np.int32), np.array([], dtype=np.int32), np.array([], dtype='|S21'), np.array([], dtype=np.float32), np.array([], dtype=np.float32), np.array([], dtype=np.float32)]
# if len(qsofits) == 1:
# print("No QSO drawn, creating a null table")
# if not random_cond:
# null_table = [np.array([], dtype=np.float32), np.array([], dtype=np.float32), np.array([], dtype=np.float32), np.array([], dtype=np.float32), np.array([], dtype=np.int32), np.array([], dtype=np.int64), np.array([], dtype=np.int64), np.array([], dtype=np.int32), np.array([], dtype=np.int32), np.array([], dtype='|S21'), np.array([], dtype=np.float32), np.array([], dtype=np.float32), np.array([], dtype=np.float32)]
# else:
# null_table = [np.array([], dtype=np.float32), np.array([], dtype=np.float32), np.array([], dtype=np.float32), np.array([], dtype=np.int32), np.array([], dtype=np.int64), np.array([], dtype=np.int64), np.array([], dtype=np.int32), np.array([], dtype=np.int32), np.array([], dtype='|S21'), np.array([], dtype=np.float32), np.array([], dtype=np.float32), np.array([], dtype=np.float32)]

qsofits.write(null_table, names=names)
# qsofits.write(null_table, names=names)

qsofits[1].write_key("seed", np.int32(seed), comment="seed used to generate randoms")
qsofits[1].write_key("ra0", ra0, comment="right ascension of the box center")
Expand All @@ -548,15 +549,11 @@ def main():
qsofits[1].write_key("XX", None, comment="position on X axis in Mpc/h")
qsofits[1].write_key("YY", None, comment="position on Y axis in Mpc/h")
qsofits[1].write_key("ZZ", None, comment="position on Z axis in Mpc/h")
# qsofits[1].write_key("XGRID", None, comment="index on X axis")
# qsofits[1].write_key("YGRID", None, comment="index on Y axis")
# qsofits[1].write_key("ZGRID", None, comment="index on Z axis")

qsofits.close()
print("File {} written in {} s".format(out_file, time() - t5))
print nQSO, "QSOs drawn"
if (not random_cond):
print nnQSO, "QSOs in the full box" # prov
print out_file, "file written"
print("Took {}s".format(time()-t_init))

if drawPlot : plt.show()
Expand Down

0 comments on commit 2ff1d74

Please sign in to comment.