Merge remote-tracking branch 'origin/develop'

ObservationSim/MockObject/CatalogBase.py → observation_sim/mock_objects/CatalogBase.py

@@ -5,7 +5,7 @@ import cmath
 
 from astropy.table import Table
 from abc import abstractmethod, ABCMeta
-from ObservationSim.MockObject._util import integrate_sed_bandpass, getNormFactorForSpecWithABMAG, getABMAG
+from observation_sim.mock_objects._util import integrate_sed_bandpass, getNormFactorForSpecWithABMAG, getABMAG
 
 
 class CatalogBase(metaclass=ABCMeta):

ObservationSim/MockObject/FlatLED.py → observation_sim/mock_objects/FlatLED.py

 
 
 import galsim
-import os, sys
+import os
+import sys
 import numpy as np
 import time
 import math
@@ -9,12 +10,11 @@ import astropy.constants as cons
 from astropy.io import fits
 from scipy.interpolate import griddata
 from astropy.table import Table
-from ObservationSim.MockObject.SpecDisperser import SpecDisperser
+from observation_sim.mock_objects.SpecDisperser import SpecDisperser
 from scipy import interpolate
 import gc
 
-from ObservationSim.MockObject.MockObject import MockObject
-# from ObservationSim.Straylight import calculateSkyMap_split_g
+from observation_sim.mock_objects.MockObject import MockObject
 
 try:
     import importlib.resources as pkg_resources
@@ -27,13 +27,13 @@ except ImportError:
 LED_name = ['LED1', 'LED2', 'LED3', 'LED4', 'LED5', 'LED6', 'LED7', 'LED8', 'LED9', 'LED10', 'LED11', 'LED12', 'LED13',
             'LED14']
 cwaves_name = {'LED1': '275', 'LED2': '310', 'LED3': '430', 'LED4': '505', 'LED5': '545', 'LED6': '590', 'LED7': '670',
-          'LED8': '760', 'LED9': '940', 'LED10': '940', 'LED11': '1050', 'LED12': '1050',
-          'LED13': '340', 'LED14': '365'}
+               'LED8': '760', 'LED9': '940', 'LED10': '940', 'LED11': '1050', 'LED12': '1050',
+               'LED13': '340', 'LED14': '365'}
 
 cwaves = {'LED1': 2750, 'LED2': 3100, 'LED3': 4300, 'LED4': 5050, 'LED5': 5250, 'LED6': 5900, 'LED7': 6700,
           'LED8': 7600, 'LED9': 8800, 'LED10': 9400, 'LED11': 10500, 'LED12': 15500, 'LED13': 3400, 'LED14': 3650}
 cwaves_fwhm = {'LED1': 110, 'LED2': 120, 'LED3': 200, 'LED4': 300, 'LED5': 300, 'LED6': 130, 'LED7': 210,
-          'LED8': 260, 'LED9': 400, 'LED10': 370, 'LED11': 500, 'LED12': 1400, 'LED13': 90, 'LED14': 100}
+               'LED8': 260, 'LED9': 400, 'LED10': 370, 'LED11': 500, 'LED12': 1400, 'LED13': 90, 'LED14': 100}
 # LED_QE = {'LED1': 0.3, 'LED2': 0.4, 'LED13': 0.5, 'LED14': 0.5, 'LED10': 0.4}
 # e-/ms
 # fluxLED = {'LED1': 0.16478729, 'LED2': 0.084220931, 'LED3': 2.263360617, 'LED4': 2.190623489, 'LED5': 0.703504768,
@@ -43,17 +43,18 @@ cwaves_fwhm = {'LED1': 110, 'LED2': 120, 'LED3': 200, 'LED4': 300, 'LED5': 300,
 
 # e-/ms
 fluxLED = {'LED1': 15, 'LED2': 15, 'LED3': 12.5, 'LED4': 9, 'LED5': 9,
-           'LED6': 9, 'LED7': 9, 'LED8': 9, 'LED9': 9, 'LED10': 12.5, 'LED11': 15, 'LED12':15, 'LED13': 12.5,
+           'LED6': 9, 'LED7': 9, 'LED8': 9, 'LED9': 9, 'LED10': 12.5, 'LED11': 15, 'LED12': 15, 'LED13': 12.5,
            'LED14': 12.5}
 # fluxLEDL = {'LED1': 10, 'LED2': 10, 'LED3': 10, 'LED4': 10, 'LED5': 10,
 #            'LED6': 10, 'LED7': 10, 'LED8': 10, 'LED9': 10, 'LED10': 10, 'LED11': 10, 'LED12':10, 'LED13': 10,
 #            'LED14': 10}
 
-mirro_eff = {'GU':0.61, 'GV':0.8, 'GI':0.8}
+mirro_eff = {'GU': 0.61, 'GV': 0.8, 'GI': 0.8}
 # mirro_eff = {'GU':1, 'GV':1, 'GI':1}
 
+
 class FlatLED(MockObject):
-    def __init__(self, chip,filt, flatDir = None, logger=None):
+    def __init__(self, chip, filt, flatDir=None, logger=None):
         # self.led_type_list = led_type_list
         self.filt = filt
         self.chip = chip
@@ -62,18 +63,19 @@ class FlatLED(MockObject):
             self.flatDir = flatDir
         else:
             try:
-                with pkg_resources.files('ObservationSim.MockObject.data.led').joinpath("") as ledDir:
+                with pkg_resources.files('observation_sim.mock_objects.data.led').joinpath("") as ledDir:
                     self.flatDir = ledDir.as_posix()
             except AttributeError:
-                with pkg_resources.path('ObservationSim.MockObject.data.led', "") as ledDir:
+                with pkg_resources.path('observation_sim.mock_objects.data.led', "") as ledDir:
                     self.flatDir = ledDir.as_posix()
 
     ###
-    ### return LED flat, e/s
+    # return LED flat, e/s
     ###
     def getLEDImage(self, led_type='LED1'):
         # cwave = cwaves[led_type]
-        flat = fits.open(os.path.join(self.flatDir, 'model_' + cwaves_name[led_type] + 'nm.fits'))
+        flat = fits.open(os.path.join(self.flatDir, 'model_' +
+                         cwaves_name[led_type] + 'nm.fits'))
         xlen = flat[0].header['NAXIS1']
         ylen = 601
         x = np.linspace(0, self.chip.npix_x * 6, xlen)
@@ -95,21 +97,23 @@ class FlatLED(MockObject):
         i = self.chip.rowID - 1
         j = self.chip.colID - 1
         U = griddata(X_, Z_, (
-            M[self.chip.npix_y * i:self.chip.npix_y * (i + 1), self.chip.npix_x * j:self.chip.npix_x * (j + 1)],
+            M[self.chip.npix_y * i:self.chip.npix_y *
+                (i + 1), self.chip.npix_x * j:self.chip.npix_x * (j + 1)],
             N[self.chip.npix_y * i:self.chip.npix_y * (i + 1), self.chip.npix_x * j:self.chip.npix_x * (j + 1)]),
-                     method='linear')
+            method='linear')
         U = U/np.mean(U)
         flatImage = U*fluxLED[led_type]*1000
         gc.collect()
         return flatImage
 
-
         ###
-    ### return LED flat, e/s
+    # return LED flat, e/s
     ###
+
     def getLEDImage1(self, led_type='LED1'):
         # cwave = cwaves[led_type]
-        flat = fits.open(os.path.join(self.flatDir, 'model_' + cwaves_name[led_type] + 'nm.fits'))
+        flat = fits.open(os.path.join(self.flatDir, 'model_' +
+                         cwaves_name[led_type] + 'nm.fits'))
         xlen = flat[0].header['NAXIS1']
         ylen = 601
 
@@ -120,22 +124,23 @@ class FlatLED(MockObject):
         y = np.linspace(0, self.chip.npix_y, int(ylen/5.))
         xx, yy = np.meshgrid(x, y)
 
-        a1 = flat[0].data[int(ylen*i/5.):int(ylen*i/5.)+int(ylen/5.), int(xlen*j/6.):int(xlen*j/6.)+int(xlen/6.)]
+        a1 = flat[0].data[int(ylen*i/5.):int(ylen*i/5.)+int(ylen/5.),
+                          int(xlen*j/6.):int(xlen*j/6.)+int(xlen/6.)]
         # z = np.sin((xx+yy+xx**2+yy**2))
         # fInterp = interp2d(xx, yy, z, kind='linear')
 
         X_ = np.hstack((xx.flatten()[:, None], yy.flatten()[:, None]))
         Z_ = a1.flatten()
 
-        n_x = np.arange(0, self.chip.npix_x , 1)
+        n_x = np.arange(0, self.chip.npix_x, 1)
         n_y = np.arange(0, self.chip.npix_y, 1)
 
         M, N = np.meshgrid(n_x, n_y)
 
         U = griddata(X_, Z_, (
             M[0:self.chip.npix_y, 0:self.chip.npix_x],
-            N[0:self.chip.npix_y, 0:self.chip.npix_x ]),
-                     method='linear')
+            N[0:self.chip.npix_y, 0:self.chip.npix_x]),
+            method='linear')
         U = U/np.mean(U)
         flatImage = U*fluxLED[led_type]*1000
         gc.collect()
@@ -143,12 +148,13 @@ class FlatLED(MockObject):
 
     def drawObj_LEDFlat_img(self, led_type_list=['LED1'], exp_t_list=[0.1]):
         if len(led_type_list) > len(exp_t_list):
-            return np.ones([self.chip.npix_y,self.chip.npix_x])
+            return np.ones([self.chip.npix_y, self.chip.npix_x])
 
-        ledFlat = np.zeros([self.chip.npix_y,self.chip.npix_x])
+        ledFlat = np.zeros([self.chip.npix_y, self.chip.npix_x])
 
         ledStat = '00000000000000'
-        ledTimes = [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]
+        ledTimes = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
+                    0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
 
         nledStat = '2'
         for i in np.arange(len(led_type_list)):
@@ -160,8 +166,10 @@ class FlatLED(MockObject):
             led_wave = cwaves[led_type]
             led_fwhm = cwaves_fwhm[led_type]
             led_spec = self.gaussian1d_profile_led(led_wave, led_fwhm)
-            speci = interpolate.interp1d(led_spec['WAVELENGTH'], led_spec['FLUX'])
-            w_list = np.arange(self.filt.blue_limit, self.filt.red_limit, 0.5) #A
+            speci = interpolate.interp1d(
+                led_spec['WAVELENGTH'], led_spec['FLUX'])
+            w_list = np.arange(self.filt.blue_limit,
+                               self.filt.red_limit, 0.5)  # A
 
             f_spec = speci(w_list)
             ccd_bp = self.chip._getChipEffCurve(self.chip.filter_type)
@@ -174,19 +182,21 @@ class FlatLED(MockObject):
             # print("DEBUG1:---------------",np.mean(unitFlatImg))
             ledFlat = ledFlat+unitFlatImg*exp_t
 
-            ledStat = ledStat[0:int(led_type[3:])-1]+nledStat+ledStat[int(led_type[3:]):]
+            ledStat = ledStat[0:int(led_type[3:])-1] + \
+                nledStat+ledStat[int(led_type[3:]):]
             ledTimes[int(led_type[3:])-1] = exp_t * 1000
             gc.collect()
         return ledFlat, ledStat, ledTimes
 
     def drawObj_LEDFlat_slitless(self, led_type_list=['LED1'], exp_t_list=[0.1]):
         if len(led_type_list) != len(exp_t_list):
-            return np.ones([self.chip.npix_y,self.chip.npix_x])
+            return np.ones([self.chip.npix_y, self.chip.npix_x])
 
-        ledFlat = np.zeros([self.chip.npix_y,self.chip.npix_x])
+        ledFlat = np.zeros([self.chip.npix_y, self.chip.npix_x])
 
         ledStat = '00000000000000'
-        ledTimes = [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]
+        ledTimes = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
+                    0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
 
         nledStat = '2'
 
@@ -213,7 +223,8 @@ class FlatLED(MockObject):
                 flat_cube=self.chip.flat_cube, led_spec=led_spec)
 
             ledFlat = ledFlat + ledspec_map
-            ledStat = ledStat[0:int(led_type[3:])-1]+nledStat+ledStat[int(led_type[3:]):]
+            ledStat = ledStat[0:int(led_type[3:])-1] + \
+                nledStat+ledStat[int(led_type[3:]):]
             ledTimes[int(led_type[3:])-1] = exp_t * 1000
         return ledFlat, ledStat, ledTimes
 
@@ -223,7 +234,6 @@ class FlatLED(MockObject):
         elif self.chip.survey_type == "spectroscopic":
             return self.drawObj_LEDFlat_slitless(led_type_list=led_type_list, exp_t_list=exp_t_list)
 
-
     def gaussian1d_profile_led(self, xc=5050, fwhm=300):
         sigma = fwhm/2.355
         x_radii = int(5*sigma + 1)
@@ -232,9 +242,10 @@ class FlatLED(MockObject):
         xlist_[1:-1] = xlist
         xlist_[0] = 2000
         xlist_[-1] = 18000
-        ids1 = xlist>xc-fwhm
-        ids2 = xlist[ids1]<xc+fwhm
-        data = np.exp((-(xlist-xc)*(xlist-xc))/(2*sigma*sigma))/(np.sqrt(2*math.pi)*sigma)
+        ids1 = xlist > xc-fwhm
+        ids2 = xlist[ids1] < xc+fwhm
+        data = np.exp((-(xlist-xc)*(xlist-xc))/(2*sigma*sigma)) / \
+            (np.sqrt(2*math.pi)*sigma)
         scale = 1/np.trapz(data[ids1][ids2], xlist[ids1][ids2])
         data_ = np.zeros(len(xlist) + 2)
         data_[1:-1] = data*scale
@@ -242,8 +253,8 @@ class FlatLED(MockObject):
         return Table(np.array([xlist_.astype(np.float32), data_.astype(np.float32)]).T, names=('WAVELENGTH', 'FLUX'))
 
     def calculateLEDSpec(self, skyMap=None, blueLimit=4200, redLimit=6500,
-                                conf=[''], pixelSize=0.074, isAlongY=0,
-                                split_pos=3685, flat_cube=None, led_spec=None):
+                         conf=[''], pixelSize=0.074, isAlongY=0,
+                         split_pos=3685, flat_cube=None, led_spec=None):
 
         conf1 = conf[0]
         conf2 = conf[0]
@@ -297,7 +308,8 @@ class FlatLED(MockObject):
                     sub_x_s = k2
                     sub_x_e = sub_x_end_arr[j]
 
-                    skyImg_sub = galsim.Image(skyImg.array[sub_y_s:sub_y_e, sub_x_s:sub_x_e])
+                    skyImg_sub = galsim.Image(
+                        skyImg.array[sub_y_s:sub_y_e, sub_x_s:sub_x_e])
                     origin_sub = [sub_y_s, sub_x_s]
                     sub_x_center = (sub_x_s + sub_x_e) / 2.
 
@@ -321,11 +333,8 @@ class FlatLED(MockObject):
                             continue
                         fImg[bounds] = fImg[bounds] + ssImg[bounds]
 
-
-
         else:
 
-
             # sdp.compute_spec_orders()
             y_len = skyMap.shape[0]
             x_len = skyMap.shape[1]
@@ -352,7 +361,8 @@ class FlatLED(MockObject):
                     sub_x_e = sub_x_end_arr[j]
                     # print(i,j,sub_y_s, sub_y_e,sub_x_s,sub_x_e)
                     T1 = time.time()
-                    skyImg_sub = galsim.Image(skyImg.array[sub_y_s:sub_y_e, sub_x_s:sub_x_e])
+                    skyImg_sub = galsim.Image(
+                        skyImg.array[sub_y_s:sub_y_e, sub_x_s:sub_x_e])
                     origin_sub = [sub_y_s, sub_x_s]
                     sub_x_center = (sub_x_s + sub_x_e) / 2.
 
@@ -398,7 +408,8 @@ class FlatLED(MockObject):
 
                     T1 = time.time()
 
-                    skyImg_sub = galsim.Image(skyImg.array[sub_y_s:sub_y_e, sub_x_s:sub_x_e])
+                    skyImg_sub = galsim.Image(
+                        skyImg.array[sub_y_s:sub_y_e, sub_x_s:sub_x_e])
                     origin_sub = [sub_y_s, sub_x_s]
                     sub_x_center = (sub_x_s + sub_x_e) / 2.
 
@@ -426,14 +437,11 @@ class FlatLED(MockObject):
                     print('time: %s ms' % ((T2 - T1) * 1000))
 
         if isAlongY == 1:
-            fimg, tmx, tmy = SpecDisperser.rotate90(array_orig=fImg.array, xc=0, yc=0, isClockwise=0)
+            fimg, tmx, tmy = SpecDisperser.rotate90(
+                array_orig=fImg.array, xc=0, yc=0, isClockwise=0)
         else:
             fimg = fImg.array
 
         # fimg = fimg * pixelSize * pixelSize
 
         return fimg
-
-
-
-

ObservationSim/MockObject/Galaxy.py → observation_sim/mock_objects/Galaxy.py

@@ -2,9 +2,9 @@ import numpy as np
 import galsim
 from astropy.table import Table
 
-from ObservationSim.MockObject._util import eObs, integrate_sed_bandpass, getNormFactorForSpecWithABMAG
-from ObservationSim.MockObject.SpecDisperser import SpecDisperser
-from ObservationSim.MockObject.MockObject import MockObject
+from observation_sim.mock_objects._util import eObs, integrate_sed_bandpass, getNormFactorForSpecWithABMAG
+from observation_sim.mock_objects.SpecDisperser import SpecDisperser
+from observation_sim.mock_objects.MockObject import MockObject
 
 # import tracemalloc
 

ObservationSim/MockObject/MockObject.py → observation_sim/mock_objects/MockObject.py

@@ -6,10 +6,10 @@ from astropy import wcs
 from astropy.table import Table
 import astropy.io.fits as fitsio
 
-from ObservationSim.MockObject._util import magToFlux, VC_A, convolveGaussXorders, convolveImg
-from ObservationSim.MockObject._util import integrate_sed_bandpass, getNormFactorForSpecWithABMAG, getObservedSED, \
+from observation_sim.mock_objects._util import magToFlux, VC_A, convolveGaussXorders, convolveImg
+from observation_sim.mock_objects._util import integrate_sed_bandpass, getNormFactorForSpecWithABMAG, getObservedSED, \
     getABMAG
-from ObservationSim.MockObject.SpecDisperser import SpecDisperser
+from observation_sim.mock_objects.SpecDisperser import SpecDisperser
 
 
 class MockObject(object):
@@ -17,21 +17,21 @@ class MockObject(object):
         self.param = param
         for key in self.param:
             setattr(self, key, self.param[key])
-        
+
         if self.param["star"] == 0:
             self.type = "galaxy"
         elif self.param["star"] == 1:
             self.type = "star"
         elif self.param["star"] == 2:
             self.type = "quasar"
-        ###mock_stamp_START
+        # mock_stamp_START
         elif self.param["star"] == 3:
             self.type = "stamp"
-        ###mock_stamp_END
-        ###for calibration
+        # mock_stamp_END
+        # for calibration
         elif self.param["star"] == 4:
             self.type = "calib"
-        ###END
+        # END
 
         self.sed = None
         self.fd_shear = None
@@ -72,19 +72,22 @@ class MockObject(object):
             if verbose:
                 print("\n")
                 print("Before field distortion:\n")
-                print("x = %.2f, y = %.2f\n" % (self.posImg.x, self.posImg.y), flush=True)
-            self.posImg, self.fd_shear = fdmodel.get_distorted(chip=chip, pos_img=self.posImg)
+                print("x = %.2f, y = %.2f\n" %
+                      (self.posImg.x, self.posImg.y), flush=True)
+            self.posImg, self.fd_shear = fdmodel.get_distorted(
+                chip=chip, pos_img=self.posImg)
             if verbose:
                 print("After field distortion:\n")
-                print("x = %.2f, y = %.2f\n" % (self.posImg.x, self.posImg.y), flush=True)
-        
+                print("x = %.2f, y = %.2f\n" %
+                      (self.posImg.x, self.posImg.y), flush=True)
+
         x, y = self.posImg.x + 0.5, self.posImg.y + 0.5
         self.x_nominal = int(np.floor(x + 0.5))
         self.y_nominal = int(np.floor(y + 0.5))
         dx = x - self.x_nominal
         dy = y - self.y_nominal
         self.offset = galsim.PositionD(dx, dy)
-        
+
         # Deal with chip rotation
         if chip_wcs is not None:
             self.chip_wcs = chip_wcs
@@ -108,7 +111,8 @@ class MockObject(object):
         # print("nphotons_tot = ", nphotons_tot)
 
         try:
-            full = integrate_sed_bandpass(sed=self.sed, bandpass=filt.bandpass_full)
+            full = integrate_sed_bandpass(
+                sed=self.sed, bandpass=filt.bandpass_full)
         except Exception as e:
             print(e)
             if self.logger:
@@ -121,7 +125,7 @@ class MockObject(object):
         else:
             folding_threshold = 5.e-3
         gsp = galsim.GSParams(folding_threshold=folding_threshold)
-        
+
         # Get real image position of object (deal with chip rotation w.r.t its center)
         self.real_pos = self.getRealPos(chip.img, global_x=self.posImg.x, global_y=self.posImg.y,
                                         img_real_wcs=self.chip_wcs)
@@ -153,7 +157,7 @@ class MockObject(object):
 
             # nphotons_sum += nphotons
             # print("nphotons_sub-band_%d = %.2f"%(i, nphotons))
-            
+
             # Get PSF model
             psf, pos_shear = psf_model.get_PSF(chip=chip, pos_img=pos_img, bandpass=bandpass,
                                                folding_threshold=folding_threshold)
@@ -166,21 +170,22 @@ class MockObject(object):
             if np.sum(np.isnan(stamp.array)) > 0:
                 continue
             stamp.setCenter(x_nominal, y_nominal)
-            bounds = stamp.bounds & galsim.BoundsI(0, chip.npix_x - 1, 0, chip.npix_y - 1)
+            bounds = stamp.bounds & galsim.BoundsI(
+                0, chip.npix_x - 1, 0, chip.npix_y - 1)
             if bounds.area() > 0:
                 chip.img.setOrigin(0, 0)
                 chip.img[bounds] += stamp[bounds]
                 is_updated = 1
                 chip.img.setOrigin(chip.bound.xmin, chip.bound.ymin)
                 del stamp
-        
+
         if is_updated == 0:
             # Return code 0: object has missed this detector
-            print("obj %s missed"%(self.id))
+            print("obj %s missed" % (self.id))
             if self.logger:
-                self.logger.info("obj %s missed"%(self.id))
+                self.logger.info("obj %s missed" % (self.id))
             return 0, pos_shear
-        return 1, pos_shear # Return code 1: draw sucesss
+        return 1, pos_shear  # Return code 1: draw sucesss
 
     def addSLStoChipImage(self, sdp=None, chip=None, xOrderSigPlus=None, local_wcs=None):
         spec_orders = sdp.compute_spec_orders()
@@ -193,7 +198,8 @@ class MockObject(object):
             nan_ids = np.isnan(img_s)
             if img_s[nan_ids].shape[0] > 0:
                 # img_s[nan_ids] = 0
-                print("DEBUG: before convolveGaussXorders specImg nan num is", img_s[nan_ids].shape[0])
+                print("DEBUG: before convolveGaussXorders specImg nan num is",
+                      img_s[nan_ids].shape[0])
             #########################################################
             img_s, orig_off = convolveGaussXorders(img_s, xOrderSigPlus[k])
             origin_order_x = v[1] - orig_off
@@ -219,7 +225,8 @@ class MockObject(object):
             stamp.wcs = local_wcs
             stamp.setOrigin(origin_order_x, origin_order_y)
 
-            bounds = stamp.bounds & galsim.BoundsI(0, chip.npix_x - 1, 0, chip.npix_y - 1)
+            bounds = stamp.bounds & galsim.BoundsI(
+                0, chip.npix_x - 1, 0, chip.npix_y - 1)
             if bounds.area() == 0:
                 continue
             chip.img.setOrigin(0, 0)
@@ -230,17 +237,18 @@ class MockObject(object):
             del stamp
         del spec_orders
 
-    def addSLStoChipImageWithPSF(self, sdp=None, chip=None, pos_img_local = [1,1], psf_model=None, bandNo = 1, grating_split_pos=3685, local_wcs=None, pos_img=None):
+    def addSLStoChipImageWithPSF(self, sdp=None, chip=None, pos_img_local=[1, 1], psf_model=None, bandNo=1, grating_split_pos=3685, local_wcs=None, pos_img=None):
         spec_orders = sdp.compute_spec_orders()
         for k, v in spec_orders.items():
             img_s = v[0]
             # print(bandNo,k)
             try:
-                psf, pos_shear = psf_model.get_PSF(chip, pos_img_local = pos_img_local, bandNo = bandNo, galsimGSObject=True, g_order = k, grating_split_pos=grating_split_pos)
+                psf, pos_shear = psf_model.get_PSF(
+                    chip, pos_img_local=pos_img_local, bandNo=bandNo, galsimGSObject=True, g_order=k, grating_split_pos=grating_split_pos)
             except:
                 psf, pos_shear = psf_model.get_PSF(chip=chip, pos_img=pos_img)
 
-            psf_img = psf.drawImage(nx=100, ny=100, wcs = local_wcs)
+            psf_img = psf.drawImage(nx=100, ny=100, wcs=local_wcs)
 
             psf_img_m = psf_img.array
 
@@ -263,7 +271,6 @@ class MockObject(object):
             origin_order_x = v[1] - orig_off[0]
             origin_order_y = v[2] - orig_off[1]
 
-
             specImg = galsim.ImageF(img_s)
             # photons = galsim.PhotonArray.makeFromImage(specImg)
             # photons.x += origin_order_x
@@ -278,7 +285,8 @@ class MockObject(object):
             specImg.wcs = local_wcs
             specImg.setOrigin(origin_order_x, origin_order_y)
 
-            bounds = specImg.bounds & galsim.BoundsI(0, chip.npix_x - 1, 0, chip.npix_y - 1)
+            bounds = specImg.bounds & galsim.BoundsI(
+                0, chip.npix_x - 1, 0, chip.npix_y - 1)
             if bounds.area() == 0:
                 continue
             chip.img.setOrigin(0, 0)
@@ -297,7 +305,8 @@ class MockObject(object):
             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]),
+                                                          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
@@ -331,9 +340,9 @@ class MockObject(object):
                          'D': 5.5684364343742825, 'E': 16.260021029735388}
         grating_split_pos_chip = 0 + grating_split_pos
 
-        branges = np.zeros([len(bandpass_list),2])
+        branges = np.zeros([len(bandpass_list), 2])
 
-        if hasattr(psf_model,'bandranges'):
+        if hasattr(psf_model, 'bandranges'):
             if psf_model.bandranges is None:
                 return 2, None
             if len(psf_model.bandranges) != len(bandpass_list):
@@ -354,22 +363,24 @@ class MockObject(object):
             star = star.withFlux(tel.pupil_area * exptime)
             psf_tmp = galsim.Gaussian(sigma=0.002)
             star = galsim.Convolve(psf_tmp, star)
-            
-            starImg = star.drawImage(nx=60, ny=60, wcs=chip_wcs_local, offset=offset)
+
+            starImg = star.drawImage(
+                nx=60, ny=60, wcs=chip_wcs_local, offset=offset)
 
             origin_star = [y_nominal - (starImg.center.y - starImg.ymin),
                            x_nominal - (starImg.center.x - starImg.xmin)]
-            starImg.setOrigin(0,0)
+            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]
+            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
+                # part img disperse
 
                 subImg_p1 = starImg.array[:, 0:subSlitPos]
                 star_p1 = galsim.Image(subImg_p1)
                 origin_p1 = origin_star
-                star_p1.setOrigin(0,0)
+                star_p1.setOrigin(0, 0)
                 xcenter_p1 = min(x_nominal, grating_split_pos_chip - 1) - 0
                 ycenter_p1 = y_nominal - 0
 
@@ -382,11 +393,12 @@ class MockObject(object):
                                        flat_cube=flat_cube)
 
                 # self.addSLStoChipImage(sdp=sdp_p1, chip=chip, xOrderSigPlus=xOrderSigPlus, local_wcs=chip_wcs_local)
-                pos_shear=self.addSLStoChipImageWithPSF(sdp=sdp_p1, chip=chip, pos_img_local = [xcenter_p1,ycenter_p1],
-                                              psf_model=psf_model, bandNo = i+1, grating_split_pos=grating_split_pos,
-                                              local_wcs=chip_wcs_local, pos_img = pos_img)
+                pos_shear = self.addSLStoChipImageWithPSF(sdp=sdp_p1, chip=chip, pos_img_local=[xcenter_p1, ycenter_p1],
+                                                          psf_model=psf_model, bandNo=i+1, grating_split_pos=grating_split_pos,
+                                                          local_wcs=chip_wcs_local, pos_img=pos_img)
 
-                subImg_p2 = starImg.array[:, subSlitPos + 1:starImg.array.shape[1]]
+                subImg_p2 = starImg.array[:,
+                                          subSlitPos + 1:starImg.array.shape[1]]
                 star_p2 = galsim.Image(subImg_p2)
                 star_p2.setOrigin(0, 0)
                 origin_p2 = [origin_star[0], grating_split_pos_chip]
@@ -402,9 +414,9 @@ class MockObject(object):
                                        flat_cube=flat_cube)
 
                 # self.addSLStoChipImage(sdp=sdp_p2, chip=chip, xOrderSigPlus=xOrderSigPlus, local_wcs=chip_wcs_local)
-                pos_shear=self.addSLStoChipImageWithPSF(sdp=sdp_p2, chip=chip, pos_img_local=[xcenter_p2, ycenter_p2],
-                                              psf_model=psf_model, bandNo=i + 1, grating_split_pos=grating_split_pos,
-                                              local_wcs=chip_wcs_local, pos_img = pos_img)
+                pos_shear = self.addSLStoChipImageWithPSF(sdp=sdp_p2, chip=chip, pos_img_local=[xcenter_p2, ycenter_p2],
+                                                          psf_model=psf_model, bandNo=i + 1, grating_split_pos=grating_split_pos,
+                                                          local_wcs=chip_wcs_local, pos_img=pos_img)
 
                 del sdp_p1
                 del sdp_p2
@@ -417,9 +429,9 @@ class MockObject(object):
                                     isAlongY=0,
                                     flat_cube=flat_cube)
                 # self.addSLStoChipImage(sdp=sdp, chip=chip, xOrderSigPlus=xOrderSigPlus, local_wcs=chip_wcs_local)
-                pos_shear=self.addSLStoChipImageWithPSF(sdp=sdp, chip=chip, pos_img_local=[x_nominal, y_nominal],
-                                              psf_model=psf_model, bandNo=i + 1, grating_split_pos=grating_split_pos,
-                                              local_wcs=chip_wcs_local, pos_img = pos_img)
+                pos_shear = self.addSLStoChipImageWithPSF(sdp=sdp, chip=chip, pos_img_local=[x_nominal, y_nominal],
+                                                          psf_model=psf_model, bandNo=i + 1, grating_split_pos=grating_split_pos,
+                                                          local_wcs=chip_wcs_local, pos_img=pos_img)
                 del sdp
             elif grating_split_pos_chip >= gal_end[1]:
                 sdp = SpecDisperser(orig_img=starImg, xcenter=x_nominal - 0,
@@ -430,9 +442,9 @@ class MockObject(object):
                                     isAlongY=0,
                                     flat_cube=flat_cube)
                 # self.addSLStoChipImage(sdp=sdp, chip=chip, xOrderSigPlus=xOrderSigPlus, local_wcs=chip_wcs_local)
-                pos_shear=self.addSLStoChipImageWithPSF(sdp=sdp, chip=chip, pos_img_local=[x_nominal, y_nominal],
-                                              psf_model=psf_model, bandNo=i + 1, grating_split_pos=grating_split_pos,
-                                              local_wcs=chip_wcs_local, pos_img = pos_img)
+                pos_shear = self.addSLStoChipImageWithPSF(sdp=sdp, chip=chip, pos_img_local=[x_nominal, y_nominal],
+                                                          psf_model=psf_model, bandNo=i + 1, grating_split_pos=grating_split_pos,
+                                                          local_wcs=chip_wcs_local, pos_img=pos_img)
                 del sdp
             # del psf
         return 1, pos_shear
@@ -447,16 +459,15 @@ class MockObject(object):
         snr_obj = img_flux / sig_obj
         return snr_obj
 
-
-
     def drawObj_PSF(self, tel, pos_img, psf_model, bandpass_list, filt, chip, nphotons_tot=None, g1=0, g2=0,
-                          exptime=150., fd_shear=None, chip_output=None):
+                    exptime=150., fd_shear=None, chip_output=None):
         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)
+            full = integrate_sed_bandpass(
+                sed=self.sed, bandpass=filt.bandpass_full)
         except Exception as e:
             print(e)
             if self.logger:
@@ -504,21 +515,21 @@ class MockObject(object):
                                                folding_threshold=folding_threshold)
             star_temp = psf.withFlux(nphotons)
 
-            if i==0:
+            if i == 0:
                 star = star_temp
             else:
                 star = star+star_temp
 
         pixelScale = 0.074
         stamp = star.drawImage(wcs=chip_wcs_local, offset=offset)
-        #stamp = star.drawImage(nx=256, ny=256, scale=pixelScale)
+        # stamp = star.drawImage(nx=256, ny=256, scale=pixelScale)
         if np.sum(np.isnan(stamp.array)) > 0:
             return None
 
-
         fn = chip_output.subdir + "/psfIDW"
         os.makedirs(fn, exist_ok=True)
-        fn = fn + "/ccd_{:}".format(chip.chipID)+"_psf_"+str(self.param['id'])+".fits"
+        fn = fn + "/ccd_{:}".format(chip.chipID) + \
+            "_psf_"+str(self.param['id'])+".fits"
         if fn != None:
             if os.path.exists(fn):
                 os.remove(fn)

ObservationSim/MockObject/Quasar.py → observation_sim/mock_objects/Quasar.py

@@ -6,8 +6,8 @@ import astropy.constants as cons
 from astropy.table import Table
 from scipy import interpolate
 
-from ObservationSim.MockObject.MockObject import MockObject
-from ObservationSim.MockObject._util import integrate_sed_bandpass, getNormFactorForSpecWithABMAG, getObservedSED, getABMAG
+from observation_sim.mock_objects.MockObject import MockObject
+from observation_sim.mock_objects._util import integrate_sed_bandpass, getNormFactorForSpecWithABMAG, getObservedSED, getABMAG
 
 
 class Quasar(MockObject):