ObservationSim.py 16.5 KB
Newer Older
1
import os
Fang Yuedong's avatar
Fang Yuedong committed
2
3
4
5
6
import numpy as np
import mpi4py.MPI as MPI
import galsim
import logging
import psutil
7
8
9
10
11
12
from astropy.io import fits
from datetime import datetime

from ObservationSim.Config import config_dir, ChipOutput
from ObservationSim.Config.Header import generatePrimaryHeader, generateExtensionHeader
from ObservationSim.Instrument import Telescope, Filter, FilterParam, FocalPlane, Chip
13
from ObservationSim.Instrument.Chip import Effects
14
15
16
from ObservationSim.MockObject import calculateSkyMap_split_g
from ObservationSim.PSF import PSFGauss, FieldDistortion, PSFInterp
from ObservationSim._util import get_shear_field, makeSubDir_PointingList
Fang Yuedong's avatar
Fang Yuedong committed
17
18

class Observation(object):
19
    def __init__(self, config, Catalog, work_dir=None, data_dir=None):
20
        self.path_dict = config_dir(config=config, work_dir=work_dir, data_dir=data_dir)
Fang Yuedong's avatar
Fang Yuedong committed
21
        self.config = config
Fang Yuedong's avatar
Fang Yuedong committed
22
23
        # self.tel = Telescope(optEffCurve_path=self.path_dict["mirror_file"]) 
        self.tel = Telescope()
24
        self.focal_plane = FocalPlane(survey_type=self.config["obs_setting"]["survey_type"]) 
Fang Yuedong's avatar
Fang Yuedong committed
25
26
        # self.filter_param = FilterParam(filter_dir=self.path_dict["filter_dir"]) 
        self.filter_param = FilterParam() 
Fang Yuedong's avatar
Fang Yuedong committed
27
28
        self.chip_list = []
        self.filter_list = []
29
        self.Catalog = Catalog
Fang Yuedong's avatar
Fang Yuedong committed
30
31

        # if we want to apply field distortion?
Fang Yuedong's avatar
Fang Yuedong committed
32
        if self.config["ins_effects"]["field_dist"] == True:
Fang Yuedong's avatar
Fang Yuedong committed
33
            self.fd_model = FieldDistortion(fdModel_path=self.path_dict["fd_path"])
Fang Yuedong's avatar
Fang Yuedong committed
34
35
36
37
38
39
40
41
42
43
44
        else:
            self.fd_model = None

        # Construct chips & filters:
        nchips = self.focal_plane.nchip_x*self.focal_plane.nchip_y
        for i in range(nchips):
            chipID = i + 1
            if self.focal_plane.isIgnored(chipID=chipID):
                continue

            # Make Chip & Filter lists
45
46
            chip = Chip(
                chipID=chipID, 
Fang Yuedong's avatar
Fang Yuedong committed
47
48
49
50
51
                # ccdEffCurve_dir=self.path_dict["ccd_dir"], 
                # CRdata_dir=self.path_dict["CRdata_dir"], 
                # normalize_dir=self.path_dict["normalize_dir"], 
                # sls_dir=self.path_dict["sls_dir"], 
                config=self.config)
Fang Yuedong's avatar
Fang Yuedong committed
52
            filter_id, filter_type = chip.getChipFilter()
53
54
55
56
            filt = Filter(filter_id=filter_id, 
                filter_type=filter_type, 
                filter_param=self.filter_param, 
                ccd_bandpass=chip.effCurve)
Fang Yuedong's avatar
Fang Yuedong committed
57
58
59
60
            self.chip_list.append(chip)
            self.filter_list.append(filt)

        # Read catalog and shear(s)
61
        self.g1_field, self.g2_field, self.nshear = get_shear_field(config=self.config)
Fang Yuedong's avatar
Fang Yuedong committed
62

63
    def run_one_chip(self, chip, filt, pointing, chip_output, wcs_fp=None, psf_model=None, shear_cat_file=None, cat_dir=None, sed_dir=None):
Fang Yuedong's avatar
Fang Yuedong committed
64

65
66
67
68
69
70
71
72
73
        print(':::::::::::::::::::Current Pointing Information::::::::::::::::::')
        print("RA: %f, DEC; %f" % (pointing.ra, pointing.dec))
        print("Time: %s" % datetime.fromtimestamp(pointing.timestamp).isoformat())
        print("Exposure time: %f" % pointing.exp_time)
        print("Satellite Position (x, y, z): (%f, %f, %f)" % (pointing.sat_x, pointing.sat_y, pointing.sat_z))
        print("Satellite Velocity (x, y, z): (%f, %f, %f)" % (pointing.sat_vx, pointing.sat_vy, pointing.sat_vz))
        print("Position Angle: %f" % pointing.img_pa.deg)
        print('Chip : %d' % chip.chipID)
        print(':::::::::::::::::::::::::::END:::::::::::::::::::::::::::::::::::')
Fang Yuedong's avatar
Fang Yuedong committed
74

Fang Yuedong's avatar
Fang Yuedong committed
75
        if self.config["psf_setting"]["psf_model"] == "Gauss":
Fang Yuedong's avatar
Fang Yuedong committed
76
            psf_model = PSFGauss(chip=chip)
Fang Yuedong's avatar
Fang Yuedong committed
77
        elif self.config["psf_setting"]["psf_model"] == "Interp":
Fang Yuedong's avatar
Fang Yuedong committed
78
            psf_model = PSFInterp(chip=chip, PSF_data_file=self.path_dict["psf_dir"])
Fang Yuedong's avatar
Fang Yuedong committed
79
80
81
82
83
        else:
            print("unrecognized PSF model type!!", flush=True)

        # Get (extra) shear fields
        if shear_cat_file is not None:
84
            self.g1_field, self.g2_field, self.nshear = get_shear_field(config=self.config, shear_cat_file=shear_cat_file)
Fang Yuedong's avatar
Fang Yuedong committed
85
86
87

        # Get WCS for the focal plane
        if wcs_fp == None:
88
            wcs_fp = self.focal_plane.getTanWCS(pointing.ra, pointing.dec, pointing.img_pa, chip.pix_scale)
Fang Yuedong's avatar
Fang Yuedong committed
89
90
91
92
93
94
95

        # Create chip Image
        chip.img = galsim.ImageF(chip.npix_x, chip.npix_y)
        chip.img.setOrigin(chip.bound.xmin, chip.bound.ymin)
        chip.img.wcs = wcs_fp
        if chip.survey_type == "photometric":
            sky_map = None
96
97
        # elif chip.survey_type == "spectroscopic":
        #     sky_map = calculateSkyMap_split_g(xLen=chip.npix_x, yLen=chip.npix_y, blueLimit=filt.blue_limit, redLimit=filt.red_limit, skyfn=self.path_dict["sky_file"], conf=chip.sls_conf, pixelSize=chip.pix_scale, isAlongY=0)
Fang Yuedong's avatar
Fang Yuedong committed
98
        elif chip.survey_type == "spectroscopic":
99
100
101
102
103
104
105
106
107
108
109
110
111
112
            flat_normal = np.ones_like(chip.img.array)
            if self.config["ins_effects"]["flat_fielding"] == True:
                print("SLS flat preprocess,CHIP %d : Creating and applying Flat-Fielding"%chip.chipID, flush=True)
                print(chip.img.bounds, flush=True)
                flat_img = Effects.MakeFlatSmooth(
                    chip.img.bounds,
                    int(self.config["random_seeds"]["seed_flat"]))
                flat_normal = flat_normal * flat_img.array / np.mean(flat_img.array)
            if self.config["ins_effects"]["shutter_effect"] == True:
                print("SLS flat preprocess,CHIP %d : Apply shutter effect"%chip.chipID, flush=True)
                shuttimg = Effects.ShutterEffectArr(chip.img, t_shutter=1.3, dist_bearing=735,
                                                    dt=1E-3)  # shutter effect normalized image for this chip
                flat_normal = flat_normal*shuttimg
                flat_normal = np.array(flat_normal,dtype='float32')
Fang Yuedong's avatar
Fang Yuedong committed
113
114
115
116
117
118
119
120
            sky_map = calculateSkyMap_split_g(
                skyMap=flat_normal, 
                blueLimit=filt.blue_limit, 
                redLimit=filt.red_limit, 
                # skyfn=self.path_dict["sky_file"], 
                conf=chip.sls_conf, 
                pixelSize=chip.pix_scale, 
                isAlongY=0)
121
            del flat_normal
Fang Yuedong's avatar
Fang Yuedong committed
122

123
        if pointing.pointing_type == 'MS':
Fang Yuedong's avatar
Fang Yuedong committed
124
            # Load catalogues and templates
125
            self.cat = self.Catalog(config=self.config, chip=chip, pointing=pointing, cat_dir=cat_dir, sed_dir=sed_dir)
Fang Yuedong's avatar
Fang Yuedong committed
126
127
128
129
130
131
132
            self.nobj = len(self.cat.objs)

            # Loop over objects
            missed_obj = 0
            bright_obj = 0
            dim_obj = 0
            for j in range(self.nobj):
133
                # if j >= 100:
Fang Yuedong's avatar
Fang Yuedong committed
134
135
                #     break
                obj = self.cat.objs[j]
Fang Yuedong's avatar
Fang Yuedong committed
136
                if obj.type == 'star' and self.config["run_option"]["galaxy_only"]:
137
                    continue
Fang Yuedong's avatar
Fang Yuedong committed
138
                elif obj.type == 'galaxy' and self.config["run_option"]["star_only"]:
139
                    continue
Fang Yuedong's avatar
Fang Yuedong committed
140
                elif obj.type == 'quasar' and self.config["run_option"]["star_only"]:
141
                    continue
Fang Yuedong's avatar
Fang Yuedong committed
142

143
144
145
146
147
148
149
150
151
152
153
154
155
                # load SED
                try:
                    sed_data = self.cat.load_sed(obj)
                    norm_filt = self.cat.load_norm_filt(obj)
                    obj.sed, obj.param["mag_%s"%filt.filter_type] = self.cat.convert_sed(
                        mag=obj.param["mag_use_normal"],
                        sed=sed_data,
                        target_filt=filt, 
                        norm_filt=norm_filt,
                    )
                except Exception as e:
                    print(e)
                    continue
Fang Yuedong's avatar
Fang Yuedong committed
156
157
158
159
160
161
162
163
164
165
166

                # Exclude very bright/dim objects (for now)
                if filt.is_too_bright(mag=obj.getMagFilter(filt)):
                    # print("obj too birght!!", flush=True)
                    if obj.type != 'galaxy':
                        bright_obj += 1
                        obj.unload_SED()
                        continue
                if filt.is_too_dim(mag=obj.getMagFilter(filt)):
                    # print("obj too dim!!", flush=True)
                    dim_obj += 1
Fang Yuedong's avatar
Fang Yuedong committed
167
                    obj.unload_SED()
Fang Yuedong's avatar
Fang Yuedong committed
168
                    # print(obj.getMagFilter(filt))
Fang Yuedong's avatar
Fang Yuedong committed
169
170
                    continue

Fang Yuedong's avatar
Fang Yuedong committed
171
                if self.config["shear_setting"]["shear_type"] == "constant":
Fang Yuedong's avatar
Fang Yuedong committed
172
                    if obj.type == 'star':
173
                        obj.param["g1"], obj.param["g2"] = 0, 0
Fang Yuedong's avatar
Fang Yuedong committed
174
                    else:
175
                        obj.param["g1"], obj.param["g2"] = self.g1_field, self.g2_field
Fang Yuedong's avatar
Fang Yuedong committed
176
177
178
                elif self.config["shear_setting"]["shear_type"] == "extra":
                    try:
                        # TODO: every object with individual shear from input catalog(s)
179
                        obj.param["g1"], obj.param["g2"] = self.g1_field[j], self.g2_field[j]
Fang Yuedong's avatar
Fang Yuedong committed
180
181
182
                    except:
                        print("failed to load external shear.")
                        pass
183
184
185
186
                elif self.config["shear_setting"]["shear_type"] == "catalog":
                    pass
                else:
                    raise ValueError("Unknown shear input")
Fang Yuedong's avatar
Fang Yuedong committed
187
188
189
190
191
192
193
194

                pos_img, offset, local_wcs = obj.getPosImg_Offset_WCS(img=chip.img, fdmodel=self.fd_model, chip=chip, verbose=False)
                if pos_img.x == -1 or pos_img.y == -1:
                    # Exclude object which is outside the chip area (after field distortion)
                    # print("obj missed!!")
                    missed_obj += 1
                    obj.unload_SED()
                    continue
Fang Yuedong's avatar
Fang Yuedong committed
195

Fang Yuedong's avatar
Fang Yuedong committed
196
197
                # Draw object & update output catalog
                try:
Fang Yuedong's avatar
Fang Yuedong committed
198
                    if self.config["run_option"]["out_cat_only"]:
Fang Yuedong's avatar
Fang Yuedong committed
199
                        isUpdated = True
Fang Yuedong's avatar
Fang Yuedong committed
200
                    if chip.survey_type == "photometric" and not self.config["run_option"]["out_cat_only"]:
Fang Yuedong's avatar
Fang Yuedong committed
201
202
203
204
205
206
207
                        isUpdated, pos_shear = obj.drawObj_multiband(
                            tel=self.tel,
                            pos_img=pos_img, 
                            psf_model=psf_model, 
                            bandpass_list=filt.bandpass_sub_list, 
                            filt=filt, 
                            chip=chip, 
208
209
210
211
                            g1=obj.param["g1"], 
                            g2=obj.param["g2"], 
                            exptime=pointing.exp_time
                            )
Fang Yuedong's avatar
Fang Yuedong committed
212
                    elif chip.survey_type == "spectroscopic" and not self.config["run_option"]["out_cat_only"]:
Fang Yuedong's avatar
Fang Yuedong committed
213
214
215
216
217
218
219
                        isUpdated, pos_shear = obj.drawObj_slitless(
                            tel=self.tel, 
                            pos_img=pos_img, 
                            psf_model=psf_model, 
                            bandpass_list=filt.bandpass_sub_list, 
                            filt=filt, 
                            chip=chip, 
220
221
222
                            g1=obj.param["g1"], 
                            g2=obj.param["g2"], 
                            exptime=pointing.exp_time,
223
224
                            normFilter=norm_filt,
                            )
Fang Yuedong's avatar
Fang Yuedong committed
225
226
                    if isUpdated:
                        # TODO: add up stats
227
                        chip_output.cat_add_obj(obj, pos_img, pos_shear, obj.param["g1"], obj.param["g2"])
Fang Yuedong's avatar
Fang Yuedong committed
228
229
230
231
232
233
                        pass
                    else:
                        # print("object omitted", flush=True)
                        continue
                except Exception as e:
                    print(e)
Fang Yuedong's avatar
Fang Yuedong committed
234
                    pass
Fang Yuedong's avatar
Fang Yuedong committed
235
236
237
                # Unload SED:
                obj.unload_SED()
                del obj
Fang Yuedong's avatar
Fang Yuedong committed
238

Fang Yuedong's avatar
Fang Yuedong committed
239
240
            del psf_model
            del self.cat
Fang Yuedong's avatar
Fang Yuedong committed
241

242
        print("check running:1: pointing-{:} chip-{:} pid-{:} memory-{:6.2}GB".format(pointing.id, chip.chipID, os.getpid(), (psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024 / 1024) ), flush=True)
Fang Yuedong's avatar
Fang Yuedong committed
243
244
245

        # Detector Effects
        # ===========================================================
Fang Yuedong's avatar
Fang Yuedong committed
246
247
248
249
250
251
        # whether to output zero, dark, flat calibration images.
        chip.img = chip.addEffects(
            config=self.config, 
            img=chip.img, 
            chip_output=chip_output, 
            filt=filt, 
252
253
254
255
256
257
            ra_cen=pointing.ra, 
            dec_cen=pointing.dec,
            img_rot=pointing.img_pa,
            pointing_ID=pointing.id,
            timestamp_obs=pointing.timestamp,
            pointing_type=pointing.pointing_type,
Zhang Xin's avatar
Zhang Xin committed
258
            sky_map=sky_map, tel = self.tel)
Fang Yuedong's avatar
Fang Yuedong committed
259
        
260
261
        if pointing.pointing_type == 'MS':
            datetime_obs = datetime.fromtimestamp(pointing.timestamp)
Fang Yuedong's avatar
Fang Yuedong committed
262
263
264
265
266
            date_obs = datetime_obs.strftime("%y%m%d")
            time_obs = datetime_obs.strftime("%H%M%S")
            h_prim = generatePrimaryHeader(
                xlen=chip.npix_x, 
                ylen=chip.npix_y, 
267
268
269
                pointNum = str(pointing.id),
                ra=pointing.ra, 
                dec=pointing.dec, 
Fang Yuedong's avatar
Fang Yuedong committed
270
271
272
273
274
                psize=chip.pix_scale, 
                row_num=chip.rowID, 
                col_num=chip.colID,
                date=date_obs,
                time_obs=time_obs,
Fang Yuedong's avatar
Fang Yuedong committed
275
                exptime=pointing.exp_time,
Xin Zhang's avatar
Xin Zhang committed
276
277
278
                im_type='MS',
                sat_pos=[pointing.sat_x, pointing.sat_y, pointing.sat_z],
                sat_vel=[pointing.sat_vx, pointing.sat_vy, pointing.sat_vz])
Fang Yuedong's avatar
Fang Yuedong committed
279
280
281
            h_ext = generateExtensionHeader(
                xlen=chip.npix_x, 
                ylen=chip.npix_y, 
282
283
284
                ra=pointing.ra, 
                dec=pointing.dec, 
                pa=pointing.img_pa.deg, 
Fang Yuedong's avatar
Fang Yuedong committed
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
                gain=chip.gain, 
                readout=chip.read_noise, 
                dark=chip.dark_noise, 
                saturation=90000, 
                psize=chip.pix_scale, 
                row_num=chip.rowID, 
                col_num=chip.colID,
                extName='raw')
            chip.img = galsim.Image(chip.img.array, dtype=np.uint16)
            hdu1 = fits.PrimaryHDU(header=h_prim)
            hdu2 = fits.ImageHDU(chip.img.array, header=h_ext)
            hdu1 = fits.HDUList([hdu1, hdu2])
            fname = os.path.join(chip_output.subdir, h_prim['FILENAME'] + '.fits')
            hdu1.writeto(fname, output_verify='ignore', overwrite=True)
            print("# objects that are too bright %d out of %d"%(bright_obj, self.nobj))
            print("# objects that are too dim %d out of %d"%(dim_obj, self.nobj))
            print("# objects that are missed %d out of %d"%(missed_obj, self.nobj))
Fang Yuedong's avatar
Fang Yuedong committed
302
303
        del chip.img

304
        print("check running:2: pointing-{:} chip-{:} pid-{:} memory-{:6.2}GB".format(pointing.id, chip.chipID, os.getpid(), (psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024 / 1024) ), flush=True)
Fang Yuedong's avatar
Fang Yuedong committed
305

306
    def runExposure_MPI_PointingList(self, pointing_list, shear_cat_file=None, chips=None, use_mpi=False):
Fang Yuedong's avatar
Fang Yuedong committed
307
308
309
310
        if use_mpi:
            comm = MPI.COMM_WORLD
            ind_thread = comm.Get_rank()
            num_thread = comm.Get_size()
Fang Yuedong's avatar
Fang Yuedong committed
311

Fang Yuedong's avatar
Fang Yuedong committed
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
        if chips is None:
            nchips_per_fp = len(self.chip_list)
            run_chips = self.chip_list
            run_filts = self.filter_list
        else:
            # Only run a particular set of chips
            run_chips = []
            run_filts = []
            nchips_per_fp = len(chips)
            for ichip in range(len(self.chip_list)):
                chip = self.chip_list[ichip]
                filt = self.filter_list[ichip]
                if chip.chipID in chips:
                    run_chips.append(chip)
                    run_filts.append(filt)
Fang Yuedong's avatar
Fang Yuedong committed
327

328
        for ipoint in range(len(pointing_list)):
Fang Yuedong's avatar
Fang Yuedong committed
329
330
            for ichip in range(nchips_per_fp):
                i = ipoint*nchips_per_fp + ichip
331
332
                pointing = pointing_list[ipoint]
                pointing_ID = pointing.id
Fang Yuedong's avatar
Fang Yuedong committed
333
334
335
                if use_mpi:
                    if i % num_thread != ind_thread:
                        continue
Fang Yuedong's avatar
Fang Yuedong committed
336
337
338
339
340

                pid = os.getpid()

                sub_img_dir, prefix = makeSubDir_PointingList(path_dict=self.path_dict, config=self.config, pointing_ID=pointing_ID)

Fang Yuedong's avatar
Fang Yuedong committed
341
342
                chip = run_chips[ichip]
                filt = run_filts[ichip]
Fang Yuedong's avatar
Fang Yuedong committed
343
344
345
346
347
348
                print("running pointing#%d, chip#%d, at PID#%d..."%(pointing_ID, chip.chipID, pid), flush=True)
                chip_output = ChipOutput(
                    config=self.config, 
                    focal_plane=self.focal_plane, 
                    chip=chip, 
                    filt=filt,  
349
350
                    exptime=pointing.exp_time,
                    pointing_type=pointing.pointing_type,
Fang Yuedong's avatar
Fang Yuedong committed
351
352
353
                    pointing_ID=pointing_ID,  
                    subdir=sub_img_dir,
                    prefix=prefix)
354
                self.run_one_chip(
Fang Yuedong's avatar
Fang Yuedong committed
355
356
357
                    chip=chip, 
                    filt=filt, 
                    chip_output=chip_output, 
358
                    pointing=pointing,
Fang Yuedong's avatar
Fang Yuedong committed
359
                    cat_dir=self.path_dict["cat_dir"])
Zhang Xin's avatar
Zhang Xin committed
360
                print("finished running chip#%d..."%(chip.chipID), flush=True)