ObservationSim.py 29 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
Fang Yuedong's avatar
Fang Yuedong committed
7
import gc
8
9
10
from astropy.io import fits
from datetime import datetime

Fang Yuedong's avatar
Fang Yuedong committed
11
12
import traceback

13
14
15
from ObservationSim.Config import config_dir, ChipOutput
from ObservationSim.Config.Header import generatePrimaryHeader, generateExtensionHeader
from ObservationSim.Instrument import Telescope, Filter, FilterParam, FocalPlane, Chip
16
from ObservationSim.Instrument.Chip import Effects
17
from ObservationSim.Straylight import calculateSkyMap_split_g
Zhang Xin's avatar
Zhang Xin committed
18
from ObservationSim.PSF import PSFGauss, FieldDistortion, PSFInterp, PSFInterpSLS
19
from ObservationSim._util import get_shear_field, makeSubDir_PointingList
20
from ObservationSim.Astrometry.Astrometry_util import on_orbit_obs_position
Zhang Xin's avatar
Zhang Xin committed
21
from ObservationSim.MockObject import FlatLED
Fang Yuedong's avatar
Fang Yuedong committed
22
23

class Observation(object):
24
    def __init__(self, config, Catalog, work_dir=None, data_dir=None):
25
        self.path_dict = config_dir(config=config, work_dir=work_dir, data_dir=data_dir)
Fang Yuedong's avatar
Fang Yuedong committed
26
        self.config = config
Fang Yuedong's avatar
Fang Yuedong committed
27
        self.tel = Telescope()
28
        self.focal_plane = FocalPlane(survey_type=self.config["obs_setting"]["survey_type"]) 
Fang Yuedong's avatar
Fang Yuedong committed
29
        self.filter_param = FilterParam() 
Fang Yuedong's avatar
Fang Yuedong committed
30
31
        self.chip_list = []
        self.filter_list = []
32
        self.all_filter = []
33
        self.Catalog = Catalog
Fang Yuedong's avatar
Fang Yuedong committed
34
35

        # Construct chips & filters:
Fang Yuedong's avatar
Fang Yuedong committed
36
        for i in range(self.focal_plane.nchips):
Fang Yuedong's avatar
Fang Yuedong committed
37
38
39
            chipID = i + 1

            # Make Chip & Filter lists
40
41
            chip = Chip(
                chipID=chipID, 
Fang Yuedong's avatar
Fang Yuedong committed
42
                config=self.config)
Fang Yuedong's avatar
Fang Yuedong committed
43
            filter_id, filter_type = chip.getChipFilter()
44
45
            filt = Filter(filter_id=filter_id, 
                filter_type=filter_type, 
46
                filter_param=self.filter_param)
47
48
49
50
            if not self.focal_plane.isIgnored(chipID=chipID):
                self.chip_list.append(chip)
                self.filter_list.append(filt)
            self.all_filter.append(filt)
Fang Yuedong's avatar
Fang Yuedong committed
51

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

Fang Yuedong's avatar
Fang Yuedong committed
54
55
56
57
58
59
60
61
62
        chip_output.Log_info(':::::::::::::::::::Current Pointing Information::::::::::::::::::')
        chip_output.Log_info("RA: %f, DEC; %f" % (pointing.ra, pointing.dec))
        chip_output.Log_info("Time: %s" % datetime.utcfromtimestamp(pointing.timestamp).isoformat())
        chip_output.Log_info("Exposure time: %f" % pointing.exp_time)
        chip_output.Log_info("Satellite Position (x, y, z): (%f, %f, %f)" % (pointing.sat_x, pointing.sat_y, pointing.sat_z))
        chip_output.Log_info("Satellite Velocity (x, y, z): (%f, %f, %f)" % (pointing.sat_vx, pointing.sat_vy, pointing.sat_vz))
        chip_output.Log_info("Position Angle: %f" % pointing.img_pa.deg)
        chip_output.Log_info('Chip : %d' % chip.chipID)
        chip_output.Log_info(':::::::::::::::::::::::::::END:::::::::::::::::::::::::::::::::::')
Fang Yuedong's avatar
Fang Yuedong committed
63

Fang Yuedong's avatar
Fang Yuedong committed
64
        if self.config["psf_setting"]["psf_model"] == "Gauss":
Fang Yuedong's avatar
Fang Yuedong committed
65
            psf_model = PSFGauss(chip=chip, psfRa=self.config["psf_setting"]["psf_rcont"])
Fang Yuedong's avatar
Fang Yuedong committed
66
        elif self.config["psf_setting"]["psf_model"] == "Interp":
Zhang Xin's avatar
Zhang Xin committed
67
68
69
70
            if chip.survey_type == "spectroscopic":
                psf_model = PSFInterpSLS(chip, filt,PSF_data_prefix=self.path_dict["psf_sls_dir"])
            else:
                psf_model = PSFInterp(chip=chip, npsf=chip.n_psf_samples, PSF_data_file=self.path_dict["psf_dir"])
Fang Yuedong's avatar
Fang Yuedong committed
71
        else:
Fang Yuedong's avatar
Fang Yuedong committed
72
            chip_output.Log_error("unrecognized PSF model type!!", flush=True)
Fang Yuedong's avatar
Fang Yuedong committed
73

Fang Yuedong's avatar
Fang Yuedong committed
74
        # Figure out shear fields
75
        self.g1_field, self.g2_field, self.nshear = get_shear_field(config=self.config)
Fang Yuedong's avatar
Fang Yuedong committed
76

77
78
        # Apply astrometric simulation for pointing
        if self.config["obs_setting"]["enable_astrometric_model"]:
Fang Yuedong's avatar
Fang Yuedong committed
79
            dt = datetime.utcfromtimestamp(pointing.timestamp)
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
            date_str = dt.date().isoformat()
            time_str = dt.time().isoformat()
            ra_cen, dec_cen = on_orbit_obs_position(
                input_ra_list=[pointing.ra],
                input_dec_list=[pointing.dec],
                input_pmra_list=[0.],
                input_pmdec_list=[0.],
                input_rv_list=[0.],
                input_parallax_list=[1e-9],
                input_nstars=1,
                input_x=pointing.sat_x,
                input_y=pointing.sat_y,
                input_z=pointing.sat_z,
                input_vx=pointing.sat_vx,
                input_vy=pointing.sat_vy,
                input_vz=pointing.sat_vz,
Fang Yuedong's avatar
Fang Yuedong committed
96
                input_epoch="J2000",
97
98
99
100
101
102
103
104
                input_date_str=date_str,
                input_time_str=time_str
            )
            ra_cen, dec_cen = ra_cen[0], dec_cen[0]
        else:
            ra_cen = pointing.ra
            dec_cen = pointing.dec

Fang Yuedong's avatar
Fang Yuedong committed
105
106
        # Get WCS for the focal plane
        if wcs_fp == None:
107
            wcs_fp = self.focal_plane.getTanWCS(ra_cen, dec_cen, pointing.img_pa, chip.pix_scale)
Fang Yuedong's avatar
Fang Yuedong committed
108
109
110
111
112

        # 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
113

Zhang Xin's avatar
Zhang Xin committed
114
115
        if self.config["obs_setting"]["enable_straylight_model"]:
            filt.setFilterStrayLightPixel(jtime = pointing.jdt, sat_pos = np.array([pointing.sat_x, pointing.sat_y, pointing.sat_z]), pointing_radec = np.array([pointing.ra,pointing.dec]), sun_pos = np.array([pointing.sun_x,pointing.sun_y,pointing.sun_z]))
116

117
118
        chip_output.Log_info("========================sky pix========================")
        chip_output.Log_info(filt.sky_background)
119

Fang Yuedong's avatar
Fang Yuedong committed
120
121
122
        if chip.survey_type == "photometric":
            sky_map = None
        elif chip.survey_type == "spectroscopic":
xin's avatar
xin committed
123
            # chip.loadSLSFLATCUBE(flat_fn='flat_cube.fits')
124
125
            flat_normal = np.ones_like(chip.img.array)
            if self.config["ins_effects"]["flat_fielding"] == True:
Fang Yuedong's avatar
Fang Yuedong committed
126
                chip_output.Log_info("SLS flat preprocess,CHIP %d : Creating and applying Flat-Fielding"%chip.chipID)
127
                msg = str(chip.img.bounds)
Fang Yuedong's avatar
Fang Yuedong committed
128
                chip_output.Log_info(msg)
129
130
131
132
133
                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:
Fang Yuedong's avatar
Fang Yuedong committed
134
                chip_output.Log_info("SLS flat preprocess,CHIP %d : Apply shutter effect"%chip.chipID)
135
136
137
138
                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
139
            sky_map = calculateSkyMap_split_g(
140
141
142
143
144
145
                skyMap=flat_normal,
                blueLimit=filt.blue_limit,
                redLimit=filt.red_limit,
                conf=chip.sls_conf,
                pixelSize=chip.pix_scale,
                isAlongY=0,
146
147
                flat_cube=chip.flat_cube, zoldial_spec = filt.zodical_spec)
            sky_map = sky_map+filt.sky_background
148
            del flat_normal
Fang Yuedong's avatar
Fang Yuedong committed
149

150
        if pointing.pointing_type == 'SCI':
Fang Yuedong's avatar
Fang Yuedong committed
151
            # Load catalogues and templates
Fang Yuedong's avatar
Fang Yuedong committed
152
            self.cat = self.Catalog(config=self.config, chip=chip, pointing=pointing, cat_dir=cat_dir, sed_dir=sed_dir, chip_output=chip_output, filt=filt)
153
            chip_output.create_output_file()
Fang Yuedong's avatar
Fang Yuedong committed
154
155
            self.nobj = len(self.cat.objs)

Fang Yuedong's avatar
Fang Yuedong committed
156
157
158
            for ifilt in range(len(self.all_filter)):
                temp_filter = self.all_filter[ifilt]
                # Update the limiting magnitude using exposure time in pointing
159
                temp_filter.update_limit_saturation_mags(exptime=pointing.get_full_depth_exptime(temp_filter.filter_type), chip=chip)
Fang Yuedong's avatar
Fang Yuedong committed
160
161
162
163
164

                # Select cutting band filter for saturation/limiting magnitude
                if temp_filter.filter_type.lower() == self.config["obs_setting"]["cut_in_band"].lower():
                    cut_filter = temp_filter

Fang Yuedong's avatar
Fang Yuedong committed
165
            if self.config["ins_effects"]["field_dist"] == True:
166
                self.fd_model = FieldDistortion(chip=chip, img_rot=pointing.img_pa.deg)
Fang Yuedong's avatar
Fang Yuedong committed
167
168
169
            else:
                self.fd_model = None

Fang Yuedong's avatar
Fang Yuedong committed
170
171
172
173
            # Loop over objects
            missed_obj = 0
            bright_obj = 0
            dim_obj = 0
Fang Yuedong's avatar
Fang Yuedong committed
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
 
            h_ext = generateExtensionHeader(
                chip=chip,
                xlen=chip.npix_x, 
                ylen=chip.npix_y, 
                ra=pointing.ra, 
                dec=pointing.dec, 
                pa=pointing.img_pa.deg, 
                gain=chip.gain, 
                readout=chip.read_noise, 
                dark=chip.dark_noise, 
                saturation=90000, 
                pixel_scale=chip.pix_scale, 
                pixel_size=chip.pix_size,
                xcen=chip.x_cen,
                ycen=chip.y_cen,
190
191
192
                extName='SCI',
                timestamp = pointing.timestamp,
                exptime = pointing.exp_time,
193
                readoutTime = chip.readout_time)
194
195
            
            chip_wcs = galsim.FitsWCS(header=h_ext)
Fang Yuedong's avatar
Fang Yuedong committed
196

Fang Yuedong's avatar
Fang Yuedong committed
197
            for j in range(self.nobj):
198
199
                
                # (DEBUG)
Zhang Xin's avatar
Zhang Xin committed
200
201
                if j >= 10:
                    break
202

Fang Yuedong's avatar
Fang Yuedong committed
203
204
                obj = self.cat.objs[j]

Zhang Xin's avatar
Zhang Xin committed
205
206
207
208
                # (DEBUG)
                # if obj.getMagFilter(filt)>20:
                #     continue

209
                # load and convert SED; also caculate object's magnitude in all CSST bands
210
211
212
                try:
                    sed_data = self.cat.load_sed(obj)
                    norm_filt = self.cat.load_norm_filt(obj)
Fang Yuedong's avatar
Fang Yuedong committed
213
                    obj.sed, obj.param["mag_%s"%filt.filter_type.lower()], obj.param["flux_%s"%filt.filter_type.lower()] = self.cat.convert_sed(
214
215
216
217
218
                        mag=obj.param["mag_use_normal"],
                        sed=sed_data,
                        target_filt=filt, 
                        norm_filt=norm_filt,
                    )
Fang Yuedong's avatar
Fang Yuedong committed
219
                    _, obj.param["mag_%s"%cut_filter.filter_type.lower()], obj.param["flux_%s"%cut_filter.filter_type.lower()] = self.cat.convert_sed(
Fang Yuedong's avatar
Fang Yuedong committed
220
221
222
223
224
                        mag=obj.param["mag_use_normal"],
                        sed=sed_data,
                        target_filt=cut_filter, 
                        norm_filt=norm_filt,
                    )
225
                except Exception as e:
Fang Yuedong's avatar
Fang Yuedong committed
226
                    traceback.print_exc()
Fang Yuedong's avatar
Fang Yuedong committed
227
                    chip_output.Log_error(e)
228
                    continue
229
                
Fang Yuedong's avatar
Fang Yuedong committed
230
                # [TODO] Testing
231
                # chip_output.Log_info("mag_%s = %.3f"%(filt.filter_type.lower(), obj.param["mag_%s"%filt.filter_type.lower()]))
Fang Yuedong's avatar
Fang Yuedong committed
232
233

                # Exclude very bright/dim objects (for now)
234
235
236
                if cut_filter.is_too_bright(
                    mag=obj.param["mag_%s"%self.config["obs_setting"]["cut_in_band"].lower()],
                    margin=self.config["obs_setting"]["mag_sat_margin"]):
Fang Yuedong's avatar
Fang Yuedong committed
237
238
239
240
                    chip_output.Log_info("obj %s too birght!! mag_%s = %.3f"%(obj.id, cut_filter.filter_type, obj.param["mag_%s"%self.config["obs_setting"]["cut_in_band"].lower()]))
                    bright_obj += 1
                    obj.unload_SED()
                    continue
241
242
243
                if filt.is_too_dim(
                    mag=obj.getMagFilter(filt),
                    margin=self.config["obs_setting"]["mag_lim_margin"]):
Fang Yuedong's avatar
Fang Yuedong committed
244
                    chip_output.Log_info("obj %s too dim!! mag_%s = %.3f"%(obj.id, filt.filter_type, obj.getMagFilter(filt)))
Fang Yuedong's avatar
Fang Yuedong committed
245
                    dim_obj += 1
Fang Yuedong's avatar
Fang Yuedong committed
246
247
248
                    obj.unload_SED()
                    continue

Fang Yuedong's avatar
Fang Yuedong committed
249
                # Get corresponding shear values
Fang Yuedong's avatar
Fang Yuedong committed
250
                if self.config["shear_setting"]["shear_type"] == "constant":
Fang Yuedong's avatar
Fang Yuedong committed
251
                    if obj.type == 'star':
252
                        obj.g1, obj.g2 = 0., 0.
Fang Yuedong's avatar
Fang Yuedong committed
253
                    else:
254
                        obj.g1, obj.g2 = self.g1_field, self.g2_field
255
256
257
                elif self.config["shear_setting"]["shear_type"] == "catalog":
                    pass
                else:
Fang Yuedong's avatar
Fang Yuedong committed
258
                    chip_output.Log_error("Unknown shear input")
259
                    raise ValueError("Unknown shear input")
Fang Yuedong's avatar
Fang Yuedong committed
260
261

                # Get position of object on the focal plane
262
                pos_img, offset, local_wcs, real_wcs, fd_shear = obj.getPosImg_Offset_WCS(img=chip.img, fdmodel=self.fd_model, chip=chip, verbose=False, chip_wcs=chip_wcs, img_header=h_ext)
Fang Yuedong's avatar
Fang Yuedong committed
263
264
265
266

                # [TODO] For now, only consider objects which their centers (after field distortion) are projected within the focal plane
                # Otherwise they will be considered missed objects
                # if pos_img.x == -1 or pos_img.y == -1 or (not chip.isContainObj(x_image=pos_img.x, y_image=pos_img.y, margin=0.)):
Fang Yuedong's avatar
Fang Yuedong committed
267
                if pos_img.x == -1 or pos_img.y == -1:
Fang Yuedong's avatar
Fang Yuedong committed
268
269
                    chip_output.Log_info('obj_ra = %.6f, obj_dec = %.6f, obj_ra_orig = %.6f, obj_dec_orig = %.6f'%(obj.ra, obj.dec, obj.ra_orig, obj.dec_orig))
                    chip_output.Log_error("Objected missed: %s"%(obj.id))
Fang Yuedong's avatar
Fang Yuedong committed
270
271
272
                    missed_obj += 1
                    obj.unload_SED()
                    continue
Fang Yuedong's avatar
Fang Yuedong committed
273

Fang Yuedong's avatar
Fang Yuedong committed
274
275
                # Draw object & update output catalog
                try:
Fang Yuedong's avatar
Fang Yuedong committed
276
                    if self.config["run_option"]["out_cat_only"]:
Fang Yuedong's avatar
Fang Yuedong committed
277
                        isUpdated = True
278
                        obj.real_pos = obj.getRealPos(chip.img, global_x=obj.posImg.x, global_y=obj.posImg.y, img_real_wcs=obj.chip_wcs)
279
                        pos_shear = 0.
280
                    elif chip.survey_type == "photometric" and not self.config["run_option"]["out_cat_only"]:
Fang Yuedong's avatar
Fang Yuedong committed
281
282
283
284
285
286
287
                        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, 
288
289
                            g1=obj.g1, 
                            g2=obj.g2, 
Fang Yuedong's avatar
Fang Yuedong committed
290
291
292
                            exptime=pointing.exp_time,
                            fd_shear=fd_shear)

Fang Yuedong's avatar
Fang Yuedong committed
293
                    elif chip.survey_type == "spectroscopic" and not self.config["run_option"]["out_cat_only"]:
Fang Yuedong's avatar
Fang Yuedong committed
294
295
296
297
298
299
300
                        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, 
301
302
                            g1=obj.g1, 
                            g2=obj.g2, 
303
                            exptime=pointing.exp_time,
Fang Yuedong's avatar
Fang Yuedong committed
304
305
306
307
                            normFilter=norm_filt,
                            fd_shear=fd_shear)

                    if isUpdated == 1:
Fang Yuedong's avatar
Fang Yuedong committed
308
                        # TODO: add up stats
309
                        chip_output.cat_add_obj(obj, pos_img, pos_shear)
Fang Yuedong's avatar
Fang Yuedong committed
310
                        pass
Fang Yuedong's avatar
Fang Yuedong committed
311
312
313
                    elif isUpdated == 0:
                        missed_obj += 1
                        chip_output.Log_error("Objected missed: %s"%(obj.id))
Fang Yuedong's avatar
Fang Yuedong committed
314
                    else:
Fang Yuedong's avatar
Fang Yuedong committed
315
                        chip_output.Log_error("Draw error, object omitted: %s"%(obj.id))
Fang Yuedong's avatar
Fang Yuedong committed
316
317
                        continue
                except Exception as e:
Fang Yuedong's avatar
Fang Yuedong committed
318
                    traceback.print_exc()
Fang Yuedong's avatar
Fang Yuedong committed
319
320
321
322
323
324
325
                    chip_output.Log_error(e)
                    
                # # [C6 TEST]
                # chip_output.Log_info("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) ))
                # chip_output.Log_info('draw object %s'%obj.id)
                # chip_output.Log_info('mag = %.3f'%obj.param['mag_use_normal'])

Fang Yuedong's avatar
Fang Yuedong committed
326
327
328
                # Unload SED:
                obj.unload_SED()
                del obj
Fang Yuedong's avatar
Fang Yuedong committed
329
                gc.collect()
Fang Yuedong's avatar
Fang Yuedong committed
330

Fang Yuedong's avatar
Fang Yuedong committed
331
332
            del psf_model
            del self.cat
Fang Yuedong's avatar
Fang Yuedong committed
333
            gc.collect()
Fang Yuedong's avatar
Fang Yuedong committed
334

Fang Yuedong's avatar
Fang Yuedong committed
335
        chip_output.Log_info("check running:1: pointing-%d chip-%d pid-%d memory-%6.2fGB"%(pointing.id, chip.chipID, os.getpid(), (psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024 / 1024) ))
Fang Yuedong's avatar
Fang Yuedong committed
336
337
338

        # Detector Effects
        # ===========================================================
Fang Yuedong's avatar
Fang Yuedong committed
339
        # whether to output zero, dark, flat calibration images.
340
341
342
343
344
345

        if not self.config["run_option"]["out_cat_only"]:
            chip.img = chip.addEffects(
                config=self.config, 
                img=chip.img, 
                chip_output=chip_output, 
346
                filt=filt,
347
348
349
                ra_cen=pointing.ra, 
                dec_cen=pointing.dec,
                img_rot=pointing.img_pa,
Fang Yuedong's avatar
Fang Yuedong committed
350
                exptime=pointing.exp_time,
351
352
353
354
355
356
                pointing_ID=pointing.id,
                timestamp_obs=pointing.timestamp,
                pointing_type=pointing.pointing_type,
                sky_map=sky_map, tel = self.tel,
                logger=chip_output.logger)
            
357
            if pointing.pointing_type == 'SCI':
Fang Yuedong's avatar
Fang Yuedong committed
358
                datetime_obs = datetime.utcfromtimestamp(pointing.timestamp)
359
360
361
362
363
364
365
366
                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, 
                    pointNum = str(pointing.id),
                    ra=pointing.ra, 
                    dec=pointing.dec, 
Fang Yuedong's avatar
Fang Yuedong committed
367
                    pixel_scale=chip.pix_scale,
368
369
370
371
372
                    date=date_obs,
                    time_obs=time_obs,
                    exptime=pointing.exp_time,
                    im_type='SCI',
                    sat_pos=[pointing.sat_x, pointing.sat_y, pointing.sat_z],
Fang Yuedong's avatar
Fang Yuedong committed
373
                    sat_vel=[pointing.sat_vx, pointing.sat_vy, pointing.sat_vz],
374
375
                    project_cycle=self.config["project_cycle"],
                    run_counter=self.config["run_counter"],
Fang Yuedong's avatar
Fang Yuedong committed
376
                    chip_name=str(chip.chipID).rjust(2, '0'))
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
                h_ext = generateExtensionHeader(
                    chip=chip,
                    xlen=chip.npix_x,
                    ylen=chip.npix_y,
                    ra=pointing.ra,
                    dec=pointing.dec,
                    pa=pointing.img_pa.deg,
                    gain=chip.gain,
                    readout=chip.read_noise,
                    dark=chip.dark_noise,
                    saturation=90000,
                    pixel_scale=chip.pix_scale,
                    pixel_size=chip.pix_size,
                    xcen=chip.x_cen,
                    ycen=chip.y_cen,
                    extName='SCI',
                    timestamp=pointing.timestamp,
                    exptime=pointing.exp_time,
395
                    readoutTime=chip.readout_time)
Fang Yuedong's avatar
Fang Yuedong committed
396
                
397
398
                chip.img = galsim.Image(chip.img.array, dtype=np.uint16)
                hdu1 = fits.PrimaryHDU(header=h_prim)
Fang Yuedong's avatar
Fang Yuedong committed
399
                hdu1.add_checksum()
400
401
                hdu1.header.comments['CHECKSUM'] = 'HDU checksum'
                hdu1.header.comments['DATASUM'] = 'data unit checksum'
402
                hdu2 = fits.ImageHDU(chip.img.array, header=h_ext)
Fang Yuedong's avatar
Fang Yuedong committed
403
                hdu2.add_checksum()
404
405
406
                hdu2.header.comments['XTENSION'] = 'extension type'
                hdu2.header.comments['CHECKSUM'] = 'HDU checksum'
                hdu2.header.comments['DATASUM'] = 'data unit checksum'
407
408
409
                hdu1 = fits.HDUList([hdu1, hdu2])
                fname = os.path.join(chip_output.subdir, h_prim['FILENAME'] + '.fits')
                hdu1.writeto(fname, output_verify='ignore', overwrite=True)
Fang Yuedong's avatar
Fang Yuedong committed
410
411
412
413

                chip_output.Log_info("# objects that are too bright %d out of %d"%(bright_obj, self.nobj))
                chip_output.Log_info("# objects that are too dim %d out of %d"%(dim_obj, self.nobj))
                chip_output.Log_info("# objects that are missed %d out of %d"%(missed_obj, self.nobj))
Fang Yuedong's avatar
Fang Yuedong committed
414
415
        del chip.img

Fang Yuedong's avatar
Fang Yuedong committed
416
        chip_output.Log_info("check running:2: pointing-%d chip-%d pid-%d memory-%6.2fGB"%(pointing.id, chip.chipID, os.getpid(), (psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024 / 1024) ))
Fang Yuedong's avatar
Fang Yuedong committed
417

Zhang Xin's avatar
Zhang Xin committed
418
    def run_one_chip_calibration(self, chip, filt, pointing, chip_output, skyback_level = 20000, sky_level_filt = 'g', wcs_fp=None, psf_model=None, cat_dir=None, sed_dir=None):
Zhang Xin's avatar
Zhang Xin committed
419
420
421
422
423
424




        # # Get WCS for the focal plane
        # if wcs_fp == None:
Zhang Xin's avatar
Zhang Xin committed
425
        # wcs_fp = self.focal_plane.getTanWCS(ra_cen, dec_cen, pointing.img_pa, chip.pix_scale)
Zhang Xin's avatar
Zhang Xin committed
426
427
428
429
430

        # 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
Zhang Xin's avatar
Zhang Xin committed
431
432
433
434
        pf_map = np.zeros_like(chip.img.array)
        if self.config["obs_setting"]["LED_TYPE"] is not None:
            if len(self.config["obs_setting"]["LED_TYPE"]) != 0:
                print("LED OPEN--------")
Zhang Xin's avatar
Zhang Xin committed
435

Zhang Xin's avatar
Zhang Xin committed
436
                led_obj = FlatLED(chip, filt)
Zhang Xin's avatar
Zhang Xin committed
437

Zhang Xin's avatar
Zhang Xin committed
438
439
                led_flat = led_obj.drawObj_LEDFlat(led_type_list=self.config["obs_setting"]["LED_TYPE"], exp_t_list=self.config["obs_setting"]["LED_TIME"])
                pf_map = led_flat
Zhang Xin's avatar
Zhang Xin committed
440
        # whether to output zero, dark, flat calibration images.
Zhang Xin's avatar
Zhang Xin committed
441
        expTime = self.config["obs_setting"]["exp_time"]
442
        norm_scaler = skyback_level/expTime / self.filter_param.param[sky_level_filt][5]
Zhang Xin's avatar
Zhang Xin committed
443
444
445

        if skyback_level == 0:
            self.config["ins_effects"]["shutter_effect"] = False
Zhang Xin's avatar
Zhang Xin committed
446

Zhang Xin's avatar
Zhang Xin committed
447
        if chip.survey_type == "photometric":
448
            sky_map = np.ones_like(chip.img.array) * norm_scaler * self.filter_param.param[chip.filter_type][5] / self.tel.pupil_area
Zhang Xin's avatar
Zhang Xin committed
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
        elif chip.survey_type == "spectroscopic":
            flat_normal = np.ones_like(chip.img.array)
            if self.config["ins_effects"]["flat_fielding"] == True:
                chip_output.Log_info("SLS flat preprocess,CHIP %d : Creating and applying Flat-Fielding" % chip.chipID)
                msg = str(chip.img.bounds)
                chip_output.Log_info(msg)
                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:
                chip_output.Log_info("SLS flat preprocess,CHIP %d : Apply shutter effect" % chip.chipID)
                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,
                conf=chip.sls_conf,
                pixelSize=chip.pix_scale,
                isAlongY=0,
                flat_cube=chip.flat_cube)
            sky_map = sky_map * norm_scaler
Zhang Xin's avatar
Zhang Xin committed
474
475
476
477
478
479
480
481
482
483
484
485
486
487

        chip.img = chip.addEffects(
            config=self.config,
            img=chip.img,
            chip_output=chip_output,
            filt=filt,
            ra_cen=pointing.ra,
            dec_cen=pointing.dec,
            img_rot=pointing.img_pa,
            exptime=self.config["obs_setting"]["exp_time"],
            pointing_ID=pointing.id,
            timestamp_obs=pointing.timestamp,
            pointing_type=pointing.pointing_type,
            sky_map=sky_map, tel=self.tel,
Zhang Xin's avatar
Zhang Xin committed
488
            post_flash_map=pf_map,
Zhang Xin's avatar
Zhang Xin committed
489
490
            logger=chip_output.logger)

Zhang Xin's avatar
Zhang Xin committed
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
        datetime_obs = datetime.utcfromtimestamp(pointing.timestamp)
        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,
            pointNum=str(pointing.id),
            ra=pointing.ra,
            dec=pointing.dec,
            pixel_scale=chip.pix_scale,
            date=date_obs,
            time_obs=time_obs,
            exptime=self.config["obs_setting"]["exp_time"],
            im_type='DARKPF',
            sat_pos=[pointing.sat_x, pointing.sat_y, pointing.sat_z],
            sat_vel=[pointing.sat_vx, pointing.sat_vy, pointing.sat_vz],
            chip_name=str(chip.chipID).rjust(2, '0'))
        h_ext = generateExtensionHeader(
            chip=chip,
            xlen=chip.npix_x,
            ylen=chip.npix_y,
            ra=pointing.ra,
            dec=pointing.dec,
            pa=pointing.img_pa.deg,
            gain=chip.gain,
            readout=chip.read_noise,
            dark=chip.dark_noise,
            saturation=90000,
            pixel_scale=chip.pix_scale,
            pixel_size=chip.pix_size,
            xcen=chip.x_cen,
            ycen=chip.y_cen,
            extName='SCI',
            timestamp=pointing.timestamp,
            exptime=self.config["obs_setting"]["exp_time"],
            readoutTime=chip.readout_time)

        chip.img = galsim.Image(chip.img.array, dtype=np.uint16)
        hdu1 = fits.PrimaryHDU(header=h_prim)
        hdu1.add_checksum()
        hdu1.header.comments['CHECKSUM'] = 'HDU checksum'
        hdu1.header.comments['DATASUM'] = 'data unit checksum'
        hdu2 = fits.ImageHDU(chip.img.array, header=h_ext)
        hdu2.add_checksum()
        hdu2.header.comments['XTENSION'] = 'extension type'
        hdu2.header.comments['CHECKSUM'] = 'HDU checksum'
        hdu2.header.comments['DATASUM'] = 'data unit checksum'
        hdu1 = fits.HDUList([hdu1, hdu2])
        fname = os.path.join(chip_output.subdir, h_prim['FILENAME'] + '.fits')
        hdu1.writeto(fname, output_verify='ignore', overwrite=True)
Zhang Xin's avatar
Zhang Xin committed
541
542
543
544
545
546
547
548
549

            # chip_output.Log_info("# objects that are too bright %d out of %d" % (bright_obj, self.nobj))
            # chip_output.Log_info("# objects that are too dim %d out of %d" % (dim_obj, self.nobj))
            # chip_output.Log_info("# objects that are missed %d out of %d" % (missed_obj, self.nobj))
        del chip.img

        chip_output.Log_info("check running:2: pointing-%d chip-%d pid-%d memory-%6.2fGB" % (
        pointing.id, chip.chipID, os.getpid(), (psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024 / 1024)))

550
    def runExposure_MPI_PointingList(self, pointing_list,chips=None, use_mpi=False):
Fang Yuedong's avatar
Fang Yuedong committed
551
552
553
554
        if use_mpi:
            comm = MPI.COMM_WORLD
            ind_thread = comm.Get_rank()
            num_thread = comm.Get_size()
Fang Yuedong's avatar
Fang Yuedong committed
555

Fang Yuedong's avatar
Fang Yuedong committed
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
        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
571

Zhang Xin's avatar
Zhang Xin committed
572

573
        for ipoint in range(len(pointing_list)):
Fang Yuedong's avatar
Fang Yuedong committed
574
575
            for ichip in range(nchips_per_fp):
                i = ipoint*nchips_per_fp + ichip
576
577
                pointing = pointing_list[ipoint]
                pointing_ID = pointing.id
Fang Yuedong's avatar
Fang Yuedong committed
578
579
580
                if use_mpi:
                    if i % num_thread != ind_thread:
                        continue
Fang Yuedong's avatar
Fang Yuedong committed
581
582
583
584
585

                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
586
587
                chip = run_chips[ichip]
                filt = run_filts[ichip]
Fang Yuedong's avatar
Fang Yuedong committed
588
                # chip_output.Log_info("running pointing#%d, chip#%d, at PID#%d..."%(pointing_ID, chip.chipID, pid))
Fang Yuedong's avatar
Fang Yuedong committed
589
590
591
592
593
                chip_output = ChipOutput(
                    config=self.config, 
                    focal_plane=self.focal_plane, 
                    chip=chip, 
                    filt=filt,  
594
595
                    exptime=pointing.exp_time,
                    pointing_type=pointing.pointing_type,
Fang Yuedong's avatar
Fang Yuedong committed
596
597
598
                    pointing_ID=pointing_ID,  
                    subdir=sub_img_dir,
                    prefix=prefix)
Fang Yuedong's avatar
Fang Yuedong committed
599
                chip_output.Log_info("running pointing#%d, chip#%d, at PID#%d..."%(pointing_ID, chip.chipID, pid))
Zhang Xin's avatar
Zhang Xin committed
600
601
                if self.config["obs_setting"]["survey_type"] == "CALIBRATION":
                    self.run_one_chip_calibration(chip=chip,
Zhang Xin's avatar
Zhang Xin committed
602
603
                                               filt=filt,
                                               chip_output=chip_output,
Zhang Xin's avatar
Zhang Xin committed
604
605
606
                                               pointing=pointing,
                                               skyback_level = self.config["obs_setting"]["FLAT_LEVEL"],
                                               sky_level_filt = self.config["obs_setting"]["FLAT_LEVEL_FIL"])
Zhang Xin's avatar
Zhang Xin committed
607
608
609
610
611
612
                else:
                    self.run_one_chip(
                        chip=chip,
                        filt=filt,
                        chip_output=chip_output,
                        pointing=pointing)
Fang Yuedong's avatar
Fang Yuedong committed
613
                chip_output.Log_info("finished running chip#%d..."%(chip.chipID))
614
615
                for handler in chip_output.logger.handlers[:]:
                    chip_output.logger.removeHandler(handler)
Fang Yuedong's avatar
Fang Yuedong committed
616
                gc.collect()