add milky way extinction for galaxy

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/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/__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

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
\ No newline at end of file
+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