-
Notifications
You must be signed in to change notification settings - Fork 0
/
apply_survey_geometry.py
288 lines (217 loc) · 10 KB
/
apply_survey_geometry.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
# Generate lightcone for DESI mocks
# from HOD catalogues generated by Shadab
import sys
import glob
import configparser
from itertools import product
import multiprocessing as mp
import numpy as np
from astropy.table import Table, vstack
import desimodel.footprint as foot
import desimodel.io
import h5py
def bits(ask="try"):
"Used"
if ask == "LC": return 0 #(0 0 0 0 0 0)
if ask == "downsample": return 1 #(0 0 0 0 0 1)
if ask == "Y5foot": return 2 #(0 0 0 0 1 0)
if ask == "downsample_LOP": return 4 #(0 0 0 1 0 0)
if ask == "Y1foot": return 8 #(0 0 1 0 0 0)
if ask == "Y1footbright": return 16 #(0 1 0 0 0 0)
print(f"You have asked for {ask}. Which does not exist. Please check bits() function in apply_survey_geometry.py.")
os._exit(1)
def mask(nz=0, Y5=0, nz_lop=0, Y1=0, Y1BRIGHT=0):
return nz * (2**0) + Y5 * (2**1) + nz_lop * (2**2) + Y1 * (2**3) + Y1BRIGHT * (2**4)
def apply_footprint(ra, dec, footprint_mask):
""" apply desi ootprint """
bitval = 0
# footprint_mask possibilities
# 0 - Y5 DESI; 1, 2, 3 - Y1 DESI
if footprint_mask == 0:
tiles = Table.read('/global/cfs/cdirs/desi/survey/ops/surveyops/trunk/ops/tiles-main.ecsv')
mask_y5 = (tiles['PROGRAM'] != 'BACKUP')
tiles = tiles[mask_y5]
bitval = bits(ask="Y5foot")
elif footprint_mask == 1:
tiles = Table.read('/global/cfs/cdirs/desi/survey/catalogs/Y1/LSS/tiles-DARK.fits')
bitval = bits(ask="Y1foot")
elif footprint_mask == 2:
tiles = Table.read('/global/cfs/cdirs/desi/survey/catalogs/Y1/LSS/tiles-BRIGHT.fits')
bitval = bits(ask="Y1foot")
elif footprint_mask == 3:
tiles_dark = Table.read('/global/cfs/cdirs/desi/survey/catalogs/Y1/LSS/tiles-DARK.fits')
tiles_bright = Table.read('/global/cfs/cdirs/desi/survey/catalogs/Y1/LSS/tiles-BRIGHT.fits')
tiles = vstack([tiles_dark, tiles_bright])
bitval = bits(ask="Y1foot")
else:
print("ERROR: Wrong footprint.", flush=True)
os._exit(1)
point = foot.is_point_in_desi(tiles, ra, dec)
idx = np.where(point)
print("FOOTPRINT: Selected {} out of {} galaxies.".format(len(idx[0]), len(ra)), flush=True)
newbits = np.zeros(len(ra), dtype=np.int32)
newbits[idx] = bitval
return newbits
def return_north(ra, dec):
'''
given a table that already includes RA,DEC, add PHOTSYS column denoting whether
the data is in the DECaLS ('S') or BASS/MzLS ('N') photometric region
'''
from astropy.coordinates import SkyCoord
import astropy.units as u
c = SkyCoord(ra* u.deg, dec* u.deg,frame='icrs')
gc = c.transform_to('galactic')
sel_ngc = gc.b > 0
seln = dec > 32.375
sel = seln&sel_ngc
return sel
class SurveyGeometry():
def __init__(self, config_file, args, galtype=None):
config = configparser.ConfigParser()
config.read(config_file)
self.config = config
self.box_length = config.getint('sim', 'box_length')
self.zmin = config.getfloat('sim', 'zmin')
self.zmax = config.getfloat('sim', 'zmax')
self.galtype = galtype
self.tracer_id = 0
self.mock_random_ic = args.mock_random_ic
if self.mock_random_ic is None:
self.mock_random_ic = config.get('sim', 'mock_random_ic')
if self.mock_random_ic != "ic":
if galtype in ("LRG", "LRG_main"):
self.tracer_id = 0
elif galtype == "ELG":
self.tracer_id = 1
elif galtype == "QSO":
self.tracer_id = 2
print(f"INFO: {self.galtype} with {self.tracer_id} ID")
def get_nz(self, z_cat, ask=None, ns=None):
''' The function where the n(z) is read
and the NZ column is computed for the given
redshifts.
'''
config = self.config
if ns is None:
z, nz = np.loadtxt(config["nz"][ask], usecols=(config.getint("nz", "col_z"), config.getint("nz", "col_nz")), unpack=True)
z_n = z
nz_n = (nz / ( 1 - config.getfloat('failurerate', ask) )) * 1.0
np.savetxt(config["nz"][ask + "_red_cor"], np.array([z_n, nz_n]).T)
return np.interp(z_cat, z_n, nz_n, left=0, right=0)
else:
ask = 'downsample_n'
z, nz = np.loadtxt(config["nz"][ask], usecols=(config.getint("nz", "col_z"), config.getint("nz", "col_nz")), unpack=True)
z_n = z
nz_n = (nz / ( 1 - config.getfloat('failurerate', 'downsample') )) * 1.0
np.savetxt(config["nz"][ask + "_red_cor"], np.array([z_n, nz_n]).T)
north_interp = np.interp(z_cat, z_n, nz_n, left=0, right=0)
ask = 'downsample_s'
z, nz = np.loadtxt(config["nz"][ask], usecols=(config.getint("nz", "col_z"), config.getint("nz", "col_nz")), unpack=True)
z_n = z
nz_n = (nz / ( 1 - config.getfloat('failurerate', 'downsample') )) * 1.0
np.savetxt(config["nz"][ask + "_red_cor"], np.array([z_n, nz_n]).T)
south_interp = np.interp(z_cat, z_n, nz_n, left=0, right=0)
return north_interp,south_interp
def downsample_aux_NS(self, z_cat, ran, n_mean, radec, ask=None):
""" downsample galaxies following n(z) model specified in galtype"""
nz_N, nz_S = self.get_nz(z_cat, ask=ask, ns='yes')
ra = radec[0]
dec = radec[1]
is_north = return_north(ra, dec)
is_south = ~return_north(ra, dec)
# downsample
nz_selected_n = (ran < nz_N / n_mean) & is_north
nz_selected_s = (ran < nz_S / n_mean) & is_south
# n_N = nz_N / n_mean
# n_S = nz_S / n_mean
idx = np.where(nz_selected_n|nz_selected_s)
print("DOWNSAMPLE: Selected {} out of {} galaxies.".format(len(idx[0]), len(z_cat)), flush=True)
bitval = bits(ask=ask)
newbits = np.zeros(len(z_cat), dtype=np.int32)
newbits[idx] = bitval
return newbits, nz_N, nz_S
def downsample_aux(self, z_cat, ran, n_mean, ask=None):
""" downsample galaxies following n(z) model specified in galtype"""
nz = self.get_nz(z_cat, ask=ask)
# downsample
nz_selected = ran < nz / n_mean
n = nz / n_mean
idx = np.where(nz_selected)
print("DOWNSAMPLE: Selected {} out of {} galaxies.".format(len(idx[0]), len(z_cat)), flush=True)
bitval = bits(ask=ask)
newbits = np.zeros(len(z_cat), dtype=np.int32)
newbits[idx] = bitval
return newbits, nz
def downsample(self, z_cat, n_mean, radec=None):
""" downsample galaxies following n(z) model specified in galtype"""
ran_i = np.random.rand(len(z_cat))
outbits = []
if self.galtype == "LRG":
outbits , _ = self.downsample_aux(z_cat, ran_i, n_mean, ask="downsample")
ran = [ran_i]
elif self.galtype == "ELG":
newbits, nz_N, nz_S = self.downsample_aux_NS(z_cat, ran_i, n_mean, radec, ask="downsample")
ran_n = np.random.rand(len(z_cat))
ran_n[newbits == 0] = np.inf
newbits_LOP, _ = self.downsample_aux(z_cat, ran_n, 1 , ask="downsample_LOP")
outbits = np.bitwise_or(newbits, newbits_LOP)
ran = [ran_i, ran_n]
elif self.galtype == "QSO":
outbits, _ = self.downsample_aux(z_cat, ran_i, n_mean, ask="downsample")
ran = [ran_i]
else:
print("Wrong galaxy type.")
os._exit(1)
return outbits, ran
def generate_shell(self, args):
infile, footprint_mask, todo = args
print(f"INFO: Read {infile}")
f = h5py.File(infile, 'r+')
n_mean = f.attrs["NGAL"] / (f.attrs["BOX_LENGTH"]**3)
shellnum = f.attrs["SHELLNUM"]
cat_seed = f.attrs["CAT_SEED"]
unique_seed = self.tracer_id * 500500 + 250 * cat_seed + shellnum
print("INFO: UNIQUE SEED:", unique_seed, flush=True)
np.random.seed(unique_seed)
data = f['galaxy']
ra = data['RA'][()]
dec = data['DEC'][()]
z_cosmo = data['Z_RSD'][()]
##TEMPz_cosmo = data['Z_COSMO'][()]
foot_bit_0 = apply_footprint(ra, dec, 0)
foot_bit_1 = apply_footprint(ra, dec, 1)
if self.mock_random_ic != "ic":
down_bit, ran_arr = self.downsample(z_cosmo, n_mean, radec=[ra, dec])
out_arr = np.bitwise_or(np.bitwise_or(foot_bit_0, foot_bit_1), down_bit)
else:
foot_bit_2 = apply_footprint(ra, dec, 2)
out_arr = np.bitwise_or(np.bitwise_or(foot_bit_0, foot_bit_1), foot_bit_2)
out_arr = out_arr.astype(np.int32)
if "STATUS" in data.keys():
print("WARNING: STATUS EXISTS. New STATUS has not been written.")
else:
f.create_dataset('galaxy/STATUS', data=out_arr, dtype=np.int32)
if self.mock_random_ic != "ic":
if "RAN_NUM_0_1" in data.keys():
print("WARNING: RAN_NUM_0_1 EXISTS. New RAN_NUM_0_1 has not been written.")
else:
f.create_dataset('galaxy/RAN_NUM_0_1', data=ran_arr[0], dtype=np.float32)
if self.galtype == "ELG":
f.create_dataset('galaxy/RAN_NUM_0_1_LOP', data=ran_arr[1], dtype=np.float32)
f.close()
def shell(self, path_instance, nproc=5, footprint_mask=0, todo=1):
infiles = glob.glob(path_instance.shells_out_path + "/*.hdf5")
args = product(infiles, [footprint_mask], [todo])
if nproc > len(infiles):
nproc = len(infiles)
with mp.Pool(processes=nproc) as pool:
pool.map_async(self.generate_shell, args)
pool.close()
pool.join()
def shell_series(self, path_instance, footprint_mask=0, todo=1):
infiles = glob.glob(path_instance.shells_out_path + "/*.hdf5")
print(infiles)
for file_ in infiles:
args = [file_, footprint_mask, todo]
# args = [infiles[1], footprint_mask, todo]
self.generate_shell(args)