Merge branch 'milky_way_extinction' into 'release_v3.0'

bug fix for guidance abberation correction

See merge request !28

catalog/C9_Catalog.py

@@ -12,7 +12,7 @@ from astropy.table import Table
 from scipy import interpolate
 from datetime import datetime
 
-from observation_sim.mock_objects import CatalogBase, Star, Galaxy, Quasar
+from observation_sim.mock_objects import CatalogBase, Star, Galaxy, Quasar, ExtinctionMW
 from observation_sim.mock_objects._util import tag_sed, getObservedSED, getABMAG, integrate_sed_bandpass, comoving_dist
 from observation_sim.astrometry.Astrometry_util import on_orbit_obs_position
 
@@ -110,6 +110,16 @@ class Catalog(CatalogBase):
 
         self.max_size = 0.
 
+        # [TODO] Milky Way extinction
+        if "enable_mw_ext_gal" in config["catalog_options"] and config["catalog_options"]["enable_mw_ext_gal"]:
+            if "planck_ebv_map" not in config["catalog_options"]:
+                raise ValueError(
+                    "Planck dust map must be given to enable Milky Way extinction calculation for galaxies.")
+            self.mw_ext = ExtinctionMW()
+            self.mw_ext.init_ext_model(model_name="odonnell")
+            self.mw_ext.load_Planck_ext(
+                file_path=config["catalog_options"]["planck_ebv_map"])
+
         if "star_cat" in config["catalog_options"]["input_path"] and config["catalog_options"]["input_path"]["star_cat"] and not config["catalog_options"]["galaxy_only"]:
             # Get the cloest star catalog file
             star_file_name = get_star_cat(
@@ -248,6 +258,12 @@ class Catalog(CatalogBase):
                 input_time_str=time_str
             )
 
+        # [TODO] get Milky Way extinction AVs
+        if "enable_mw_ext_gal" in self.config["catalog_options"] and self.config["catalog_options"]["enable_mw_ext_gal"]:
+            MW_Av_arr = self.mw_ext.Av_from_Planck(ra=ra_arr, dec=dec_arr)
+        else:
+            MW_Av_arr = np.zeros(len(ra_arr))
+
         for igals in range(ngals):
             # # (TEST)
             # if igals > 100:
@@ -259,6 +275,9 @@ class Catalog(CatalogBase):
             param['ra_orig'] = gals['ra'][igals]
             param['dec_orig'] = gals['dec'][igals]
 
+            # [TODO] get Milky Way extinction AVs
+            param['mw_Av'] = MW_Av_arr[igals]
+
             if not self.chip.isContainObj(ra_obj=param['ra'], dec_obj=param['dec'], margin=200):
                 continue
 
@@ -530,11 +549,30 @@ class Catalog(CatalogBase):
         speci = interpolate.interp1d(wave, flux)
         lamb = np.arange(2000, 11001+0.5, 0.5)
         y = speci(lamb)
+
+        # [TODO] Apply Milky Way extinction
+        if obj.type != 'star' and ("enable_mw_ext_gal" in self.config["catalog_options"] and self.config["catalog_options"]["enable_mw_ext_gal"]):
+            y = self.mw_ext.apply_extinction(y, Av=obj.mw_Av)
+
         # erg/s/cm2/A --> photon/s/m2/A
         all_sed = y * lamb / (cons.h.value * cons.c.value) * 1e-13
         sed = Table(np.array([lamb, all_sed]).T, names=('WAVELENGTH', 'FLUX'))
-        if obj.type == 'quasar':
+        # if obj.type == 'quasar':
+        #     # integrate to get the magnitudes
+        #     sed_photon = np.array([sed['WAVELENGTH'], sed['FLUX']]).T
+        #     sed_photon = galsim.LookupTable(x=np.array(sed_photon[:, 0]), f=np.array(
+        #         sed_photon[:, 1]), interpolant='nearest')
+        #     sed_photon = galsim.SED(
+        #         sed_photon, wave_type='A', flux_type='1', fast=False)
+        #     interFlux = integrate_sed_bandpass(
+        #         sed=sed_photon, bandpass=self.filt.bandpass_full)
+        #     obj.param['mag_use_normal'] = getABMAG(
+        #         interFlux, self.filt.bandpass_full)
+        #     # mag = getABMAG(interFlux, self.filt.bandpass_full)
+        #     # print("mag diff = %.3f"%(mag - obj.param['mag_use_normal']))
+
         # integrate to get the magnitudes
+        if obj.type == 'quasar' or obj.type == 'galaxy':
             sed_photon = np.array([sed['WAVELENGTH'], sed['FLUX']]).T
             sed_photon = galsim.LookupTable(x=np.array(sed_photon[:, 0]), f=np.array(
                 sed_photon[:, 1]), interpolant='nearest')

config/config_overall.yaml

@@ -11,7 +11,7 @@
 # can add some of the command-line arguments here as well;
 # ok to pass either way or both, as long as they are consistent
 work_dir: "/public/home/fangyuedong/project/workplace/"
-run_name: "no_nonlinearity_test"
+run_name: "ext_on"
 
 # Project cycle and run counter are used to name the outputs
 project_cycle: 9
@@ -49,6 +49,10 @@ catalog_options:
   # rotate galaxy ellipticity
   rotateEll: 0. # [degree]
 
+  # Whether to apply milky way extinction to galaxies
+  enable_mw_ext_gal: YES
+  planck_ebv_map: "/public/home/fangyuedong/project/ext_maps/planck/HFI_CompMap_ThermalDustModel_2048_R1.20.fits"
+
 ###############################################
 # Observation setting
 ###############################################
@@ -65,10 +69,10 @@ obs_setting:
   # - give a list of indexes of pointings: [ip_1, ip_2...]
   # - run all pointings: null
   # Note: only valid when a pointing list is specified
-  run_pointings: [1]
+  run_pointings: [0, 1, 2, 3, 4]
 
   # Whether to enable astrometric modeling
-  enable_astrometric_model: True
+  enable_astrometric_model: YES
 
   # Cut by saturation magnitude in which band?
   cut_in_band: "z"

config/obs_config_SCI.yaml

@@ -16,7 +16,7 @@ obs_id: "00000001" # this setting will only be used if pointing list file is not
 # Define list of chips
 # run_chips: [6,7,8,9,11,12,13,14,15,16,17,18,19,20,22,23,24,25] # Photometric chips
 #run_chips: [1,2,3,4,5,10,21,26,27,28,29,30] # Spectroscopic chips
-run_chips: [1]
+run_chips: [17, 22]
 
 # Define observation sequence
 call_sequence:

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,27 @@ 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.)
@@ -45,7 +66,7 @@ class Observation(object):
             chip.shutter_img = np.ones_like(chip.img.array)
         else:
             chip.shutter_img = effects.ShutterEffectArr(
-                chip.img, t_shutter=1.3, dist_bearing=735, dt=1E-3)
+                chip.img, t_exp=pointing.exp_time, t_shutter=1.3, dist_bearing=735, dt=1E-3)
         chip.prnu_img = effects.PRNU_Img(xsize=chip.npix_x, ysize=chip.npix_y, sigma=0.01,
                                          seed=int(self.config["random_seeds"]["seed_prnu"]+chip.chipID))
 
@@ -91,17 +112,24 @@ class Observation(object):
                 input_date_str=date_str,
                 input_time_str=time_str
             )
-            ra_cen, dec_cen = ra_cen[0], dec_cen[0]
+            ra_offset, dec_offset = pointing.ra - \
+                ra_cen[0], pointing.dec - dec_cen[0]
         else:
+            ra_offset, dec_offset = 0., 0.
         ra_cen = pointing.ra
         dec_cen = pointing.dec
 
+        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,
-                             chip_output=chip_output, all_filters=self.all_filters)
+                             chip_output=chip_output,
+                             all_filters=self.all_filters,
+                             ra_offset=ra_offset,
+                             dec_offset=dec_offset)
 
         for step in pointing.obs_param["call_sequence"]:
             if self.config["run_option"]["out_cat_only"]:
@@ -130,6 +158,9 @@ class Observation(object):
         chip_output.Log_info("check running:1: pointing-%d chip-%d pid-%d memory-%6.2fGB" % (pointing.id,
                              chip.chipID, os.getpid(), (psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024 / 1024)))
         del chip.img
+        del chip.flat_img
+        del chip.prnu_img
+        del chip.shutter_img
 
     def runExposure_MPI_PointingList(self, pointing_list, chips=None):
         comm = MPI.COMM_WORLD

observation_sim/config/ChipOutput.py

@@ -78,12 +78,12 @@ class ChipOutput(object):
                 self.hdr += "\n"
             self.cat.write(self.hdr)
 
-    def cat_add_obj(self, obj, pos_img, pos_shear):
+    def cat_add_obj(self, obj, pos_img, pos_shear, ra_offset=0., dec_offset=0.):
         ximg = obj.real_pos.x + 1.0
         yimg = obj.real_pos.y + 1.0
 
         line = self.fmt % (
-            obj.id, int(self.chip_label), self.filt.filter_type, ximg, yimg, obj.ra, obj.dec, obj.ra_orig, obj.dec_orig, obj.z, obj.getMagFilter(
+            obj.id, int(self.chip_label), self.filt.filter_type, ximg, yimg, obj.ra + ra_offset, obj.dec + dec_offset, obj.ra_orig, obj.dec_orig, obj.z, obj.getMagFilter(
                 self.filt), obj.type,
             obj.pmra, obj.pmdec, obj.rv, obj.parallax)
         line += obj.additional_output_str

observation_sim/config/header/ImageHeader.py

@@ -434,7 +434,7 @@ def generatePrimaryHeader(xlen=9216, ylen=9232, pointing_id='00000001', pointing
     co = coord.SkyCoord(ra, dec, unit='deg')
 
     ra_hms = format(co.ra.hms.h, '02.0f') + format(co.ra.hms.m,
-                                                   '02.0f') + format(co.ra.hms.s, '02.1f')
+                                                   '02.0f') + format(co.ra.hms.s, '04.1f')
     dec_hms = format(co.dec.dms.d, '02.0f') + format(abs(co.dec.dms.m),
                                                      '02.0f') + format(abs(co.dec.dms.s), '02.0f')
     if dec >= 0:

observation_sim/instruments/chip/effects.py

@@ -793,7 +793,7 @@ def ShutterEffectArr(GSImage, t_exp=150, t_shutter=1.3, dist_bearing=735, dt=1E-
     sizey = ymax-ymin+1
     xnewgrid = np.mgrid[xmin:(xmin+sizex)]
     expeffect = interpolate.splev(xnewgrid, intp, der=0)
-    expeffect /= np.max(expeffect)
+    expeffect /= t_exp
     exparrnormal = np.tile(expeffect, (sizey,1))
     # GSImage *= exparrnormal
 

observation_sim/instruments/data/filters/fgs_sub.list

-3000
+2000
 4500
 4750
 5000

observation_sim/instruments/data/filters/g_sub.list

-3800
+2000
 4217
 4432
 4631
@@ -6,4 +6,4 @@
 5002
 5179
 5354
-5799
\ No newline at end of file
+11000
\ No newline at end of file

observation_sim/instruments/data/filters/i_sub.list

-6600
+2000
 7061
 7255
 7448
@@ -6,4 +6,4 @@
 7833
 8027
 8226
-8999
\ No newline at end of file
+11000
\ No newline at end of file

observation_sim/instruments/data/filters/nuv_sub.list

-2513
+2000
 2621
 2716
 2805
@@ -6,4 +6,4 @@
 2969
 3050
 3132
-3499
\ No newline at end of file
+11000
\ No newline at end of file

observation_sim/instruments/data/filters/r_sub.list

-5100
+2000
 5642
 5821
 6001
@@ -6,4 +6,4 @@
 6363
 6547
 6735
-7199
\ No newline at end of file
+11000
\ No newline at end of file

observation_sim/instruments/data/filters/u_sub.list

-3000
+2000
 3277
 3380
 3485
@@ -6,4 +6,4 @@
 3703
 3813
 3918
-4499
\ No newline at end of file
+11000
\ No newline at end of file

observation_sim/instruments/data/filters/y_sub.list

-9000
+2000
 9322
 9405
 9489
@@ -6,4 +6,4 @@
 9695
 9832
 10024
-10590
\ No newline at end of file
+11000
\ No newline at end of file

observation_sim/instruments/data/filters/z_sub.list

-7800
+2000
 8494
 8638
 8790
@@ -6,4 +6,4 @@
 9141
 9363
 9663
-10590
\ No newline at end of file
+11000
\ No newline at end of file

observation_sim/mock_objects/ExtinctionMW.py

+import os
+import numpy as np
+import healpy as hp
+from astropy.io import fits
+from astropy import units as u
+from astropy.coordinates import SkyCoord
+
+
+class ExtinctionMW(object):
+    def __init__(self):
+        self.Rv = 3.1
+        self.lamb = np.arange(2000, 11001+0.5, 0.5)
+
+    @staticmethod
+    def radec2pix(ra, dec, NSIDE=2048):
+        return hp.pixelfunc.ang2pix(nside=NSIDE, theta=ra, phi=dec, nest=True, lonlat=True)
+
+    def init_ext_model(self, model_name="odonnell", Av=1.0, Rv=3.1, lamb=None):
+        self.model_name = model_name
+        self.Av = Av
+        self.Rv = Rv
+        if lamb is not None:
+            self.lamb = lamb
+
+        if self.model_name == "odonnell":
+            alpha = np.zeros(self.lamb.shape)
+            beta = np.zeros(self.lamb.shape)
+            x = 1.e4 / self.lamb
+            def f_alpha_1(x): return 0.574 * (x**1.61)
+            def f_beta_1(x): return -0.527 * (x**1.61)
+
+            def f_alpha_2(x):
+                y = x - 1.82
+                return 1 + 0.104*y - 0.609*(y**2) + 0.701*(y**3) + \
+                    1.137*(y**4) - 1.718*(y**5) - 0.827 * \
+                    (y**6) + 1.647*(y**7) - 0.505*(y**8)
+
+            def f_beta_2(x):
+                y = x - 1.82
+                return 1.952*y + 2.908*(y**2) - 3.989*(y**3) - 7.985 * \
+                    (y**4) + 11.102*(y**5) + 5.491 * \
+                    (y**6) - 10.805*(y**7) + 3.347*(y**8)
+
+            def f_alpha_3(x): return 1.752 - 0.316*x - \
+                0.104 / ((x - 4.67)**2 + 0.341)
+
+            def f_beta_3(x): return -3.090 + 1.825*x + \
+                1.206 / ((x - 4.62)**2 + 0.262)
+
+            def f_alpha_4(x): return f_alpha_3(x) + -0.04473 * \
+                (x - 5.9)**2 - 0.009779 * (x - 5.9)**3
+
+            def f_beta_4(x): return f_beta_3(x) + 0.2130 * \
+                (x - 5.9)**2 + 0.1207 * (x - 5.9)**3
+
+            def f_alpha_5(x): return -1.073 - 0.628*(x - 8) + \
+                0.137*(x - 8)**2 - 0.070*(x - 8)**3
+
+            def f_beta_5(x): return 13.670 + 4.257*(x - 8) - \
+                0.420*(x - 8)**2 + 0.374 * (x - 8)**3
+
+            alpha = np.piecewise(
+                x, [x <= 1.1, (x > 1.1)*(x <= 3.3), (x > 3.3)*(x <= 5.9), (x > 5.9)*(x <= 8.), (x > 8.)], [f_alpha_1, f_alpha_2, f_alpha_3, f_alpha_4, f_alpha_5])
+            beta = np.piecewise(
+                x, [x <= 1.1, (x > 1.1)*(x <= 3.3), (x > 3.3)*(x <= 5.9), (x > 5.9)*(x <= 8.), (x > 8.)], [f_beta_1, f_beta_2, f_beta_3, f_beta_4, f_beta_5])
+
+            self.ext = (alpha + beta / self.Rv) * self.Av
+
+    def load_Planck_ext(self, file_path):
+        hdu = fits.open(file_path)
+        self.ebv_planck = hdu[1].data["EBV"]
+
+    def Av_from_Planck(self, ra, dec):
+        if not hasattr(self, 'ebv_planck'):
+            raise ValueError(
+                "Need to load planck dust map first")
+        # Convert to galactic coordinates
+        c = SkyCoord(ra=ra * u.degree, dec=dec *
+                     u.degree, frame='icrs').galactic
+        l, b = c.l.radian, c.b.radian
+        NSIDE = hp.pixelfunc.get_nside(self.ebv_planck)
+        pix = hp.pixelfunc.ang2pix(
+            nside=NSIDE, theta=np.pi/2. - b, phi=l, nest=True)
+        return self.ebv_planck[pix] * self.Rv
+
+    def apply_extinction(self, spec, Av=1.0):
+        if len(spec) != len(self.lamb):
+            raise ValueError(
+                "Dimension of spec do not match that of ExtinctionMW.lamb: try initilize (init_ext_model) by the actual size of spec which you want to apply the extinction to")
+        if not hasattr(self, 'ext'):
+            raise ValueError(
+                "Need to initialize the extinction model (init_ext_model) first")
+        scale = 10**(-.4 * self.ext * Av)
+        # print("scale = ", scale)
+        spec *= scale
+        return spec

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')
+            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 - (starImg.center.y - starImg.ymin),
-                           x_nominal - (starImg.center.x - starImg.xmin)]
-            starImg.setOrigin(0, 0)
+            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_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]]
+                    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

@@ -21,6 +21,29 @@ 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)
@@ -54,6 +77,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)
         nflat = len(flat)
@@ -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,11 +146,14 @@ 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
 
 @cython.boundscheck(False)