Merge branch 'develop' of https://csst-tb.bao.ac.cn/code/csst-sims/csst_msc_sim into develop

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,12 @@ class Catalog(CatalogBase):
 
         self.max_size = 0.
 
+        # [TODO] Milky Way extinction
+        self.mw_ext = ExtinctionMW()
+        self.mw_ext.init_ext_model(model_name="odonnell")
+        self.mw_ext.load_Planck_ext(
+            file_path="/public/home/fangyuedong/project/ext_maps/planck/HFI_CompMap_ThermalDustModel_2048_R1.20.fits")
+
         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 +254,9 @@ class Catalog(CatalogBase):
                 input_time_str=time_str
             )
 
+        # [TODO] get Milky Way extinction AVs
+        MW_Av_arr = self.mw_ext.Av_from_Planck(ra=ra_arr, dec=dec_arr)
+
         for igals in range(ngals):
             # # (TEST)
             # if igals > 100:
@@ -259,6 +268,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,6 +542,11 @@ 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':
+            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'))

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_test"
 
 # Project cycle and run counter are used to name the outputs
 project_cycle: 9
@@ -44,7 +44,7 @@ catalog_options:
   star_only: NO
 
   # Only simulate galaxies?
-  galaxy_only: NO
+  galaxy_only: YES
 
   # rotate galaxy ellipticity
   rotateEll: 0. # [degree]
@@ -65,7 +65,7 @@ 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

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

@@ -92,16 +92,21 @@ class Observation(object):
                 input_time_str=time_str
             )
             ra_cen, dec_cen = ra_cen[0], dec_cen[0]
+            ra_offset, dec_offset = pointing.ra - ra_cen, pointing.dec - dec_cen
         else:
             ra_cen = pointing.ra
             dec_cen = pointing.dec
+            ra_offset, dec_offset = 0., 0.
 
         # Prepare necessary chip properties for simulation
         chip = self.prepare_chip_for_exposure(chip, ra_cen, dec_cen, pointing)
 
         # 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 +135,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/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, 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)
+        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/MockObject.py

@@ -59,13 +59,13 @@ class MockObject(object):
         flux = self.getFluxFilter(filt)
         return flux * tel.pupil_area * exptime
 
-    def getPosWorld(self):
-        ra = self.param["ra"]
-        dec = self.param["dec"]
+    def getPosWorld(self, ra_offset=0., dec_offset=0.):
+        ra = self.param["ra"] + ra_offset
+        dec = self.param["dec"] + dec_offset
         return galsim.CelestialCoord(ra=ra * galsim.degrees, dec=dec * galsim.degrees)
 
-    def getPosImg_Offset_WCS(self, img, fdmodel=None, chip=None, verbose=True, chip_wcs=None, img_header=None):
-        self.posImg = img.wcs.toImage(self.getPosWorld())
+    def getPosImg_Offset_WCS(self, img, fdmodel=None, chip=None, verbose=True, chip_wcs=None, img_header=None, ra_offset=0., dec_offset=0.):
+        self.posImg = img.wcs.toImage(self.getPosWorld(ra_offset, dec_offset))
         self.localWCS = img.wcs.local(self.posImg)
         # Apply field distortion model
         if (fdmodel is not None) and (chip is not None):

observation_sim/mock_objects/__init__.py

@@ -5,3 +5,4 @@ from .Quasar import Quasar
 from .Star import Star
 from .Stamp import Stamp
 from .FlatLED import FlatLED
+from .ExtinctionMW import ExtinctionMW

observation_sim/sim_steps/__init__.py

 import os
 
+
 class SimSteps:
-    def __init__(self, overall_config, chip_output, all_filters):
+    def __init__(self, overall_config, chip_output, all_filters, ra_offset=0., dec_offset=0.):
         self.overall_config = overall_config
         self.chip_output = chip_output
         self.all_filters = all_filters
+        self.ra_offset = ra_offset
+        self.dec_offset = dec_offset
 
     from .prepare_headers import prepare_headers, updateHeaderInfo
     from .add_sky_background import add_sky_background_sci, add_sky_flat_calibration, add_sky_background
@@ -15,6 +18,7 @@ class SimSteps:
     from .readout_output import add_prescan_overscan, add_readout_noise, apply_gain, quantization_and_output
     from .add_LED_flat import add_LED_Flat
 
+
 SIM_STEP_TYPES = {
     "scie_obs": "add_objects",
     "sky_background": "add_sky_background",
@@ -31,6 +35,6 @@ SIM_STEP_TYPES = {
     "readout_noise": "add_readout_noise",
     "gain": "apply_gain",
     "quantization_and_output": "quantization_and_output",
-    "led_calib_model":"add_LED_Flat",
-    "sky_flatField":"add_sky_flat_calibration",
+    "led_calib_model": "add_LED_Flat",
+    "sky_flatField": "add_sky_flat_calibration",
 }

observation_sim/sim_steps/add_objects.py

@@ -145,7 +145,7 @@ def add_objects(self, chip, filt, tel, pointing, catalog, obs_param):
 
         # Get position of object on the focal plane
         pos_img, _, _, _, fd_shear = obj.getPosImg_Offset_WCS(
-            img=chip.img, fdmodel=fd_model, chip=chip, verbose=False, chip_wcs=chip_wcs, img_header=self.h_ext)
+            img=chip.img, fdmodel=fd_model, chip=chip, verbose=False, chip_wcs=chip_wcs, img_header=self.h_ext, ra_offset=self.ra_offset, dec_offset=self.dec_offset)
 
         # [TODO] For now, only consider objects which their centers (after field distortion) are projected within the focal plane
         # Otherwise they will be considered missed objects
@@ -216,7 +216,7 @@ def add_objects(self, chip, filt, tel, pointing, catalog, obs_param):
         # Unload SED:
         obj.unload_SED()
         del obj
-        gc.collect()
+        # gc.collect()
     del psf_model
     gc.collect()
 

submit_jobs.slurm

 #!/bin/bash
 
 #SBATCH -J CSSTSim
-#SBATCH -N 1
+#SBATCH -N 2
 #SBATCH --ntasks-per-node=6
 #SBATCH -p debug
-#SBATCH --mem=60G
+#SBATCH --mem=96G
 
 module load mpi/hpcx/2.4.1/gcc-7.3.1
 date
@@ -12,4 +12,4 @@ date
 #限定单节点任务数
 srun hostname -s | sort -n | awk -F"-" '{print $2}' | uniq > pnodes
 
-mpirun -mca pml ucx -x  UCX_NET_DEVICES=mlx5_0:1 -machinefile pnodes -np 6 --map-by node python3 /public/home/fangyuedong/project/csst_msc_sim/run_sim.py  --config_file config_overall.yaml --catalog C9_Catalog -c /public/home/fangyuedong/project/csst_msc_sim/config
+mpirun -mca pml ucx -x  UCX_NET_DEVICES=mlx5_0:1 -machinefile pnodes -np 12 --map-by node python3 /public/home/fangyuedong/project/csst_msc_sim/run_sim.py  --config_file config_overall.yaml --catalog C9_Catalog -c /public/home/fangyuedong/project/csst_msc_sim/config
\ No newline at end of file