merge sls

observation_sim/ObservationSim.py

@@ -23,7 +23,7 @@ class Observation(object):
         self.filter_param = FilterParam()
         self.Catalog = Catalog
 
-    def prepare_chip_for_exposure(self, chip, ra_cen, dec_cen, pointing, wcs_fp=None):
+    def prepare_chip_for_exposure(self, chip, ra_cen, dec_cen, pointing, wcs_fp=None, slsPSFOptim = False):
         # Get WCS for the focal plane
         if wcs_fp == None:
             wcs_fp = self.focal_plane.getTanWCS(
@@ -34,6 +34,26 @@ class Observation(object):
         chip.img.setOrigin(chip.bound.xmin, chip.bound.ymin)
         chip.img.wcs = wcs_fp
 
+        chip.slsPSFOptim = slsPSFOptim
+        if chip.chipID in [1,2,3,4,5,10,21,26,27,28,29,30] and slsPSFOptim:
+            chip.img_stack = {}
+            for id1 in np.arange(2):
+                gn = chip_utils.getChipSLSGratingID(chip.chipID)[id1]
+                orders = {}
+                # for id2 in ['-2','-1','0','1','2']:
+                for id2 in ['0','1']:
+                    o_n = "order"+id2
+                    allbands = {}
+                    for id3 in ['1','2','3','4']:
+                        w_n = "w"+id3
+                        allbands[w_n] = galsim.ImageF(chip.npix_x, chip.npix_y)
+                        allbands[w_n].setOrigin(chip.bound.xmin, chip.bound.ymin)
+                        allbands[w_n].wcs = wcs_fp
+                    orders[o_n] = allbands      
+                chip.img_stack[gn] = orders
+        else:
+            chip.img_stack = {}
+
         # Get random generators for this chip
         chip.rng_poisson, chip.poisson_noise = chip_utils.get_poisson(
             seed=int(self.config["random_seeds"]["seed_poisson"]) + pointing.id*30 + chip.chipID, sky_level=0.)
@@ -97,9 +117,10 @@ class Observation(object):
             ra_cen = pointing.ra
             dec_cen = pointing.dec
             ra_offset, dec_offset = 0., 0.
-
+            
+        slsPSFOpt = False
         # Prepare necessary chip properties for simulation
-        chip = self.prepare_chip_for_exposure(chip, ra_cen, dec_cen, pointing)
+        chip = self.prepare_chip_for_exposure(chip, ra_cen, dec_cen, pointing, slsPSFOptim = slsPSFOpt)
 
         # Initialize SimSteps
         sim_steps = SimSteps(overall_config=self.config,

observation_sim/mock_objects/Galaxy.py

@@ -323,28 +323,73 @@ class Galaxy(MockObject):
             #     # if fd_shear is not None:
             #     #     gal = gal.shear(fd_shear)
 
-            starImg = gal.drawImage(
-                wcs=chip_wcs_local, offset=offset, method='real_space')
-
-            origin_star = [y_nominal - (starImg.center.y - starImg.ymin),
-                           x_nominal - (starImg.center.x - starImg.xmin)]
-            starImg.setOrigin(0, 0)
+            galImg_List = []
+            try:
+                pos_img_local = [0,0]
+                x_start = chip.x_cen/chip.pix_size - chip.npix_x / 2.
+                y_start = chip.y_cen/chip.pix_size - chip.npix_y / 2.
+                pos_img_local[0] = pos_img.x - x_start
+                pos_img_local[1] = pos_img.y - y_start
+                nnx = 0
+                nny = 0
+                for order in ["A","B"]:
+                    psf, pos_shear = psf_model.get_PSF(
+                        chip, pos_img_local=pos_img_local, bandNo=i+1, galsimGSObject=True, g_order=order, grating_split_pos=grating_split_pos)
+                    star_p = galsim.Convolve(psf, gal)
+                    if nnx == 0:
+                        galImg = star_p.drawImage(wcs=chip_wcs_local, offset=offset)
+                        nnx = galImg.xmax - galImg.xmin + 1
+                        nny = galImg.ymax - galImg.ymin + 1
+                    else:
+                        galImg = star_p.drawImage(nx = nnx, ny = nny, wcs=chip_wcs_local, offset=offset)
+                    galImg.setOrigin(0, 0)
+                    # n1 = np.sum(np.isinf(galImg.array))
+                    # n2 = np.sum(np.isnan(galImg.array))
+                    # if n1>0 or n2 > 0:
+                    #     print("DEBUG: Galaxy, inf:%d, nan:%d"%(n1, n2))
+                    if np.sum(np.isnan(galImg.array)) > 0:
+                        # ERROR happens
+                        return 2, pos_shear
+                    galImg_List.append(galImg)
+                for order in ["C","D","E"]:
+                    galImg_List.append(galImg)
+            except:
+                psf, pos_shear = psf_model.get_PSF(chip=chip, pos_img=pos_img)
+                star_p = galsim.Convolve(psf, gal)
+                galImg = star_p.drawImage(wcs=chip_wcs_local, offset=offset)
+                galImg.setOrigin(0, 0)
+                if np.sum(np.isnan(galImg.array)) > 0:
+                    # ERROR happens
+                    return 2, pos_shear
+                for order in ["A","B","C","D","E"]:
+                    galImg_List.append(galImg)
+
+            # starImg = gal.drawImage(
+            #     wcs=chip_wcs_local, offset=offset, method='real_space')
+
+            origin_star = [y_nominal - (galImg.center.y - galImg.ymin),
+                           x_nominal - (galImg.center.x - galImg.xmin)]
+            galImg.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] + galImg.array.shape[0] -
+                       1, origin_star[1] + galImg.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)
-                star_p1.setOrigin(0, 0)
+                star_p1s=[]
+                for galImg in galImg_List:
+
+                    subImg_p1 = galImg.array[:, 0:subSlitPos]
+                    star_p1 = galsim.Image(subImg_p1)
+                    star_p1.setOrigin(0, 0)
+                    star_p1s.append(star_p1)
                 origin_p1 = origin_star
                 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,
+                sdp_p1 = SpecDisperser(orig_img=star_p1s, xcenter=xcenter_p1,
                                        ycenter=ycenter_p1, origin=origin_p1,
                                        tar_spec=normalSED,
                                        band_start=brange[0], band_end=brange[1],
@@ -352,21 +397,25 @@ class Galaxy(MockObject):
                                        isAlongY=0,
                                        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)
+                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)
+
+                star_p2s=[]
+                for galImg in galImg_List:
 
-                subImg_p2 = starImg.array[:,
-                                          subSlitPos+1:starImg.array.shape[1]]
-                star_p2 = galsim.Image(subImg_p2)
-                star_p2.setOrigin(0, 0)
+                    subImg_p2 = galImg.array[:,
+                                          subSlitPos + 1:galImg.array.shape[1]]
+                    star_p2 = galsim.Image(subImg_p2)
+                    star_p2.setOrigin(0, 0)
+                    star_p2s.append(star_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,
+                sdp_p2 = SpecDisperser(orig_img=star_p2s, xcenter=xcenter_p2,
                                        ycenter=ycenter_p2, origin=origin_p2,
                                        tar_spec=normalSED,
                                        band_start=brange[0], band_end=brange[1],
@@ -374,41 +423,41 @@ class Galaxy(MockObject):
                                        isAlongY=0,
                                        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)
+                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)
 
                 del sdp_p1
                 del sdp_p2
             elif grating_split_pos_chip <= gal_origin[1]:
-                sdp = SpecDisperser(orig_img=starImg, xcenter=x_nominal - 0,
+                sdp = SpecDisperser(orig_img=galImg_List, xcenter=x_nominal - 0,
                                     ycenter=y_nominal - 0, origin=origin_star,
                                     tar_spec=normalSED,
                                     band_start=brange[0], band_end=brange[1],
                                     conf=chip.sls_conf[1],
                                     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)
+                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)
                 del sdp
             elif grating_split_pos_chip >= gal_end[1]:
-                sdp = SpecDisperser(orig_img=starImg, xcenter=x_nominal - 0,
+                sdp = SpecDisperser(orig_img=galImg_List, xcenter=x_nominal - 0,
                                     ycenter=y_nominal - 0, origin=origin_star,
                                     tar_spec=normalSED,
                                     band_start=brange[0], band_end=brange[1],
                                     conf=chip.sls_conf[0],
                                     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)
+                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)
                 del sdp
 
             # print(self.y_nominal, starImg.center.y, starImg.ymin)

observation_sim/mock_objects/SpecDisperser/SpecDisperser.py

@@ -65,10 +65,30 @@ class SpecDisperser(object):
         # self.img_x = orig_img.shape[1]
         # self.img_y = orig_img.shape[0]
 
-        self.thumb_img = np.abs(orig_img.array)
-        self.thumb_x = orig_img.center.x
-        self.thumb_y = orig_img.center.y
-        self.img_sh = orig_img.array.shape
+        # 5 orders, A, B ,
+        orderName=["A","B","C","D","E"]
+        self.orig_img_orders = OrderedDict()
+        if isinstance(orig_img, list):
+            orig_img_list = orig_img
+            list_len = len(orig_img_list)
+            if list_len < 5:
+                for i in np.arange(5-list_len):
+                    orig_img_list.append(orig_img_list[list_len-1])
+            for i, k in enumerate(orig_img_list):
+                self.orig_img_orders[orderName[i]] = k
+
+        
+        if isinstance(orig_img, galsim.Image):
+            for i in np.arange(5):
+                self.orig_img_orders[orderName[i]] = orig_img
+
+
+        orig_img_one = self.orig_img_orders["A"]
+
+        self.thumb_img = np.abs(orig_img_one.array)
+        self.thumb_x = orig_img_one.center.x
+        self.thumb_y = orig_img_one.center.y
+        self.img_sh = orig_img_one.array.shape
 
         self.id = gid
 
@@ -78,10 +98,13 @@ class SpecDisperser(object):
         self.isAlongY = isAlongY
         self.flat_cube = flat_cube
         if self.isAlongY == 1:
-            self.thumb_img, self.thumb_x, self.thumb_y = rotate90(array_orig=self.thumb_img, xc=orig_img.center.x,
-                                                                  yc=orig_img.center.y, isClockwise=1)
+            for order in orderName:
+                self.orig_img_orders[order], self.thumb_x, self.thumb_y = rotate90(array_orig=self.orig_img_orders[order], xc=orig_img_one.center.x,
+                                                                  yc=orig_img_one.center.y, isClockwise=1)
+            # self.thumb_img, self.thumb_x, self.thumb_y = rotate90(array_orig=self.thumb_img, xc=orig_img_one.center.x,
+            #                                                       yc=orig_img_one.center.y, isClockwise=1)
 
-            self.img_sh = orig_img.array.T.shape
+            self.img_sh = self.orig_img_orders[order].array.T.shape
             self.xcenter = ycenter
             self.ycenter = xcenter
 
@@ -111,10 +134,16 @@ class SpecDisperser(object):
 
     def compute_spec(self, beam):
 
+        # if beam == "B":
+        #     return self.thumb_img, self.origin[1], self.origin[0], None, None, None
+
         from .disperse_c import interp
         from .disperse_c import disperse
         # from MockObject.disperse_c import disperse
-
+        self.thumb_img = np.abs(self.orig_img_orders[beam].array)
+        self.thumb_x = self.orig_img_orders[beam].center.x
+        self.thumb_y = self.orig_img_orders[beam].center.y
+        self.img_sh = self.orig_img_orders[beam].array.shape
         dx = self.grating_conf.dxlam[beam]
         xoff = 0
         ytrace_beam, lam_beam = self.grating_conf.get_beam_trace(x=self.xcenter, y=self.ycenter, dx=(dx + xoff),
@@ -169,7 +198,8 @@ class SpecDisperser(object):
 
         dyc = cast[int](np.floor(ytrace_beam+0.5))
 
-        dypix = cast[int](np.floor(ytrace_beam - dyc[0] + x0[0] + 0.5))
+        # dypix = cast[int](np.floor(ytrace_beam - dyc[0] + x0[0] + 0.5))
+        dypix = dyc - dyc[0] + x0[0]
 
         frac_ids = yfrac_beam < 0
 
@@ -248,7 +278,8 @@ class SpecDisperser(object):
 #                        beam_flat[k] = self.flat_cube[:, originOut_y + i, originOut_x + j]
 
         status = disperse.disperse_grism_object(self.thumb_img.astype(np.float32),
-                                                flat_index[nonz], yfrac_beam[nonz],
+                                                flat_index[nonz],
+                                                yfrac_beam[nonz],
                                                 sensitivity_beam[nonz],
                                                 modelf, x0,
                                                 array(self.img_sh,
@@ -258,11 +289,24 @@ class SpecDisperser(object):
                                                 lam_beam[lam_index][nonz])
 
         model = modelf.reshape(beam_sh)
+        # n1 = np.sum(np.isinf(model))
+        # n2 = np.sum(np.isnan(model))
+        # n3 = np.sum(np.isinf(modelf))
+        # n4 = np.sum(np.isnan(modelf))
+        # if n1>0 or n2 > 0:
+        #     print("DEBUG: SpecDisperser, inf:%d, nan:%d--------%d,%d"%(n1, n2, n3, n4))
+        #     print(dypix)
+            # n1 = np.sum(np.isinf(self.thumb_img.astype(np.float32)))
+            # n2 = np.sum(np.isnan(self.thumb_img.astype(np.float32)))
+            # n3 = np.sum(np.isinf(yfrac_beam))
+            # n4 = np.sum(np.isnan(yfrac_beam))
+            # n5 = np.sum(np.isinf(sensitivity_beam))
+            # n6 = np.sum(np.isnan(sensitivity_beam))
+            # print("DEBUG: SpecDisperser, innput ---inf:%d, nan:%d, yfrac_beam:%d/%d, sensitivity_beam:%d/%d"%(n1, n2, n3, n4, n5, n6))
         self.beam_flux[beam] = sum(modelf)
 
         if self.isAlongY == 1:
             model, _, _ = rotate90(array_orig=model, isClockwise=0)
-
         return model, originOut_x, originOut_y, dxpix, dypix, lam_beam, ysens
 
     def writerSensitivityFile(self, conffile='', beam='', w=None, sens=None):

observation_sim/mock_objects/SpecDisperser/disperse_c/disperse.pyx

@@ -20,7 +20,30 @@ import cython
 cdef extern from "math.h":
     double sqrt(double x)
     double exp(double x)
+
+
+def check_nan2D(np.ndarray[FTYPE_t, ndim=2] arr):
+    cdef int i, j
+    cdef int nrows = arr.shape[0]
+    cdef int ncols = arr.shape[1]
     
+    # 遍历数组的每个元素并检查是否存在 NaN
+    for i in range(nrows):
+        for j in range(ncols):
+            if np.isnan(arr[i, j]) | np.isinf(arr[i, j]):
+                return True
+    return False
+
+def check_nan1d(np.ndarray[DTYPE_t, ndim=1] arr):
+    cdef int i
+    cdef int n = arr.shape[0]
+
+    # 遍历数组的每个元素并检查是否存在 NaN
+    for i in range(n):
+        if np.isnan(arr[i]) | np.isinf(arr[i]):
+            return True
+    return False
+
 @cython.boundscheck(False)
 @cython.wraparound(False)
 @cython.embedsignature(True)
@@ -53,6 +76,18 @@ def disperse_grism_object(np.ndarray[FTYPE_t, ndim=2] flam,
 
     nk = len(idxl)
     nl = len(full)
+    
+
+    #if check_nan2D(flam):
+    #    print("DEBUG: disperse, input Array 'flam' contains NaN.")
+    #if check_nan1d(ysens):
+    #    print("DEBUG: disperse, input Array 'ysens' contains NaN.")
+    #if check_nan1d(yfrac):
+    #    print("DEBUG: disperse, input Array 'yfrac' contains NaN.")
+    #if check_nan1d(full):
+    #    print("DEBUG: disperse, input Array 'full' contains NaN before processing.")
+
+
 
     if (flat is not None):
         nlamb = len(wlambda)
@@ -95,14 +130,15 @@ def disperse_grism_object(np.ndarray[FTYPE_t, ndim=2] flam,
 
     else:
         for i in range(0-x0[1], x0[1]):
-            if (x0[1]+i < 0) | (x0[1]+i >= shd[1]):
+            x_pos = x0[1]+i
+            if (x_pos < 0) | (x_pos >= shd[1]):
                 continue
 
             for j in range(0-x0[0], x0[0]):
-                if (x0[0]+j < 0) | (x0[0]+j >= shd[0]):
+                y_pos = x0[0]+j
+                if (y_pos < 0) | (y_pos >= shd[0]):
                     continue
-
-                fl_ij = flam[x0[0]+j, x0[1]+i] #/1.e-17
+                fl_ij = flam[y_pos, x_pos] #/1.e-17
                 if (fl_ij == 0):
                     continue
 
@@ -110,10 +146,13 @@ def disperse_grism_object(np.ndarray[FTYPE_t, ndim=2] flam,
                     k1 = idxl[k]+j*shg[1]+i
                     if (k1 >= 0) & (k1 < nl):
                         full[k1] += ysens[k]*fl_ij*(1-yfrac[k])
-
                     k2 = idxl[k]+(j+1)*shg[1]+i
                     if (k2 >= 0) & (k2 < nl):
                         full[k2] += ysens[k]*fl_ij*yfrac[k]
+  
+        #if (check_nan1d(full)):
+        #    print("DEBUG: disperse, output Array 'full' contains NaN after processing.+++++++++++++++++++++++++++")
+
 
     return True
 

observation_sim/psf/PSFGauss.py

@@ -17,7 +17,7 @@ class PSFGauss(PSFModel):
             self.fwhm = self.fwhmGauss(r=psfRa)
             self.sigGauss = psfRa  # 80% light radius
         self.sigSpin = sigSpin
-        self.psf = galsim.Gaussian(flux=1.0, fwhm=fwhm)
+        self.psf = galsim.Gaussian(flux=1.0, fwhm=self.fwhm)
 
     def perfGauss(self, r, sig):
         """
@@ -122,4 +122,4 @@ class PSFGauss(PSFModel):
 
         # return ell, beta, qr
         PSFshear = galsim.Shear(e=ell, beta=beta*galsim.radians)
-        return self.psf.shear(PSFshear), PSFshear
+        return self.psf.shear(PSFshear), PSFshear
\ No newline at end of file

observation_sim/psf/PSFInterpSLS.py

@@ -20,8 +20,10 @@ import os
 from astropy.io import fits
 
 from astropy.modeling.models import Gaussian2D
-from scipy import signal
-
+from scipy import signal, interpolate
+import datetime
+import gc
+from jax import numpy as jnp
 
 LOG_DEBUG = False  # ***#
 NPSF = 900  # ***# 30*30
@@ -433,7 +435,16 @@ class PSFInterpSLS(PSFModel):
         # PSF_int_trans[ids_szero] = 0
         # print(PSF_int_trans[ids_szero].shape[0],PSF_int_trans.shape)
         PSF_int_trans = PSF_int_trans/np.sum(PSF_int_trans)
+        ###DEBGU
+        ids_szero = PSF_int_trans<0
+        n01 = PSF_int_trans[ids_szero].shape[0]
+
+        n1 = np.sum(np.isinf(PSF_int_trans))
+        n2 = np.sum(np.isnan(PSF_int_trans))
+        if n1>0 or n2>0:
+            print("DEBUG: PSFInterpSLS, inf:%d, nan:%d, 0 num:%d"%(n1, n2, n01))
 
+        ####
         # from astropy.io import fits
         # fits.writeto(str(bandNo) + '_' + g_order+ '_psf_o.fits', PSF_int_trans)
 
@@ -479,6 +490,215 @@ class PSFInterpSLS(PSFModel):
 
         return PSF_int_trans, PSF_int
 
+    def get_PSF_AND_convolve_withsubImg(self, chip, cutImg=None, pos_img_local=[1000, 1000], bandNo=1, g_order='A', grating_split_pos=3685):
+        """
+        Get the PSF at a given image position
+
+        Parameters:
+            chip: A 'Chip' object representing the chip we want to extract PSF from.
+            pos_img: A 'galsim.Position' object representing the image position.
+            bandpass: A 'galsim.Bandpass' object representing the wavelength range.
+            pixSize: The pixels size of psf matrix
+            findNeighMode: 'treeFind' or 'hoclistFind'
+        Returns:
+            PSF: A 'galsim.GSObject'.
+        """
+        order_IDs = {'A': '1', 'B': '0', 'C': '0', 'D': '0', 'E': '0'}
+        contam_order_sigma = {'C': 0.28032344707964174,
+                              'D': 0.39900182912061344, 'E': 1.1988309797685412}  # arcsec
+        x_start = chip.x_cen/chip.pix_size - chip.npix_x / 2.
+        y_start = chip.y_cen/chip.pix_size - chip.npix_y / 2.
+        # print(pos_img.x - x_start)
+        # pos_img_x = pos_img_local[0] + x_start
+        # pos_img_y = pos_img_local[1] + y_start
+        # pos_img = galsim.PositionD(pos_img_x, pos_img_y)
+        # centerPos_local = cutImg.ncol/2.
+        if pos_img_local[0] < grating_split_pos:
+            psf_data = self.grating1_data
+        else:
+            psf_data = self.grating2_data
+
+        grating_order = order_IDs[g_order]
+        # if grating_order in ['-2','-1','2']:
+        #     grating_order = '1'
+
+        # if grating_order in ['0', '1']:
+        psf_order = psf_data['order'+grating_order]
+        psf_order_b = psf_order['band'+str(bandNo)]
+        psf_b_dat = psf_order_b['band_data']
+        # pos_p = psf_b_dat[1].data
+        pos_p = psf_b_dat[1].data/chip.pix_size - np.array([y_start, x_start])
+
+        pc_coeff = psf_b_dat[2].data
+        pcs = psf_b_dat[0].data
+
+        npc = 10
+        m_size = int(pcs.shape[0]**0.5)
+        sumImg = np.sum(cutImg.array)
+        tmp_img = cutImg*0
+        for j in np.arange(npc):
+            X_ = jnp.hstack((pos_p[:,1].flatten()[:, None], pos_p[:,0].flatten()[:, None]),dtype=np.float32)
+            Z_ = (pc_coeff[j].astype(np.float32)).flatten()
+            # print(pc_coeff[j].shape[0], pos_p[:,1].shape[0], pos_p[:,0].shape[0])
+            cx_len = int(chip.npix_x)
+            cy_len = int(chip.npix_y)
+            n_x = jnp.arange(0, cx_len, 1, dtype = int)
+            n_y = jnp.arange(0, cy_len, 1, dtype = int)
+            M, N = jnp.meshgrid(n_x, n_y)
+            # t1=datetime.datetime.now()
+    #         U = interpolate.griddata(X_, Z_, (M[0:cy_len, 0:cx_len],N[0:cy_len, 0:cx_len]),
+    # method='nearest',fill_value=1.0)
+            b_img = galsim.Image(cx_len, cy_len)
+            b_img.setOrigin(0,0)
+            bounds = cutImg.bounds & b_img.bounds
+            if bounds.area() == 0:
+                continue
+
+            # ys = cutImg.ymin
+            # if ys < 0: 
+            #     ys = 0
+            # ye = cutImg.ymin+cutImg.nrow
+            # if ye >= cy_len-1: 
+            #     ye = cy_len-1
+            # if ye - ys <=0:
+            #     continue
+            # xs = cutImg.xmin
+            # if xs < 0: 
+            #     xs = 0
+            # xe = cutImg.xmin+cutImg.ncol
+            # if xe >= cx_len-1: 
+            #     xe = cx_len-1
+            # if xe - xs <=0:
+            #     continue
+            ys = bounds.ymin
+            ye = bounds.ymax+1
+            xs = bounds.xmin
+            xe = bounds.xmax+1
+            U = interpolate.griddata(X_, Z_, (M[ys:ye, xs:xe],N[ys:ye, xs:xe]),
+    method='nearest',fill_value=1.0)
+            # t2=datetime.datetime.now()
+            
+            # print("time interpolate:", t2-t1)
+
+            # if U.shape != cutImg.array.shape:
+            #     print('DEBUG:SHAPE',cutImg.ncol,cutImg.nrow,cutImg.xmin, cutImg.ymin)
+            #     continue
+            img_tmp = cutImg
+            img_tmp[bounds] = img_tmp[bounds]*U
+            psf = pcs[:, j].reshape(m_size, m_size)
+            tmp_img = tmp_img + signal.fftconvolve(img_tmp.array, psf, mode='same', axes=None)
+
+            # t3=datetime.datetime.now()
+            # print("time convole:", t3-t2)
+            del U
+            del img_tmp
+        if np.sum(tmp_img.array)==0:
+            tmp_img = cutImg
+        else:
+            tmp_img = tmp_img/np.sum(tmp_img.array)*sumImg
+        return tmp_img
+
+
+    def convolveFullImgWithPCAPSF(self, chip, folding_threshold=5.e-3):
+        keys_L1= chip_utils.getChipSLSGratingID(chip.chipID)
+        # keys_L2 = ['order-2','order-1','order0','order1','order2']
+        keys_L2 = ['order0','order1']
+        keys_L3 = ['w1','w2','w3','w4']
+
+        npca = 10
+
+        x_start = chip.x_cen/chip.pix_size - chip.npix_x / 2.
+        y_start = chip.y_cen/chip.pix_size - chip.npix_y / 2.
+
+        for i,gt in enumerate(keys_L1):
+            psfCo = self.grating1_data
+            if i > 0:
+                psfCo = self.grating2_data
+            for od in keys_L2:
+                psfCo_L2 = psfCo['order1']
+                if od in ['order-2','order-1','order0','order2']:
+                    psfCo_L2 = psfCo['order0']
+                for w in keys_L3:
+                    img = chip.img_stack[gt][od][w]
+                    pcs = psfCo_L2['band'+w[1]]['band_data'][0].data
+                    pos_p = psfCo_L2['band'+w[1]]['band_data'][1].data/chip.pix_size - np.array([y_start, x_start])
+                    pc_coeff = psfCo_L2['band'+w[1]]['band_data'][2].data
+                    # print("DEBUG-----------",np.max(pos_p[:,1]),np.min(pos_p[:,1]), np.max(pos_p[:,0]),np.min(pos_p[:,0]))
+                    sum_img = np.sum(img.array)
+                    
+                    
+                    # coeff_mat = np.zeros([npca, chip.npix_y, chip.npix_x])
+                    # for m in np.arange(chip.npix_y):
+                    #     for n in np.arange(chip.npix_x):
+                    #         px = n
+                    #         py = m
+                            
+                    #         dist2 = (pos_p[:, 1] - px)*(pos_p[:, 1] - px) + (pos_p[:, 0] - py)*(pos_p[:, 0] - py)
+                    #         temp_sort_dist = np.zeros([dist2.shape[0], 2])
+                    #         temp_sort_dist[:, 0] = np.arange(0, dist2.shape[0], 1)
+                    #         temp_sort_dist[:, 1] = dist2
+                    #         # print(temp_sort_dist)
+                    #         dits2_sortlist = sorted(temp_sort_dist, key=lambda x: x[1])
+                    #         # print(dits2_sortlist)
+                    #         nearest4p = np.zeros([4, 3])
+                    #         pc_coeff_4p = np.zeros([npca, 4])
+
+                    #         for i in np.arange(4):
+                    #             smaller_ids = int(dits2_sortlist[i][0])
+                    #             nearest4p[i, 0] = pos_p[smaller_ids, 1]
+                    #             nearest4p[i, 1] = pos_p[smaller_ids, 0]
+                    #             # print(pos_p[smaller_ids, 1],pos_p[smaller_ids, 0])
+                    #             nearest4p[i, 2] = dits2_sortlist[i][1]
+                    #             pc_coeff_4p[:, i] = pc_coeff[npca, smaller_ids]
+                    #         # idw_dist = 1/(np.sqrt((px-nearest4p[:, 0]) * (px-nearest4p[:, 0]) + (
+                    #         #     py-nearest4p[:, 1]) * (py-nearest4p[:, 1])))
+                    #         idw_dist = 1/(np.sqrt(nearest4p[:, 2]))
+
+                    #         coeff_int = np.zeros(npca)
+                    #         for i in np.arange(4):
+                    #             coeff_int = coeff_int + pc_coeff_4p[:, i]*idw_dist[i]
+                    #         coeff_mat[:, m, n] = coeff_int
+
+                    m_size = int(pcs.shape[0]**0.5)
+                    tmp_img = np.zeros_like(img.array,dtype=np.float32)
+                    for j in np.arange(npca):
+                        print(gt, od, w, j)
+                        X_ = jnp.hstack((pos_p[:,1].flatten()[:, None], pos_p[:,0].flatten()[:, None]),dtype=np.float32)
+                        Z_ = (pc_coeff[j].astype(np.float32)).flatten()
+                        # print(pc_coeff[j].shape[0], pos_p[:,1].shape[0], pos_p[:,0].shape[0])
+                        sub_size = 4
+                        cx_len = int(chip.npix_x/sub_size)
+                        cy_len = int(chip.npix_y/sub_size)
+                        n_x = jnp.arange(0, chip.npix_x, sub_size, dtype = int)
+                        n_y = jnp.arange(0, chip.npix_y, sub_size, dtype = int)
+
+                        M, N = jnp.meshgrid(n_x, n_y)
+                        t1=datetime.datetime.now()
+                #         U = interpolate.griddata(X_, Z_, (M[0:cy_len, 0:cx_len],N[0:cy_len, 0:cx_len]),
+                # method='nearest',fill_value=1.0)
+                        U1 = interpolate.griddata(X_, Z_, (M, N),
+                method='nearest',fill_value=1.0)
+                        U = np.zeros_like(chip.img.array, dtype=np.float32)
+                        for mi in np.arange(cy_len):
+                            for mj in np.arange(cx_len):
+                                U[mi*sub_size:(mi+1)*sub_size, mj*sub_size:(mj+1)*sub_size]=U1[mi,mj]
+                        t2=datetime.datetime.now()
+                        
+                        print("time interpolate:", t2-t1)
+
+                        img_tmp = img.array*U
+                        psf = pcs[:, j].reshape(m_size, m_size)
+                        tmp_img = tmp_img + signal.fftconvolve(img_tmp, psf, mode='same', axes=None)
+
+                        t3=datetime.datetime.now()
+                        print("time convole:", t3-t2)
+                        del U
+                        del U1
+                        
+                    chip.img = chip.img + tmp_img*sum_img/np.sum(tmp_img)
+                    del tmp_img
+                    gc.collect()
+
         # pixSize = np.rad2deg(self.pixsize*1e-3/28)*3600  #set psf pixsize
         #
         # # assert self.iccd == int(chip.getChipLabel(chipID=chip.chipID)), 'ERROR: self.iccd != chip.chipID'

observation_sim/sim_steps/add_objects.py

@@ -217,6 +217,26 @@ def add_objects(self, chip, filt, tel, pointing, catalog, obs_param):
         obj.unload_SED()
         del obj
         # gc.collect()
+
+    if chip.survey_type == "spectroscopic" and not self.overall_config["run_option"]["out_cat_only"] and chip.slsPSFOptim:
+        # from observation_sim.instruments.chip import chip_utils as chip_utils
+        # gn = chip_utils.getChipSLSGratingID(chip.chipID)[0]
+        # img1 = np.zeros([2,chip.img.array.shape[0],chip.img.array.shape[1]])
+
+        # for id1 in np.arange(2):
+        #     gn = chip_utils.getChipSLSGratingID(chip.chipID)[id1]
+        #     img_i = 0
+        #     for id2 in ['0','1']:
+        #         o_n = "order"+id2
+        #         for id3 in ['1','2','3','4']:
+        #             w_n = "w"+id3
+        #             img1[img_i] = img1[img_i] + chip.img_stack[gn][o_n][w_n].array
+        #         img_i = img_i + 1
+        # from astropy.io import fits
+        # fits.writeto('order0.fits',img1[0],overwrite=True)
+        # fits.writeto('order1.fits',img1[1],overwrite=True)
+
+        psf_model.convolveFullImgWithPCAPSF(chip) 
     del psf_model
     gc.collect()