Newer
Older
import os, sys
import astropy.constants as cons
from astropy.table import Table
from scipy import interpolate
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
committed
# import tracemalloc
Fang Yuedong
committed
def __init__(self, param, rotation=None, logger=None):
super().__init__(param, logger=logger)
# self.thetaR = self.param["theta"]
# self.bfrac = self.param["bfrac"]
# self.hlr_disk = self.param["hlr_disk"]
# self.hlr_bulge = self.param["hlr_bulge"]
# 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
if rotation is not None:
self.rotateEllipticity(rotation)
def unload_SED(self):
"""(Test) free up SED memory
"""
del self.sed
def getGSObj_multiband(self, tel, psf_list, bandpass_list, filt, nphotons_tot=None, g1=0, g2=0, exptime=150.):
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)
self.logger.error(e)
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)
self.logger.error(e)
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
return -1
ratio = sub/full
if not (ratio == -1 or (ratio != ratio)):
nphotons = ratio * nphotons_tot
else:
return -1
psf = psf_list[i]
disk = galsim.Sersic(n=1.0, half_light_radius=self.hlr_disk, flux=1.0)
disk_shape = galsim.Shear(g1=self.e1_disk, g2=self.e2_disk)
disk = disk.shear(disk_shape)
bulge = galsim.Sersic(n=4.0, half_light_radius=self.hlr_bulge, flux=1.0)
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)
gal_shear = galsim.Shear(g1=g1, g2=g2)
gal = gal.shear(gal_shear)
gal = galsim.Convolve(psf, gal)
objs.append(gal)
final = galsim.Sum(objs)
return final
def drawObj_multiband(self, tel, pos_img, psf_model, bandpass_list, filt, chip, nphotons_tot=None, g1=0, g2=0, exptime=150.):
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)
self.logger.error(e)
return False
nphotons_sum = 0
photons_list = []
xmax, ymax = 0, 0
Fang Yuedong
committed
# # [C6 TEST]
# print('hlr_disk = %.4f, hlr_bulge = %.4f'%(self.hlr_disk, self.hlr_bulge))
Fang Yuedong
committed
# tracemalloc.start()
Fang Yuedong
committed
if self.hlr_disk > 3.0 or self.hlr_bulge > 3.0: # Very big galaxy
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)
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)
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)
self.logger.error(e)
# return False
continue
ratio = sub/full
if not (ratio == -1 or (ratio != ratio)):
nphotons = ratio * nphotons_tot
else:
# return False
continue
nphotons_sum += nphotons
Fang Yuedong
committed
# # [C6 TEST]
# 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)
disk = galsim.Sersic(n=1.0, 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=4.0, 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
committed
# (TEST) Random knots
# 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
gal = gal.withFlux(nphotons)
gal_shear = galsim.Shear(g1=g1, g2=g2)
gal = gal.shear(gal_shear)
if self.hlr_disk < 10.0: # Not apply PSF for very big galaxy
gal = galsim.Convolve(psf, gal)
# Use (explicit) stamps to draw
# stamp = gal.drawImage(wcs=self.localWCS, method='phot', offset=self.offset, save_photons=True)
# xmax = max(xmax, stamp.xmax)
# ymax = max(ymax, stamp.ymax)
# photons = stamp.photons
# photons.x += self.x_nominal
# photons.y += self.y_nominal
# photons_list.append(photons)
stamp = gal.drawImage(wcs=real_wcs_local, method='phot', offset=offset, save_photons=True)
Fang Yuedong
committed
xmax = max(xmax, stamp.xmax - stamp.xmin)
ymax = max(ymax, stamp.ymax - stamp.ymin)
Fang Yuedong
committed
# # [C6 TEST]
Fang Yuedong
committed
# # Output memory usage
# snapshot = tracemalloc.take_snapshot()
# top_stats = snapshot.statistics('lineno')
# for stat in top_stats[:10]:
# print(stat)
stamp = galsim.ImageF(int(xmax * 1.1), int(ymax * 1.1))
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)
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)
chip.img.setOrigin(chip.bound.xmin, chip.bound.ymin)
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
# stamp = galsim.ImageF(int(xmax*1.1), int(ymax*1.1))
# stamp.wcs = self.localWCS
# stamp.setCenter(self.x_nominal, self.y_nominal)
# bounds = stamp.bounds & chip.img.bounds
# 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)
#
# # print(stamp.array.sum())
# # chip.img[bounds] += stamp[bounds]
# chip.img[bounds] = stamp[bounds]
Fang Yuedong
committed
# # [C6 TEST]
# print("nphotons_sum = ", nphotons_sum)
return True, pos_shear
def drawObj_slitless(self, tel, pos_img, psf_model, bandpass_list, filt, chip, nphotons_tot=None, g1=0, g2=0,
exptime=150., normFilter=None, grating_split_pos=3685):
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 False
else:
sedNormFactor = 1.
normalSED = Table(np.array([self.sed['WAVELENGTH'], self.sed['FLUX'] * sedNormFactor]).T,
names=('WAVELENGTH', 'FLUX'))
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)
big_galaxy = False
if self.hlr_disk > 3.0: # Very big galaxy
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
flat_cube = chip.flat_cube
xOrderSigPlus = {'A':1.3909419820029296,'B':1.4760376591236062,'C':4.035447379743442,'D':5.5684364343742825,'E':16.260021029735388}
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)
disk = galsim.Sersic(n=1.0, 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=4.0, 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)
# 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
gal = gal.withFlux(tel.pupil_area * exptime)
gal_shear = galsim.Shear(g1=g1, g2=g2)
gal = gal.shear(gal_shear)
gal = galsim.Convolve(psf, gal)
starImg = gal.drawImage(wcs=real_wcs_local, offset=offset)
origin_star = [y_nominal - (starImg.center.y - starImg.ymin),
x_nominal - (starImg.center.x - starImg.xmin)]
starImg.setOrigin(0, 0)
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)
xcenter_p1 = min(x_nominal,grating_split_pos_chip-1) - 0
ycenter_p1 = y_nominal-0
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],
isAlongY=0,
flat_cube=flat_cube)
self.addSLStoChipImage(sdp=sdp_p1, chip=chip, xOrderSigPlus = xOrderSigPlus, local_wcs=real_wcs_local)
subImg_p2 = starImg.array[:, subSlitPos+1:starImg.array.shape[1]]
star_p2 = galsim.Image(subImg_p2)
origin_p2 = [origin_star[0], grating_split_pos_chip]
xcenter_p2 = max(x_nominal, grating_split_pos_chip - 1) - 0
ycenter_p2 = y_nominal - 0
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],
isAlongY=0,
flat_cube=flat_cube)
self.addSLStoChipImage(sdp=sdp_p2, chip=chip, xOrderSigPlus = xOrderSigPlus, local_wcs=real_wcs_local)
del sdp_p1
del sdp_p2
elif grating_split_pos_chip<=gal_origin[1]:
sdp = SpecDisperser(orig_img=starImg, xcenter=x_nominal - 0,
ycenter=y_nominal - 0, origin=origin_star,
tar_spec=normalSED,
band_start=bandpass.blue_limit * 10, band_end=bandpass.red_limit * 10,
conf=chip.sls_conf[1],
isAlongY=0,
flat_cube=flat_cube)
self.addSLStoChipImage(sdp=sdp, chip=chip, xOrderSigPlus = xOrderSigPlus, local_wcs=real_wcs_local)
del sdp
elif grating_split_pos_chip>=gal_end[1]:
sdp = SpecDisperser(orig_img=starImg, xcenter=x_nominal - 0,
ycenter=y_nominal - 0, origin=origin_star,
tar_spec=normalSED,
band_start=bandpass.blue_limit * 10, band_end=bandpass.red_limit * 10,
conf=chip.sls_conf[0],
isAlongY=0,
flat_cube=flat_cube)
self.addSLStoChipImage(sdp=sdp, chip=chip, xOrderSigPlus = xOrderSigPlus, local_wcs=real_wcs_local)
# print(self.y_nominal, starImg.center.y, starImg.ymin)
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
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
del psf
return True, pos_shear
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)
disk = galsim.Sersic(n=1.0, half_light_radius=self.hlr_disk, flux=1.0)
disk_shape = galsim.Shear(g1=self.e1_disk, g2=self.e2_disk)
disk = disk.shear(disk_shape)
bulge = galsim.Sersic(n=4.0, half_light_radius=self.hlr_bulge, flux=1.0)
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)
return self.e1_total, self.e2_total, g1, g2, e1_obs, e2_obs