Galaxy.py 21.7 KB
Newer Older
Fang Yuedong's avatar
Fang Yuedong committed
1
2
3
4
5
6
7
import numpy as np
import galsim
import os, sys
import astropy.constants as cons
from astropy.table import Table
from scipy import interpolate

8
9
10
11
from ObservationSim.MockObject._util import eObs, integrate_sed_bandpass, getNormFactorForSpecWithABMAG, getObservedSED, getABMAG,convolveGaussXorders
from ObservationSim.MockObject.SpecDisperser import SpecDisperser
from ObservationSim.MockObject.MockObject import MockObject

Fang Yuedong's avatar
Fang Yuedong committed
12
13
# import tracemalloc

Fang Yuedong's avatar
Fang Yuedong committed
14
class Galaxy(MockObject):
15
16
    def __init__(self, param, rotation=None, logger=None):
        super().__init__(param, logger=logger)
Fang Yuedong's avatar
Fang Yuedong committed
17
18
19
20
        # self.thetaR = self.param["theta"]
        # self.bfrac = self.param["bfrac"]
        # self.hlr_disk = self.param["hlr_disk"]
        # self.hlr_bulge = self.param["hlr_bulge"]
Fang Yuedong's avatar
Fang Yuedong committed
21
22

        # Extract ellipticity components
Fang Yuedong's avatar
Fang Yuedong committed
23
24
25
26
27
28
        # self.e_disk = galsim.Shear(g=self.param["ell_disk"], beta=self.thetaR*galsim.degrees)
        # self.e_bulge = galsim.Shear(g=self.param["ell_bulge"], beta=self.thetaR*galsim.degrees)
        # self.e_total = galsim.Shear(g=self.param["ell_tot"], beta=self.thetaR*galsim.degrees)
        # self.e1_disk, self.e2_disk = self.e_disk.g1, self.e_disk.g2
        # self.e1_bulge, self.e2_bulge = self.e_bulge.g1, self.e_bulge.g2
        # self.e1_total, self.e2_total = self.e_total.g1, self.e_total.g2
Fang Yuedong's avatar
Fang Yuedong committed
29
30
31

        if rotation is not None:
            self.rotateEllipticity(rotation)
Fang Yuedong's avatar
Fang Yuedong committed
32
33
34
35
        if not hasattr(self, "disk_sersic_idx"):
            self.disk_sersic_idx = 1.
        if not hasattr(self, "bulge_sersic_idx"):
            self.bulge_sersic_idx = 4.
Fang Yuedong's avatar
Fang Yuedong committed
36
37
38
39
40
41

    def unload_SED(self):
        """(Test) free up SED memory
        """
        del self.sed

Fang Yuedong's avatar
Fang Yuedong committed
42
    def getGSObj_multiband(self, tel, psf_list, bandpass_list, filt, nphotons_tot=None, g1=0, g2=0, exptime=150., fd_shear=None):
Fang Yuedong's avatar
Fang Yuedong committed
43
44
45
46
47
48
49
50
51
52
53
        if len(psf_list) != len(bandpass_list):
            raise ValueError("!!!The number of PSF profiles and the number of bandpasses must be equal.")
        objs = []
        if nphotons_tot == None:
            nphotons_tot = self.getElectronFluxFilt(filt, tel, exptime)
        # print("nphotons_tot = ", nphotons_tot)

        try:
            full = integrate_sed_bandpass(sed=self.sed, bandpass=filt.bandpass_full)
        except Exception as e:
            print(e)
54
55
            if self.logger:
                self.logger.error(e)
Fang Yuedong's avatar
Fang Yuedong committed
56
57
58
59
60
61
62
            return -1
        for i in range(len(bandpass_list)):
            bandpass = bandpass_list[i]
            try:
                sub = integrate_sed_bandpass(sed=self.sed, bandpass=bandpass)
            except Exception as e:
                print(e)
63
64
                if self.logger:
                    self.logger.error(e)
Fang Yuedong's avatar
Fang Yuedong committed
65
66
67
68
69
70
71
72
73
                return -1
            
            ratio = sub/full
            if not (ratio == -1 or (ratio != ratio)):
                nphotons = ratio * nphotons_tot
            else:
                return -1

            psf = psf_list[i]
Fang Yuedong's avatar
Fang Yuedong committed
74
            disk = galsim.Sersic(n=self.disk_sersic_idx, half_light_radius=self.hlr_disk, flux=1.0)
Fang Yuedong's avatar
Fang Yuedong committed
75
76
            disk_shape = galsim.Shear(g1=self.e1_disk, g2=self.e2_disk)
            disk = disk.shear(disk_shape)
Fang Yuedong's avatar
Fang Yuedong committed
77
            bulge = galsim.Sersic(n=self.bulge_sersic_idx, half_light_radius=self.hlr_bulge, flux=1.0)
Fang Yuedong's avatar
Fang Yuedong committed
78
79
80
81
82
            bulge_shape = galsim.Shear(g1=self.e1_bulge, g2=self.e2_bulge)
            bulge = bulge.shear(bulge_shape)

            gal = self.bfrac * bulge + (1.0 - self.bfrac) * disk
            gal = gal.withFlux(nphotons)
83
84
85
            if fd_shear is not None:
                g1 += fd_shear.g1
                g2 += fd_shear.g2
Fang Yuedong's avatar
Fang Yuedong committed
86
87
88
            gal_shear = galsim.Shear(g1=g1, g2=g2)
            gal = gal.shear(gal_shear)
            gal = galsim.Convolve(psf, gal)
Fang Yuedong's avatar
Fang Yuedong committed
89
            
90
91
            # if fd_shear is not None:
            #     gal = gal.shear(fd_shear)
Fang Yuedong's avatar
Fang Yuedong committed
92
93
94
95
            objs.append(gal)
        final = galsim.Sum(objs)
        return final

Fang Yuedong's avatar
Fang Yuedong committed
96
    def drawObj_multiband(self, tel, pos_img, psf_model, bandpass_list, filt, chip, nphotons_tot=None, g1=0, g2=0, exptime=150., fd_shear=None):
Fang Yuedong's avatar
Fang Yuedong committed
97
98
99
100
101
102
103
104
        if nphotons_tot == None:
            nphotons_tot = self.getElectronFluxFilt(filt, tel, exptime)
        # print("nphotons_tot = ", nphotons_tot)

        try:
            full = integrate_sed_bandpass(sed=self.sed, bandpass=filt.bandpass_full)
        except Exception as e:
            print(e)
105
106
            if self.logger:
                self.logger.error(e)
Fang Yuedong's avatar
Fang Yuedong committed
107
            return 2, None
Fang Yuedong's avatar
Fang Yuedong committed
108
109
110
111

        nphotons_sum = 0
        photons_list = []
        xmax, ymax = 0, 0
Fang Yuedong's avatar
Fang Yuedong committed
112
113
        
        # # [C6 TEST]
Fang Yuedong's avatar
Fang Yuedong committed
114
        # print('hlr_disk = %.4f, hlr_bulge = %.4f'%(self.hlr_disk, self.hlr_bulge))
Fang Yuedong's avatar
Fang Yuedong committed
115
116
        # tracemalloc.start()

Fang Yuedong's avatar
Fang Yuedong committed
117
        big_galaxy = False
Fang Yuedong's avatar
Fang Yuedong committed
118
        if self.hlr_disk > 3.0 or self.hlr_bulge > 3.0: # Very big galaxy
Fang Yuedong's avatar
Fang Yuedong committed
119
120
121
122
123
124
125
126
127
            big_galaxy = True

        # (TEST) Galsim Parameters
        if self.getMagFilter(filt) <= 15 and (not big_galaxy):
            folding_threshold = 5.e-4
        else:
            folding_threshold = 5.e-3
        gsp = galsim.GSParams(folding_threshold=folding_threshold)

xin's avatar
xin committed
128
129
130
131
132
133
134
135
136
137
138
139
        self.real_pos = self.getRealPos(chip.img, global_x=self.posImg.x, global_y=self.posImg.y,
                                        img_real_wcs=self.real_wcs)

        x, y = self.real_pos.x + 0.5, self.real_pos.y + 0.5
        x_nominal = int(np.floor(x + 0.5))
        y_nominal = int(np.floor(y + 0.5))
        dx = x - x_nominal
        dy = y - y_nominal
        offset = galsim.PositionD(dx, dy)

        real_wcs_local = self.real_wcs.local(self.real_pos)

140
141
142
143
144
145
146
        disk = galsim.Sersic(n=self.disk_sersic_idx, half_light_radius=self.hlr_disk, flux=1.0, gsparams=gsp)
        disk_shape = galsim.Shear(g1=self.e1_disk, g2=self.e2_disk)
        disk = disk.shear(disk_shape)
        bulge = galsim.Sersic(n=self.bulge_sersic_idx, half_light_radius=self.hlr_bulge, flux=1.0, gsparams=gsp)
        bulge_shape = galsim.Shear(g1=self.e1_bulge, g2=self.e2_bulge)
        bulge = bulge.shear(bulge_shape)

Fang Yuedong's avatar
Fang Yuedong committed
147
148
149
150
151
152
153
        for i in range(len(bandpass_list)):
            bandpass = bandpass_list[i]

            try:
                sub = integrate_sed_bandpass(sed=self.sed, bandpass=bandpass)
            except Exception as e:
                print(e)
154
155
                if self.logger:
                    self.logger.error(e)
Fang Yuedong's avatar
Fang Yuedong committed
156
157
158
159
160
161
162
163
164
                # return False
                continue
            
            ratio = sub/full
            if not (ratio == -1 or (ratio != ratio)):
                nphotons = ratio * nphotons_tot
            else:
                continue
            nphotons_sum += nphotons
Fang Yuedong's avatar
Fang Yuedong committed
165
166

            # # [C6 TEST]
Fang Yuedong's avatar
Fang Yuedong committed
167
168
169
            # print("nphotons_sub-band_%d = %.2f"%(i, nphotons))

            psf, pos_shear = psf_model.get_PSF(chip=chip, pos_img=pos_img, bandpass=bandpass, folding_threshold=folding_threshold)
170
171
172
173
174
175
176
177
            # disk = galsim.Sersic(n=self.disk_sersic_idx, half_light_radius=self.hlr_disk, flux=1.0, gsparams=gsp)
            # disk_shape = galsim.Shear(g1=self.e1_disk, g2=self.e2_disk)
            # disk = disk.shear(disk_shape)
            # bulge = galsim.Sersic(n=self.bulge_sersic_idx, half_light_radius=self.hlr_bulge, flux=1.0, gsparams=gsp)
            # bulge_shape = galsim.Shear(g1=self.e1_bulge, g2=self.e2_bulge)
            # bulge = bulge.shear(bulge_shape)
            
            gal_temp = self.bfrac * bulge + (1.0 - self.bfrac) * disk
Fang Yuedong's avatar
Fang Yuedong committed
178

179
180
181
182
183
184
185
186
            gal_temp = gal_temp.withFlux(nphotons)
            if not big_galaxy: # Not apply PSF for very big galaxy
                gal_temp = galsim.Convolve(psf, gal_temp)
            
            if i == 0:
                gal = gal_temp
            else:
                gal = gal + gal_temp
Fang Yuedong's avatar
Fang Yuedong committed
187

Fang Yuedong's avatar
Fang Yuedong committed
188
            # (TEST) Random knots
189
190
191
            # knots = galsim.RandomKnots(npoints=100, profile=disk)
            # kfrac = np.random.random()*(1.0 - self.bfrac)
            # gal = self.bfrac * bulge + (1.0 - self.bfrac - kfrac) * disk + kfrac * knots
Fang Yuedong's avatar
Fang Yuedong committed
192

193
194
195
            # gal = gal.withFlux(nphotons)
            # gal_shear = galsim.Shear(g1=g1, g2=g2)
            # gal = gal.shear(gal_shear)
Fang Yuedong's avatar
Fang Yuedong committed
196

197
198
199
200
            # if not big_galaxy: # Not apply PSF for very big galaxy
            #     gal = galsim.Convolve(psf, gal)
            #     # if fd_shear is not None:
            #     #     gal = gal.shear(fd_shear)
Fang Yuedong's avatar
Fang Yuedong committed
201

202
203
            # # Use (explicit) stamps to draw
            # stamp = gal.drawImage(wcs=real_wcs_local, method='phot', offset=offset, save_photons=True)
Fang Yuedong's avatar
Fang Yuedong committed
204
            
205
206
            # xmax = max(xmax, stamp.xmax - stamp.xmin)
            # ymax = max(ymax, stamp.ymax - stamp.ymin)
Fang Yuedong's avatar
Fang Yuedong committed
207

208
209
210
211
212
            # photons = stamp.photons
            # photons.x += x_nominal
            # photons.y += y_nominal
            # photons_list.append(photons)
            # del gal
Fang Yuedong's avatar
Fang Yuedong committed
213

Fang Yuedong's avatar
Fang Yuedong committed
214
        # # [C6 TEST]
Fang Yuedong's avatar
Fang Yuedong committed
215
        # print('xmax = %d, ymax = %d '%(xmax, ymax))
Fang Yuedong's avatar
Fang Yuedong committed
216
217
218
219
220
        # # Output memory usage
        # snapshot = tracemalloc.take_snapshot()
        # top_stats = snapshot.statistics('lineno')
        # for stat in top_stats[:10]:
        #     print(stat)
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
        if fd_shear:
            g1 += fd_shear.g1
            g2 += fd_shear.g2
        gal_shear = galsim.Shear(g1=g1, g2=g2)
        gal = gal.shear(gal_shear)
        # if fd_shear is not None:
        #     gal = gal.shear(fd_shear)

        stamp = gal.drawImage(wcs=real_wcs_local, method='phot', offset=offset, save_photons=True)
        photons = stamp.photons
        photons.x += x_nominal
        photons.y += y_nominal
        photons_list.append(photons)

        # stamp = galsim.ImageF(int(xmax * 1.1), int(ymax * 1.1))
xin's avatar
xin committed
236
237
238
        stamp.wcs = real_wcs_local
        stamp.setCenter(x_nominal, y_nominal)
        bounds = stamp.bounds & galsim.BoundsI(0, chip.npix_x - 1, 0, chip.npix_y - 1)
Fang Yuedong's avatar
Fang Yuedong committed
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261

        if bounds.area() > 0:
            chip.img.setOrigin(0, 0)
            stamp[bounds] = chip.img[bounds]

            if not big_galaxy:
                for i in range(len(photons_list)):
                    if i == 0:
                        chip.sensor.accumulate(photons_list[i], stamp)
                    else:
                        chip.sensor.accumulate(photons_list[i], stamp, resume=True)
            else:
                sensor = galsim.Sensor()
                for i in range(len(photons_list)):
                    if i == 0:
                        sensor.accumulate(photons_list[i], stamp)
                    else:
                        sensor.accumulate(photons_list[i], stamp, resume=True)
                del sensor

            chip.img[bounds] = stamp[bounds]

            chip.img.setOrigin(chip.bound.xmin, chip.bound.ymin)
Fang Yuedong's avatar
Fang Yuedong committed
262
        else:
Fang Yuedong's avatar
Fang Yuedong committed
263
264
            # Return code 0: object photons missed this detector
            print("obj %s missed"%(self.id))
265
266
            if self.logger:
                self.logger.info("obj %s missed"%(self.id))
Fang Yuedong's avatar
Fang Yuedong committed
267
268
269
270
            return 0, pos_shear
        
        # # [C6 TEST]
        # print("nphotons_sum = ", nphotons_sum)
Fang Yuedong's avatar
Fang Yuedong committed
271
272
        del photons_list
        del stamp
Fang Yuedong's avatar
Fang Yuedong committed
273
        return 1, pos_shear
Fang Yuedong's avatar
Fang Yuedong committed
274
275

    def drawObj_slitless(self, tel, pos_img, psf_model, bandpass_list, filt, chip, nphotons_tot=None, g1=0, g2=0,
Fang Yuedong's avatar
Fang Yuedong committed
276
277
278
279
280
281
282
283
284
285
286
                         exptime=150., normFilter=None, grating_split_pos=3685, fd_shear=None):
        if normFilter is not None:
            norm_thr_rang_ids = normFilter['SENSITIVITY'] > 0.001
            sedNormFactor = getNormFactorForSpecWithABMAG(ABMag=self.param['mag_use_normal'], spectrum=self.sed,
                                                        norm_thr=normFilter,
                                                        sWave=np.floor(normFilter[norm_thr_rang_ids][0][0]),
                                                        eWave=np.ceil(normFilter[norm_thr_rang_ids][-1][0]))
            if sedNormFactor == 0:
                return 2, None
        else:
            sedNormFactor = 1.
Fang Yuedong's avatar
Fang Yuedong committed
287
288
289
        normalSED = Table(np.array([self.sed['WAVELENGTH'], self.sed['FLUX'] * sedNormFactor]).T,
                          names=('WAVELENGTH', 'FLUX'))

xin's avatar
xin committed
290
291
292
293
294
295
296
297
298
299
300
301
302
        self.real_pos = self.getRealPos(chip.img, global_x=self.posImg.x, global_y=self.posImg.y,
                                        img_real_wcs=self.real_wcs)

        x, y = self.real_pos.x + 0.5, self.real_pos.y + 0.5
        x_nominal = int(np.floor(x + 0.5))
        y_nominal = int(np.floor(y + 0.5))
        dx = x - x_nominal
        dy = y - y_nominal
        offset = galsim.PositionD(dx, dy)

        real_wcs_local = self.real_wcs.local(self.real_pos)


Fang Yuedong's avatar
Fang Yuedong committed
303
        big_galaxy = False
Fang Yuedong's avatar
Fang Yuedong committed
304
        if self.hlr_disk > 3.0 or self.hlr_bulge > 3.0: # Very big galaxy
Fang Yuedong's avatar
Fang Yuedong committed
305
306
307
308
309
310
311
312
            big_galaxy = True

        if self.getMagFilter(filt) <= 15 and (not big_galaxy):
            folding_threshold = 5.e-4
        else:
            folding_threshold = 5.e-3
        gsp = galsim.GSParams(folding_threshold=folding_threshold)
        # nphotons_sum = 0
313
314
315

        flat_cube = chip.flat_cube

Zhang Xin's avatar
Zhang Xin committed
316
        xOrderSigPlus = {'A':1.3909419820029296,'B':1.4760376591236062,'C':4.035447379743442,'D':5.5684364343742825,'E':16.260021029735388}
xin's avatar
xin committed
317
        grating_split_pos_chip = 0 + grating_split_pos
Fang Yuedong's avatar
Fang Yuedong committed
318
319
320
321
        for i in range(len(bandpass_list)):
            bandpass = bandpass_list[i]

            psf, pos_shear = psf_model.get_PSF(chip=chip, pos_img=pos_img, bandpass=bandpass, folding_threshold=folding_threshold)
Fang Yuedong's avatar
Fang Yuedong committed
322
            disk = galsim.Sersic(n=self.disk_sersic_idx, half_light_radius=self.hlr_disk, flux=1.0, gsparams=gsp)
Fang Yuedong's avatar
Fang Yuedong committed
323
324
            disk_shape = galsim.Shear(g1=self.e1_disk, g2=self.e2_disk)
            disk = disk.shear(disk_shape)
Fang Yuedong's avatar
Fang Yuedong committed
325
            bulge = galsim.Sersic(n=self.bulge_sersic_idx, half_light_radius=self.hlr_bulge, flux=1.0, gsparams=gsp)
Fang Yuedong's avatar
Fang Yuedong committed
326
327
328
            bulge_shape = galsim.Shear(g1=self.e1_bulge, g2=self.e2_bulge)
            bulge = bulge.shear(bulge_shape)

329
            gal = self.bfrac * bulge + (1.0 - self.bfrac) * disk
Fang Yuedong's avatar
Fang Yuedong committed
330
331

            # (TEST) Random knots
332
333
334
            # knots = galsim.RandomKnots(npoints=100, profile=disk)
            # kfrac = np.random.random()*(1.0 - self.bfrac)
            # gal = self.bfrac * bulge + (1.0 - self.bfrac - kfrac) * disk + kfrac * knots
Fang Yuedong's avatar
Fang Yuedong committed
335

Fang Yuedong's avatar
Fang Yuedong committed
336
            gal = gal.withFlux(tel.pupil_area * exptime)
337
338
339
            if fd_shear:
                g1 += fd_shear.g1
                g2 += fd_shear.g2
Fang Yuedong's avatar
Fang Yuedong committed
340
341
342
343
            gal_shear = galsim.Shear(g1=g1, g2=g2)
            gal = gal.shear(gal_shear)
            gal = galsim.Convolve(psf, gal)

Fang Yuedong's avatar
Fang Yuedong committed
344
345
            if not big_galaxy: # Not apply PSF for very big galaxy
                gal = galsim.Convolve(psf, gal)
346
347
                # if fd_shear is not None:
                #     gal = gal.shear(fd_shear)
Fang Yuedong's avatar
Fang Yuedong committed
348
349

            starImg = gal.drawImage(wcs=real_wcs_local, offset=offset)
Fang Yuedong's avatar
Fang Yuedong committed
350

xin's avatar
xin committed
351
352
            origin_star = [y_nominal - (starImg.center.y - starImg.ymin),
                           x_nominal - (starImg.center.x - starImg.xmin)]
Fang Yuedong's avatar
Fang Yuedong committed
353
            starImg.setOrigin(0, 0)
Xin Zhang's avatar
Xin Zhang committed
354
355
356
357
358
359
360
361
362
            gal_origin = [origin_star[0], origin_star[1]]
            gal_end = [origin_star[0] + starImg.array.shape[0] - 1, origin_star[1] + starImg.array.shape[1] - 1]

            if gal_origin[1] < grating_split_pos_chip < gal_end[1]:
                subSlitPos = int(grating_split_pos_chip - gal_origin[1] + 1)
                ## part img disperse

                subImg_p1 = starImg.array[:, 0:subSlitPos]
                star_p1 = galsim.Image(subImg_p1)
Fang Yuedong's avatar
Fang Yuedong committed
363
                star_p1.setOrigin(0, 0)
Xin Zhang's avatar
Xin Zhang committed
364
                origin_p1 = origin_star
xin's avatar
xin committed
365
366
                xcenter_p1 = min(x_nominal,grating_split_pos_chip-1) - 0
                ycenter_p1 = y_nominal-0
Xin Zhang's avatar
Xin Zhang committed
367
368
369
370
371
372

                sdp_p1 = SpecDisperser(orig_img=star_p1, xcenter=xcenter_p1,
                                    ycenter=ycenter_p1, origin=origin_p1,
                                    tar_spec=normalSED,
                                    band_start=bandpass.blue_limit * 10, band_end=bandpass.red_limit * 10,
                                    conf=chip.sls_conf[0],
373
374
                                    isAlongY=0,
                                    flat_cube=flat_cube)
Xin Zhang's avatar
Xin Zhang committed
375

xin's avatar
xin committed
376
                self.addSLStoChipImage(sdp=sdp_p1, chip=chip, xOrderSigPlus = xOrderSigPlus, local_wcs=real_wcs_local)
Xin Zhang's avatar
Xin Zhang committed
377
378
379

                subImg_p2 = starImg.array[:, subSlitPos+1:starImg.array.shape[1]]
                star_p2 = galsim.Image(subImg_p2)
Fang Yuedong's avatar
Fang Yuedong committed
380
                star_p2.setOrigin(0, 0)
Xin Zhang's avatar
Xin Zhang committed
381
                origin_p2 = [origin_star[0], grating_split_pos_chip]
xin's avatar
xin committed
382
383
                xcenter_p2 = max(x_nominal, grating_split_pos_chip - 1) - 0
                ycenter_p2 = y_nominal - 0
Xin Zhang's avatar
Xin Zhang committed
384
385
386
387
388
389

                sdp_p2 = SpecDisperser(orig_img=star_p2, xcenter=xcenter_p2,
                                       ycenter=ycenter_p2, origin=origin_p2,
                                       tar_spec=normalSED,
                                       band_start=bandpass.blue_limit * 10, band_end=bandpass.red_limit * 10,
                                       conf=chip.sls_conf[1],
390
391
                                       isAlongY=0,
                                       flat_cube=flat_cube)
Xin Zhang's avatar
Xin Zhang committed
392

xin's avatar
xin committed
393
                self.addSLStoChipImage(sdp=sdp_p2, chip=chip, xOrderSigPlus = xOrderSigPlus, local_wcs=real_wcs_local)
Xin Zhang's avatar
Xin Zhang committed
394
395
396
397

                del sdp_p1
                del sdp_p2
            elif grating_split_pos_chip<=gal_origin[1]:
xin's avatar
xin committed
398
399
                sdp = SpecDisperser(orig_img=starImg, xcenter=x_nominal - 0,
                                    ycenter=y_nominal - 0, origin=origin_star,
Xin Zhang's avatar
Xin Zhang committed
400
401
402
                                    tar_spec=normalSED,
                                    band_start=bandpass.blue_limit * 10, band_end=bandpass.red_limit * 10,
                                    conf=chip.sls_conf[1],
403
404
                                    isAlongY=0,
                                    flat_cube=flat_cube)
xin's avatar
xin committed
405
                self.addSLStoChipImage(sdp=sdp, chip=chip, xOrderSigPlus = xOrderSigPlus, local_wcs=real_wcs_local)
Xin Zhang's avatar
Xin Zhang committed
406
407
                del sdp
            elif grating_split_pos_chip>=gal_end[1]:
xin's avatar
xin committed
408
409
                sdp = SpecDisperser(orig_img=starImg, xcenter=x_nominal - 0,
                                    ycenter=y_nominal - 0, origin=origin_star,
Xin Zhang's avatar
Xin Zhang committed
410
411
412
                                    tar_spec=normalSED,
                                    band_start=bandpass.blue_limit * 10, band_end=bandpass.red_limit * 10,
                                    conf=chip.sls_conf[0],
413
414
                                    isAlongY=0,
                                    flat_cube=flat_cube)
xin's avatar
xin committed
415
                self.addSLStoChipImage(sdp=sdp, chip=chip, xOrderSigPlus = xOrderSigPlus, local_wcs=real_wcs_local)
Xin Zhang's avatar
Xin Zhang committed
416
                del sdp
Fang Yuedong's avatar
Fang Yuedong committed
417

Xin Zhang's avatar
Xin Zhang committed
418
            # print(self.y_nominal, starImg.center.y, starImg.ymin)
Fang Yuedong's avatar
Fang Yuedong committed
419
            del psf
Fang Yuedong's avatar
Fang Yuedong committed
420
        return 1, pos_shear
Fang Yuedong's avatar
Fang Yuedong committed
421
422
423
424

    def getGSObj(self, psf, g1=0, g2=0, flux=None, filt=None, tel=None, exptime=150.):
        if flux == None:
            flux = self.getElectronFluxFilt(filt, tel, exptime)
Fang Yuedong's avatar
Fang Yuedong committed
425
        disk = galsim.Sersic(n=self.disk_sersic_idx, half_light_radius=self.hlr_disk, flux=1.0)
Fang Yuedong's avatar
Fang Yuedong committed
426
427
428
        disk_shape = galsim.Shear(g1=self.e1_disk, g2=self.e2_disk)
        disk = disk.shear(disk_shape)

Fang Yuedong's avatar
Fang Yuedong committed
429
        bulge = galsim.Sersic(n=self.bulge_sersic_idx, half_light_radius=self.hlr_bulge, flux=1.0)
Fang Yuedong's avatar
Fang Yuedong committed
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
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
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
        bulge_shape = galsim.Shear(g1=self.e1_bulge, g2=self.e2_bulge)
        bulge = bulge.shear(bulge_shape)

        gal = self.bfrac * bulge + (1.0 - self.bfrac) * disk
        gal = gal.withFlux(flux)
        gal_shear = galsim.Shear(g1=g1, g2=g2)
        gal = gal.shear(gal_shear)
        final = galsim.Convolve(psf, gal)
        return final

    def rotateEllipticity(self, rotation):
        if rotation == 1:
            self.e1_disk, self.e2_disk, self.e1_bulge, self.e2_bulge, self.e1_total, self.e2_total = -self.e2_disk, self.e1_disk, -self.e2_bulge, self.e1_bulge, -self.e2_total, self.e1_total
        if rotation == 2:
            self.e1_disk, self.e2_disk, self.e1_bulge, self.e2_bulge, self.e1_total, self.e2_total = -self.e1_disk, -self.e2_disk, -self.e1_bulge, -self.e2_bulge, -self.e1_total, -self.e2_total
        if rotation == 3:
            self.e1_disk, self.e2_disk, self.e1_bulge, self.e2_bulge, self.e1_total, self.e2_total = self.e2_disk, -self.e1_disk, self.e2_bulge, -self.e1_bulge, self.e2_total, -self.e1_total

    def drawObject(self, img, final, noise_level=0.0, flux=None, filt=None, tel=None, exptime=150.):
        """ Override the method in parent class 
        Need to constrain the size of image stamp for extended objects
        """
        isUpdated = True
        if flux == None:
            flux = self.getElectronFluxFilt(filt, tel, exptime)
        stamp = final.drawImage(wcs=self.localWCS, offset=self.offset)
        stamp_arr = stamp.array
        mask = (stamp_arr >= 0.001*noise_level) # why 0.001?
        err = int(np.sqrt(mask.sum()))
        if np.mod(err, 2) == 1:
            err += 1
        # if err == 1:
        if err == 0:
            subSize = 16 # why 16?
        else:
            subSize = max([err, 16])
            fluxRatio = flux / stamp_arr[mask].sum()
            final = final.withScaledFlux(fluxRatio)

        imgSub = galsim.ImageF(subSize, subSize)

        # Draw with FFT
        # stamp = final.drawImage(image=imgSub, wcs=self.localWCS, offset=self.offset)

        # Draw with Photon Shoot
        stamp = final.drawImage(image=imgSub, wcs=self.localWCS, method='phot', offset=self.offset)
        
        stamp.setCenter(self.x_nominal, self.y_nominal)
        if np.sum(np.isnan(stamp.array)) >= 1:
            stamp.setZero()
        bounds = stamp.bounds & img.bounds
        if bounds.area() == 0:
            isUpdated = False
        else:
            img[bounds] += stamp[bounds]
        return img, stamp, isUpdated


    def getObservedEll(self, g1=0, g2=0):
        e1_obs, e2_obs, e_obs, theta = eObs(self.e1_total, self.e2_total, g1, g2)
Zhang Xin's avatar
Zhang Xin committed
490
        return self.e1_total, self.e2_total, g1, g2, e1_obs, e2_obs