ObservationSim.py 29.6 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
200
201
                # # (DEBUG)
                # 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
                            normFilter=norm_filt,
                            fd_shear=fd_shear)

Wei Chengliang's avatar
Wei Chengliang committed
307
308
309
310
311
312
313
314
315
316
317
318
319
320
                    if isUpdated == 1 and self.config["run_option"]["out_psf"]:
                        obj.drawObj_PSF(
                            tel=self.tel,
                            pos_img=pos_img,
                            psf_model=psf_model,
                            bandpass_list=filt.bandpass_sub_list,
                            filt=filt,
                            chip=chip,
                            g1=obj.g1,
                            g2=obj.g2,
                            exptime=pointing.exp_time,
                            fd_shear=fd_shear,
                            chip_output=chip_output)

Fang Yuedong's avatar
Fang Yuedong committed
321
                    if isUpdated == 1:
Fang Yuedong's avatar
Fang Yuedong committed
322
                        # TODO: add up stats
323
                        chip_output.cat_add_obj(obj, pos_img, pos_shear)
Fang Yuedong's avatar
Fang Yuedong committed
324
                        pass
Fang Yuedong's avatar
Fang Yuedong committed
325
326
327
                    elif isUpdated == 0:
                        missed_obj += 1
                        chip_output.Log_error("Objected missed: %s"%(obj.id))
Fang Yuedong's avatar
Fang Yuedong committed
328
                    else:
Fang Yuedong's avatar
Fang Yuedong committed
329
                        chip_output.Log_error("Draw error, object omitted: %s"%(obj.id))
Fang Yuedong's avatar
Fang Yuedong committed
330
331
                        continue
                except Exception as e:
Fang Yuedong's avatar
Fang Yuedong committed
332
                    traceback.print_exc()
Fang Yuedong's avatar
Fang Yuedong committed
333
334
335
336
337
338
339
                    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
340
341
342
                # Unload SED:
                obj.unload_SED()
                del obj
Fang Yuedong's avatar
Fang Yuedong committed
343
                gc.collect()
Fang Yuedong's avatar
Fang Yuedong committed
344

Fang Yuedong's avatar
Fang Yuedong committed
345
346
            del psf_model
            del self.cat
Fang Yuedong's avatar
Fang Yuedong committed
347
            gc.collect()
Fang Yuedong's avatar
Fang Yuedong committed
348

Fang Yuedong's avatar
Fang Yuedong committed
349
        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
350
351
352

        # Detector Effects
        # ===========================================================
Fang Yuedong's avatar
Fang Yuedong committed
353
        # whether to output zero, dark, flat calibration images.
354
355
356
357
358
359

        if not self.config["run_option"]["out_cat_only"]:
            chip.img = chip.addEffects(
                config=self.config, 
                img=chip.img, 
                chip_output=chip_output, 
360
                filt=filt,
361
362
363
                ra_cen=pointing.ra, 
                dec_cen=pointing.dec,
                img_rot=pointing.img_pa,
Fang Yuedong's avatar
Fang Yuedong committed
364
                exptime=pointing.exp_time,
365
366
367
368
369
370
                pointing_ID=pointing.id,
                timestamp_obs=pointing.timestamp,
                pointing_type=pointing.pointing_type,
                sky_map=sky_map, tel = self.tel,
                logger=chip_output.logger)
            
371
            if pointing.pointing_type == 'SCI':
Fang Yuedong's avatar
Fang Yuedong committed
372
                datetime_obs = datetime.utcfromtimestamp(pointing.timestamp)
373
374
375
376
377
378
379
380
                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
381
                    pixel_scale=chip.pix_scale,
382
383
384
385
386
                    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
387
                    sat_vel=[pointing.sat_vx, pointing.sat_vy, pointing.sat_vz],
388
389
                    project_cycle=self.config["project_cycle"],
                    run_counter=self.config["run_counter"],
Fang Yuedong's avatar
Fang Yuedong committed
390
                    chip_name=str(chip.chipID).rjust(2, '0'))
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
                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,
409
                    readoutTime=chip.readout_time)
Fang Yuedong's avatar
Fang Yuedong committed
410
                
411
412
                chip.img = galsim.Image(chip.img.array, dtype=np.uint16)
                hdu1 = fits.PrimaryHDU(header=h_prim)
Fang Yuedong's avatar
Fang Yuedong committed
413
                hdu1.add_checksum()
414
415
                hdu1.header.comments['CHECKSUM'] = 'HDU checksum'
                hdu1.header.comments['DATASUM'] = 'data unit checksum'
416
                hdu2 = fits.ImageHDU(chip.img.array, header=h_ext)
Fang Yuedong's avatar
Fang Yuedong committed
417
                hdu2.add_checksum()
418
419
420
                hdu2.header.comments['XTENSION'] = 'extension type'
                hdu2.header.comments['CHECKSUM'] = 'HDU checksum'
                hdu2.header.comments['DATASUM'] = 'data unit checksum'
421
422
423
                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
424
425
426
427

                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
428
429
        del chip.img

Fang Yuedong's avatar
Fang Yuedong committed
430
        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
431

Zhang Xin's avatar
Zhang Xin committed
432
    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
433
434
435
436
437
438




        # # Get WCS for the focal plane
        # if wcs_fp == None:
Zhang Xin's avatar
Zhang Xin committed
439
        # wcs_fp = self.focal_plane.getTanWCS(ra_cen, dec_cen, pointing.img_pa, chip.pix_scale)
Zhang Xin's avatar
Zhang Xin committed
440
441
442
443
444

        # 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
445
446
447
448
        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
449

Zhang Xin's avatar
Zhang Xin committed
450
                led_obj = FlatLED(chip, filt)
Zhang Xin's avatar
Zhang Xin committed
451

Zhang Xin's avatar
Zhang Xin committed
452
453
                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
454
        # whether to output zero, dark, flat calibration images.
Zhang Xin's avatar
Zhang Xin committed
455
        expTime = self.config["obs_setting"]["exp_time"]
456
        norm_scaler = skyback_level/expTime / self.filter_param.param[sky_level_filt][5]
Zhang Xin's avatar
Zhang Xin committed
457
458
459

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

Zhang Xin's avatar
Zhang Xin committed
461
        if chip.survey_type == "photometric":
462
            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
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
        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
488
489
490
491
492
493
494
495
496
497
498
499
500
501

        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
502
            post_flash_map=pf_map,
Zhang Xin's avatar
Zhang Xin committed
503
504
            logger=chip_output.logger)

Zhang Xin's avatar
Zhang Xin committed
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
541
542
543
544
545
546
547
548
549
550
551
552
553
554
        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
555
556
557
558
559
560
561
562
563

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

564
    def runExposure_MPI_PointingList(self, pointing_list,chips=None, use_mpi=False):
Fang Yuedong's avatar
Fang Yuedong committed
565
566
567
568
        if use_mpi:
            comm = MPI.COMM_WORLD
            ind_thread = comm.Get_rank()
            num_thread = comm.Get_size()
Fang Yuedong's avatar
Fang Yuedong committed
569

Fang Yuedong's avatar
Fang Yuedong committed
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
        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
585

Zhang Xin's avatar
Zhang Xin committed
586

587
        for ipoint in range(len(pointing_list)):
Fang Yuedong's avatar
Fang Yuedong committed
588
589
            for ichip in range(nchips_per_fp):
                i = ipoint*nchips_per_fp + ichip
590
591
                pointing = pointing_list[ipoint]
                pointing_ID = pointing.id
Fang Yuedong's avatar
Fang Yuedong committed
592
593
594
                if use_mpi:
                    if i % num_thread != ind_thread:
                        continue
Fang Yuedong's avatar
Fang Yuedong committed
595
596
597
598
599

                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
600
601
                chip = run_chips[ichip]
                filt = run_filts[ichip]
Fang Yuedong's avatar
Fang Yuedong committed
602
                # chip_output.Log_info("running pointing#%d, chip#%d, at PID#%d..."%(pointing_ID, chip.chipID, pid))
Fang Yuedong's avatar
Fang Yuedong committed
603
604
605
606
607
                chip_output = ChipOutput(
                    config=self.config, 
                    focal_plane=self.focal_plane, 
                    chip=chip, 
                    filt=filt,  
608
609
                    exptime=pointing.exp_time,
                    pointing_type=pointing.pointing_type,
Fang Yuedong's avatar
Fang Yuedong committed
610
611
612
                    pointing_ID=pointing_ID,  
                    subdir=sub_img_dir,
                    prefix=prefix)
Fang Yuedong's avatar
Fang Yuedong committed
613
                chip_output.Log_info("running pointing#%d, chip#%d, at PID#%d..."%(pointing_ID, chip.chipID, pid))
Zhang Xin's avatar
Zhang Xin committed
614
615
                if self.config["obs_setting"]["survey_type"] == "CALIBRATION":
                    self.run_one_chip_calibration(chip=chip,
Zhang Xin's avatar
Zhang Xin committed
616
617
                                               filt=filt,
                                               chip_output=chip_output,
Zhang Xin's avatar
Zhang Xin committed
618
619
620
                                               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
621
622
623
624
625
626
                else:
                    self.run_one_chip(
                        chip=chip,
                        filt=filt,
                        chip_output=chip_output,
                        pointing=pointing)
Fang Yuedong's avatar
Fang Yuedong committed
627
                chip_output.Log_info("finished running chip#%d..."%(chip.chipID))
628
629
                for handler in chip_output.logger.handlers[:]:
                    chip_output.logger.removeHandler(handler)
Fang Yuedong's avatar
Fang Yuedong committed
630
                gc.collect()