ObservationSim.py 16.1 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
22
23
24
        self.tel = Telescope(optEffCurve_path=self.path_dict["mirror_file"]) 
        self.focal_plane = FocalPlane(survey_type=self.config["obs_setting"]["survey_type"]) 
        self.filter_param = FilterParam(filter_dir=self.path_dict["filter_dir"]) 
Fang Yuedong's avatar
Fang Yuedong committed
25
26
        self.chip_list = []
        self.filter_list = []
27
        self.Catalog = Catalog
Fang Yuedong's avatar
Fang Yuedong committed
28
29

        # if we want to apply field distortion?
Fang Yuedong's avatar
Fang Yuedong committed
30
        if self.config["ins_effects"]["field_dist"] == True:
Fang Yuedong's avatar
Fang Yuedong committed
31
            self.fd_model = FieldDistortion(fdModel_path=self.path_dict["fd_path"])
Fang Yuedong's avatar
Fang Yuedong committed
32
33
34
35
36
37
38
39
40
41
42
        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
43
44
45
46
47
48
            chip = Chip(
                chipID=chipID, 
                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
49
            filter_id, filter_type = chip.getChipFilter()
50
51
52
53
            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
54
55
56
57
            self.chip_list.append(chip)
            self.filter_list.append(filt)

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

60
    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
61

62
63
64
65
66
67
68
69
70
        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
71

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

        # Get (extra) shear fields
        if shear_cat_file is not None:
81
            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
82
83
84

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

        # 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
93
94
        # 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
95
        elif chip.survey_type == "spectroscopic":
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
            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')
            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)
            del flat_normal
Fang Yuedong's avatar
Fang Yuedong committed
112

113
        if pointing.pointing_type == 'MS':
Fang Yuedong's avatar
Fang Yuedong committed
114
            # Load catalogues and templates
115
            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
116
117
118
119
120
121
122
            self.nobj = len(self.cat.objs)

            # Loop over objects
            missed_obj = 0
            bright_obj = 0
            dim_obj = 0
            for j in range(self.nobj):
123
                # if j >= 100:
Fang Yuedong's avatar
Fang Yuedong committed
124
125
                #     break
                obj = self.cat.objs[j]
Fang Yuedong's avatar
Fang Yuedong committed
126
                if obj.type == 'star' and self.config["run_option"]["galaxy_only"]:
127
                    continue
Fang Yuedong's avatar
Fang Yuedong committed
128
                elif obj.type == 'galaxy' and self.config["run_option"]["star_only"]:
129
                    continue
Fang Yuedong's avatar
Fang Yuedong committed
130
                elif obj.type == 'quasar' and self.config["run_option"]["star_only"]:
131
                    continue
Fang Yuedong's avatar
Fang Yuedong committed
132

133
134
135
136
137
138
139
140
141
142
143
144
145
                # 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
146
147
148
149
150
151
152
153
154
155
156

                # 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
157
                    obj.unload_SED()
Fang Yuedong's avatar
Fang Yuedong committed
158
                    # print(obj.getMagFilter(filt))
Fang Yuedong's avatar
Fang Yuedong committed
159
160
                    continue

Fang Yuedong's avatar
Fang Yuedong committed
161
                if self.config["shear_setting"]["shear_type"] == "constant":
Fang Yuedong's avatar
Fang Yuedong committed
162
                    if obj.type == 'star':
163
                        obj.param["g1"], obj.param["g2"] = 0, 0
Fang Yuedong's avatar
Fang Yuedong committed
164
                    else:
165
                        obj.param["g1"], obj.param["g2"] = self.g1_field, self.g2_field
Fang Yuedong's avatar
Fang Yuedong committed
166
167
168
                elif self.config["shear_setting"]["shear_type"] == "extra":
                    try:
                        # TODO: every object with individual shear from input catalog(s)
169
                        obj.param["g1"], obj.param["g2"] = self.g1_field[j], self.g2_field[j]
Fang Yuedong's avatar
Fang Yuedong committed
170
171
172
                    except:
                        print("failed to load external shear.")
                        pass
173
174
175
176
                elif self.config["shear_setting"]["shear_type"] == "catalog":
                    pass
                else:
                    raise ValueError("Unknown shear input")
Fang Yuedong's avatar
Fang Yuedong committed
177
178
179
180
181
182
183
184

                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
185

Fang Yuedong's avatar
Fang Yuedong committed
186
187
                # Draw object & update output catalog
                try:
Fang Yuedong's avatar
Fang Yuedong committed
188
                    if self.config["run_option"]["out_cat_only"]:
Fang Yuedong's avatar
Fang Yuedong committed
189
                        isUpdated = True
Fang Yuedong's avatar
Fang Yuedong committed
190
                    if chip.survey_type == "photometric" and not self.config["run_option"]["out_cat_only"]:
Fang Yuedong's avatar
Fang Yuedong committed
191
192
193
194
195
196
197
                        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, 
198
199
200
201
                            g1=obj.param["g1"], 
                            g2=obj.param["g2"], 
                            exptime=pointing.exp_time
                            )
Fang Yuedong's avatar
Fang Yuedong committed
202
                    elif chip.survey_type == "spectroscopic" and not self.config["run_option"]["out_cat_only"]:
Fang Yuedong's avatar
Fang Yuedong committed
203
204
205
206
207
208
209
                        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, 
210
211
212
                            g1=obj.param["g1"], 
                            g2=obj.param["g2"], 
                            exptime=pointing.exp_time,
213
214
                            normFilter=norm_filt,
                            )
Fang Yuedong's avatar
Fang Yuedong committed
215
216
                    if isUpdated:
                        # TODO: add up stats
217
                        chip_output.cat_add_obj(obj, pos_img, pos_shear, obj.param["g1"], obj.param["g2"])
Fang Yuedong's avatar
Fang Yuedong committed
218
219
220
221
222
223
                        pass
                    else:
                        # print("object omitted", flush=True)
                        continue
                except Exception as e:
                    print(e)
Fang Yuedong's avatar
Fang Yuedong committed
224
                    pass
Fang Yuedong's avatar
Fang Yuedong committed
225
226
227
                # Unload SED:
                obj.unload_SED()
                del obj
Fang Yuedong's avatar
Fang Yuedong committed
228

Fang Yuedong's avatar
Fang Yuedong committed
229
230
            del psf_model
            del self.cat
Fang Yuedong's avatar
Fang Yuedong committed
231

232
        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
233
234
235

        # Detector Effects
        # ===========================================================
Fang Yuedong's avatar
Fang Yuedong committed
236
237
238
239
240
241
        # whether to output zero, dark, flat calibration images.
        chip.img = chip.addEffects(
            config=self.config, 
            img=chip.img, 
            chip_output=chip_output, 
            filt=filt, 
242
243
244
245
246
247
            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
248
            sky_map=sky_map, tel = self.tel)
Fang Yuedong's avatar
Fang Yuedong committed
249
        
250
251
        if pointing.pointing_type == 'MS':
            datetime_obs = datetime.fromtimestamp(pointing.timestamp)
Fang Yuedong's avatar
Fang Yuedong committed
252
253
254
255
256
            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, 
257
258
259
                pointNum = str(pointing.id),
                ra=pointing.ra, 
                dec=pointing.dec, 
Fang Yuedong's avatar
Fang Yuedong committed
260
261
262
263
264
265
266
267
268
                psize=chip.pix_scale, 
                row_num=chip.rowID, 
                col_num=chip.colID,
                date=date_obs,
                time_obs=time_obs,
                im_type='MS')
            h_ext = generateExtensionHeader(
                xlen=chip.npix_x, 
                ylen=chip.npix_y, 
269
270
271
                ra=pointing.ra, 
                dec=pointing.dec, 
                pa=pointing.img_pa.deg, 
Fang Yuedong's avatar
Fang Yuedong committed
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
                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
289
290
        del chip.img

291
        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
292

293
    def runExposure_MPI_PointingList(self, pointing_list, shear_cat_file=None, chips=None, use_mpi=False):
Fang Yuedong's avatar
Fang Yuedong committed
294
295
296
297
        if use_mpi:
            comm = MPI.COMM_WORLD
            ind_thread = comm.Get_rank()
            num_thread = comm.Get_size()
Fang Yuedong's avatar
Fang Yuedong committed
298

Fang Yuedong's avatar
Fang Yuedong committed
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
        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
314

315
        for ipoint in range(len(pointing_list)):
Fang Yuedong's avatar
Fang Yuedong committed
316
317
            for ichip in range(nchips_per_fp):
                i = ipoint*nchips_per_fp + ichip
318
319
                pointing = pointing_list[ipoint]
                pointing_ID = pointing.id
Fang Yuedong's avatar
Fang Yuedong committed
320
321
322
                if use_mpi:
                    if i % num_thread != ind_thread:
                        continue
Fang Yuedong's avatar
Fang Yuedong committed
323
324
325
326
327

                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
328
329
                chip = run_chips[ichip]
                filt = run_filts[ichip]
Fang Yuedong's avatar
Fang Yuedong committed
330
331
332
333
334
335
                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,  
336
337
                    exptime=pointing.exp_time,
                    pointing_type=pointing.pointing_type,
Fang Yuedong's avatar
Fang Yuedong committed
338
339
340
                    pointing_ID=pointing_ID,  
                    subdir=sub_img_dir,
                    prefix=prefix)
341
                self.run_one_chip(
Fang Yuedong's avatar
Fang Yuedong committed
342
343
344
                    chip=chip, 
                    filt=filt, 
                    chip_output=chip_output, 
345
                    pointing=pointing,
Fang Yuedong's avatar
Fang Yuedong committed
346
                    cat_dir=self.path_dict["cat_dir"])
Zhang Xin's avatar
Zhang Xin committed
347
                print("finished running chip#%d..."%(chip.chipID), flush=True)