add milky way extinction for galaxy

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

        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(

                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:

            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

        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'))

# 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

  star_only: NO

  # Only simulate galaxies?

  galaxy_only: NO

  galaxy_only: YES

  # rotate galaxy ellipticity

  rotateEll: 0. # [degree]

  # - 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

# 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:

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

from .Star import Star

from .Stamp import Stamp

from .FlatLED import FlatLED

from .ExtinctionMW import ExtinctionMW