ObservationSim.py 22.7 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
17
from ObservationSim.Astrometry.Astrometry_util import on_orbit_obs_position
Fang Yuedong's avatar
Fang Yuedong committed
18
19

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

        # if we want to apply field distortion?
Fang Yuedong's avatar
Fang Yuedong committed
32
        if self.config["ins_effects"]["field_dist"] == True:
Fang Yuedong's avatar
Fang Yuedong committed
33
            self.fd_model = FieldDistortion(fdModel_path=self.path_dict["fd_path"])
Fang Yuedong's avatar
Fang Yuedong committed
34
35
36
37
38
39
40
41
42
        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

            # Make Chip & Filter lists
43
44
            chip = Chip(
                chipID=chipID, 
Fang Yuedong's avatar
Fang Yuedong committed
45
                config=self.config)
Fang Yuedong's avatar
Fang Yuedong committed
46
            filter_id, filter_type = chip.getChipFilter()
47
48
            filt = Filter(filter_id=filter_id, 
                filter_type=filter_type, 
49
                filter_param=self.filter_param)
50
51
52
53
            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
54
55

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

58
    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
59

60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
        # 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:::::::::::::::::::::::::::::::::::')
        chip_output.logger.info(':::::::::::::::::::Current Pointing Information::::::::::::::::::')
        chip_output.logger.info("RA: %f, DEC; %f" % (pointing.ra, pointing.dec))
        chip_output.logger.info("Time: %s" % datetime.fromtimestamp(pointing.timestamp).isoformat())
        chip_output.logger.info("Exposure time: %f" % pointing.exp_time)
        chip_output.logger.info("Satellite Position (x, y, z): (%f, %f, %f)" % (pointing.sat_x, pointing.sat_y, pointing.sat_z))
        chip_output.logger.info("Satellite Velocity (x, y, z): (%f, %f, %f)" % (pointing.sat_vx, pointing.sat_vy, pointing.sat_vz))
        chip_output.logger.info("Position Angle: %f" % pointing.img_pa.deg)
        chip_output.logger.info('Chip : %d' % chip.chipID)
        chip_output.logger.info(':::::::::::::::::::::::::::END:::::::::::::::::::::::::::::::::::')
Fang Yuedong's avatar
Fang Yuedong committed
78

Fang Yuedong's avatar
Fang Yuedong committed
79
        if self.config["psf_setting"]["psf_model"] == "Gauss":
Fang Yuedong's avatar
Fang Yuedong committed
80
            psf_model = PSFGauss(chip=chip)
Fang Yuedong's avatar
Fang Yuedong committed
81
        elif self.config["psf_setting"]["psf_model"] == "Interp":
Fang Yuedong's avatar
Fang Yuedong committed
82
            psf_model = PSFInterp(chip=chip, PSF_data_file=self.path_dict["psf_dir"])
Fang Yuedong's avatar
Fang Yuedong committed
83
        else:
84
85
            # print("unrecognized PSF model type!!", flush=True)
            chip_output.logger.error("unrecognized PSF model type!!", flush=True)
Fang Yuedong's avatar
Fang Yuedong committed
86
87
88

        # Get (extra) shear fields
        if shear_cat_file is not None:
89
            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
90

91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
        # Apply astrometric simulation for pointing
        if self.config["obs_setting"]["enable_astrometric_model"]:
            dt = datetime.fromtimestamp(pointing.timestamp)
            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,
                input_epoch="J2015.5",
                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
119
120
        # Get WCS for the focal plane
        if wcs_fp == None:
121
            wcs_fp = self.focal_plane.getTanWCS(ra_cen, dec_cen, pointing.img_pa, chip.pix_scale)
Fang Yuedong's avatar
Fang Yuedong committed
122
123
124
125
126
127
128

        # 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
129
130
        # 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
131
        elif chip.survey_type == "spectroscopic":
132
133
            flat_normal = np.ones_like(chip.img.array)
            if self.config["ins_effects"]["flat_fielding"] == True:
134
135
136
137
138
                # print("SLS flat preprocess,CHIP %d : Creating and applying Flat-Fielding"%chip.chipID, flush=True)
                # print(chip.img.bounds, flush=True)
                chip_output.logger.info("SLS flat preprocess,CHIP %d : Creating and applying Flat-Fielding"%chip.chipID)
                msg = str(chip.img.bounds)
                chip_output.logger.info(msg)
139
140
141
142
143
                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:
144
145
                # print("SLS flat preprocess,CHIP %d : Apply shutter effect"%chip.chipID, flush=True)
                chip_output.logger.info("SLS flat preprocess,CHIP %d : Apply shutter effect"%chip.chipID)
146
147
148
149
                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
150
151
152
153
154
155
156
            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)
157
            del flat_normal
Fang Yuedong's avatar
Fang Yuedong committed
158

159
        if pointing.pointing_type == 'MS':
Fang Yuedong's avatar
Fang Yuedong committed
160
            # Load catalogues and templates
161
162
            self.cat = self.Catalog(config=self.config, chip=chip, pointing=pointing, cat_dir=cat_dir, sed_dir=sed_dir, chip_output=chip_output)
            chip_output.create_output_file()
Fang Yuedong's avatar
Fang Yuedong committed
163
164
            self.nobj = len(self.cat.objs)

Fang Yuedong's avatar
Fang Yuedong committed
165
166
167
168
169
170
171
172
173
            for ifilt in range(len(self.all_filter)):
                temp_filter = self.all_filter[ifilt]
                # Update the limiting magnitude using exposure time in pointing
                temp_filter.update_limit_saturation_mags(exptime=pointing.exp_time, chip=chip)

                # 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
174
175
176
177
178
            # Loop over objects
            missed_obj = 0
            bright_obj = 0
            dim_obj = 0
            for j in range(self.nobj):
179
180
                
                # (DEBUG)
181
                # if j >= 10:
Xin Zhang's avatar
Xin Zhang committed
182
                #     break
183

Fang Yuedong's avatar
Fang Yuedong committed
184
                obj = self.cat.objs[j]
Fang Yuedong's avatar
Fang Yuedong committed
185
                if obj.type == 'star' and self.config["run_option"]["galaxy_only"]:
186
                    continue
Fang Yuedong's avatar
Fang Yuedong committed
187
                elif obj.type == 'galaxy' and self.config["run_option"]["star_only"]:
188
                    continue
Fang Yuedong's avatar
Fang Yuedong committed
189
                elif obj.type == 'quasar' and self.config["run_option"]["star_only"]:
190
                    continue
Fang Yuedong's avatar
Fang Yuedong committed
191

192
                # load and convert SED; also caculate object's magnitude in all CSST bands
193
194
195
                try:
                    sed_data = self.cat.load_sed(obj)
                    norm_filt = self.cat.load_norm_filt(obj)
196
                    obj.sed, obj.param["mag_%s"%filt.filter_type], obj.param["flux_%s"%filt.filter_type] = self.cat.convert_sed(
197
198
199
200
201
                        mag=obj.param["mag_use_normal"],
                        sed=sed_data,
                        target_filt=filt, 
                        norm_filt=norm_filt,
                    )
202
                    _, obj.param["mag_%s"%cut_filter.filter_type], obj.param["flux_%s"%cut_filter.filter_type] = self.cat.convert_sed(
Fang Yuedong's avatar
Fang Yuedong committed
203
204
205
206
207
                        mag=obj.param["mag_use_normal"],
                        sed=sed_data,
                        target_filt=cut_filter, 
                        norm_filt=norm_filt,
                    )
208

209
                except Exception as e:
210
                    print(e)
211
                    chip_output.logger.error(e)
212
                    continue
213
214
                
                # chip_output.logger.info("debug point #1")
Fang Yuedong's avatar
Fang Yuedong committed
215
216

                # Exclude very bright/dim objects (for now)
217
                # if filt.is_too_bright(mag=obj.getMagFilter(filt)):
218
                # if filt.is_too_bright(mag=obj.mag_use_normal):
219
220
221
                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
222
223
224
225
226
                    # print("obj too birght!!", flush=True)
                    if obj.type != 'galaxy':
                        bright_obj += 1
                        obj.unload_SED()
                        continue
227
228
229
                if filt.is_too_dim(
                    mag=obj.getMagFilter(filt),
                    margin=self.config["obs_setting"]["mag_lim_margin"]):
230
                # if cut_filter.is_too_dim(mag=obj.param["mag_%s"%self.config["obs_setting"]["cut_in_band"].lower()]):
Fang Yuedong's avatar
Fang Yuedong committed
231
232
                    # print("obj too dim!!", flush=True)
                    dim_obj += 1
Fang Yuedong's avatar
Fang Yuedong committed
233
                    obj.unload_SED()
Fang Yuedong's avatar
Fang Yuedong committed
234
                    # print(obj.getMagFilter(filt))
Fang Yuedong's avatar
Fang Yuedong committed
235
236
                    continue

237
238
                # chip_output.logger.info("debug point #2")

Fang Yuedong's avatar
Fang Yuedong committed
239
                if self.config["shear_setting"]["shear_type"] == "constant":
Fang Yuedong's avatar
Fang Yuedong committed
240
                    if obj.type == 'star':
241
                        obj.g1, obj.g2 = 0., 0.
Fang Yuedong's avatar
Fang Yuedong committed
242
                    else:
243
                        obj.g1, obj.g2 = self.g1_field, self.g2_field
Fang Yuedong's avatar
Fang Yuedong committed
244
245
246
                elif self.config["shear_setting"]["shear_type"] == "extra":
                    try:
                        # TODO: every object with individual shear from input catalog(s)
247
                        obj.g1, obj.g2 = self.g1_field[j], self.g2_field[j]
Fang Yuedong's avatar
Fang Yuedong committed
248
                    except:
249
250
                        # print("failed to load external shear.")
                        chip_output.logger.error("failed to load external shear.")
Fang Yuedong's avatar
Fang Yuedong committed
251
                        pass
252
253

                    # chip_output.logger.info("debug point #3")
254
255
256
                elif self.config["shear_setting"]["shear_type"] == "catalog":
                    pass
                else:
257
                    chip_output.logger.error("Unknown shear input")
258
                    raise ValueError("Unknown shear input")
259
260
                
                # chip_output.logger.info("debug point #4")
xin's avatar
xin committed
261
262
263
264
265
266
267
268
269
270
271
272
273
274
                h_ext = generateExtensionHeader(
                    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,
                    psize=chip.pix_scale,
                    row_num=chip.rowID,
                    col_num=chip.colID,
                    extName='raw')
Fang Yuedong's avatar
Fang Yuedong committed
275

xin's avatar
xin committed
276
                pos_img, offset, local_wcs, real_wcs = obj.getPosImg_Offset_WCS(img=chip.img, fdmodel=self.fd_model, chip=chip, verbose=False)
Fang Yuedong's avatar
Fang Yuedong committed
277
278
279
280
281
282
                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
283

284
285
                # chip_output.logger.info("debug point #5")

Fang Yuedong's avatar
Fang Yuedong committed
286
287
                # Draw object & update output catalog
                try:
288
289
                    # chip_output.logger.info("debug point #6")
                    # chip_output.logger.info("current filter type: %s"%filt.filter_type)
Fang Yuedong's avatar
Fang Yuedong committed
290
                    if self.config["run_option"]["out_cat_only"]:
Fang Yuedong's avatar
Fang Yuedong committed
291
                        isUpdated = True
292
                        pos_shear = 0.
293
                    elif chip.survey_type == "photometric" 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_multiband(
                            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
304
                            exptime=pointing.exp_time
                            )
Fang Yuedong's avatar
Fang Yuedong committed
305
                    elif chip.survey_type == "spectroscopic" and not self.config["run_option"]["out_cat_only"]:
Fang Yuedong's avatar
Fang Yuedong committed
306
307
308
309
310
311
312
                        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, 
313
314
                            g1=obj.g1, 
                            g2=obj.g2, 
315
                            exptime=pointing.exp_time,
316
317
                            normFilter=norm_filt,
                            )
318
                    # chip_output.logger.info("debug point #7")
Fang Yuedong's avatar
Fang Yuedong committed
319
320
                    if isUpdated:
                        # TODO: add up stats
321
322
                        # print("updating output catalog...")
                        chip_output.cat_add_obj(obj, pos_img, pos_shear)
Fang Yuedong's avatar
Fang Yuedong committed
323
324
325
326
327
                        pass
                    else:
                        # print("object omitted", flush=True)
                        continue
                except Exception as e:
328
                    print(e)
329
                    chip_output.logger.error(e)
Fang Yuedong's avatar
Fang Yuedong committed
330
                    pass
Fang Yuedong's avatar
Fang Yuedong committed
331
332
333
                # Unload SED:
                obj.unload_SED()
                del obj
Fang Yuedong's avatar
Fang Yuedong committed
334

Fang Yuedong's avatar
Fang Yuedong committed
335
336
            del psf_model
            del self.cat
Fang Yuedong's avatar
Fang Yuedong committed
337

338
339
        # 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)
        chip_output.logger.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
340
341
342

        # Detector Effects
        # ===========================================================
Fang Yuedong's avatar
Fang Yuedong committed
343
        # whether to output zero, dark, flat calibration images.
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404

        if not self.config["run_option"]["out_cat_only"]:
            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,
                pointing_ID=pointing.id,
                timestamp_obs=pointing.timestamp,
                pointing_type=pointing.pointing_type,
                sky_map=sky_map, tel = self.tel,
                logger=chip_output.logger)
            
            if pointing.pointing_type == 'MS':
                datetime_obs = datetime.fromtimestamp(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, 
                    psize=chip.pix_scale, 
                    row_num=chip.rowID, 
                    col_num=chip.colID,
                    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],
                    sat_vel=[pointing.sat_vx, pointing.sat_vy, pointing.sat_vz])
                h_ext = generateExtensionHeader(
                    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, 
                    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))
                chip_output.logger.info("# objects that are too bright %d out of %d"%(bright_obj, self.nobj))
                chip_output.logger.info("# objects that are too dim %d out of %d"%(dim_obj, self.nobj))
                chip_output.logger.info("# objects that are missed %d out of %d"%(missed_obj, self.nobj))
Fang Yuedong's avatar
Fang Yuedong committed
405
406
        del chip.img

407
408
        # 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)
        chip_output.logger.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
409

410
    def runExposure_MPI_PointingList(self, pointing_list, shear_cat_file=None, chips=None, use_mpi=False):
Fang Yuedong's avatar
Fang Yuedong committed
411
412
413
414
        if use_mpi:
            comm = MPI.COMM_WORLD
            ind_thread = comm.Get_rank()
            num_thread = comm.Get_size()
Fang Yuedong's avatar
Fang Yuedong committed
415

Fang Yuedong's avatar
Fang Yuedong committed
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
        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
431

432
        for ipoint in range(len(pointing_list)):
Fang Yuedong's avatar
Fang Yuedong committed
433
434
            for ichip in range(nchips_per_fp):
                i = ipoint*nchips_per_fp + ichip
435
436
                pointing = pointing_list[ipoint]
                pointing_ID = pointing.id
Fang Yuedong's avatar
Fang Yuedong committed
437
438
439
                if use_mpi:
                    if i % num_thread != ind_thread:
                        continue
Fang Yuedong's avatar
Fang Yuedong committed
440
441
442
443
444

                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
445
446
                chip = run_chips[ichip]
                filt = run_filts[ichip]
447
                # print("running pointing#%d, chip#%d, at PID#%d..."%(pointing_ID, chip.chipID, pid), flush=True)
Fang Yuedong's avatar
Fang Yuedong committed
448
449
450
451
452
                chip_output = ChipOutput(
                    config=self.config, 
                    focal_plane=self.focal_plane, 
                    chip=chip, 
                    filt=filt,  
453
454
                    exptime=pointing.exp_time,
                    pointing_type=pointing.pointing_type,
Fang Yuedong's avatar
Fang Yuedong committed
455
456
457
                    pointing_ID=pointing_ID,  
                    subdir=sub_img_dir,
                    prefix=prefix)
458
                chip_output.logger.info("running pointing#%d, chip#%d, at PID#%d..."%(pointing_ID, chip.chipID, pid))
459
                self.run_one_chip(
Fang Yuedong's avatar
Fang Yuedong committed
460
461
462
                    chip=chip, 
                    filt=filt, 
                    chip_output=chip_output, 
463
                    pointing=pointing,
Fang Yuedong's avatar
Fang Yuedong committed
464
                    cat_dir=self.path_dict["cat_dir"])
Zhang Xin's avatar
Zhang Xin committed
465
                print("finished running chip#%d..."%(chip.chipID), flush=True)
466
467
468
                chip_output.logger.info("finished running chip#%d..."%(chip.chipID))
                for handler in chip_output.logger.handlers[:]:
                    chip_output.logger.removeHandler(handler)