Skip to content
C3Catalog.py 10.9 KiB
Newer Older
import os
import galsim
import random
import numpy as np
import h5py as h5
import healpy as hp
import astropy.constants as cons
from astropy.coordinates import spherical_to_cartesian
from astropy.table import Table
from scipy import interpolate
from datetime import datetime

from ObservationSim.MockObject import CatalogBase, Star, Galaxy, Quasar
from ObservationSim.MockObject._util import seds, sed_assign, extAv, tag_sed, getObservedSED
from ObservationSim.Astrometry.Astrometry_util import on_orbit_obs_position
try:
    import importlib.resources as pkg_resources
except ImportError:
    # Try backported to PY<37 'importlib_resources'
    import importlib_resources as pkg_resources

NSIDE = 128

class C3Catalog(CatalogBase):
    def __init__(self, config, chip, pointing, **kwargs):
        super().__init__()
        self.cat_dir = os.path.join(config["data_dir"], config["input_path"]["cat_dir"])
        self.seed_Av = config["random_seeds"]["seed_Av"]
        # self.normalize_dir = os.path.join(config["data_dir"], config["SLS_path"]["SLS_norm"])
        # self.normF_star = Table.read(os.path.join(self.normalize_dir, 'SLOAN_SDSS.g.fits'))
        # self.normF_galaxy = Table.read(os.path.join(self.normalize_dir, 'lsst_throuput_g.fits'))
        with pkg_resources.path('Catalog.data', 'SLOAN_SDSS.g.fits') as filter_path:
                self.normF_star = filter_path
        with pkg_resources.path('Catalog.data', 'lsst_throuput_g.fits') as filter_path:
                self.normF_star = filter_path
        self.config = config
        self.chip = chip
        self.pointing = pointing
        if "star_cat" in config["input_path"] and config["input_path"]["star_cat"] and not config["run_option"]["galaxy_only"]:
            star_file = config["input_path"]["star_cat"]
            star_SED_file = config["SED_templates_path"]["star_SED"]
            self.star_path = os.path.join(self.cat_dir, star_file)
            self.star_SED_path = os.path.join(config["data_dir"], star_SED_file)
            self._load_SED_lib_star()
        if "galaxy_cat" in config["input_path"] and config["input_path"]["galaxy_cat"] and not config["run_option"]["star_only"]:
            galaxy_file = config["input_path"]["galaxy_cat"]
            self.galaxy_path = os.path.join(self.cat_dir, galaxy_file)
            self.galaxy_SED_path = os.path.join(config["data_dir"], config["SED_templates_path"]["galaxy_SED"])
            self._load_SED_lib_gals()
        if "rotateEll" in config["shear_setting"]:
            self.rotation = float(int(config["shear_setting"]["rotateEll"]/45.))
        else:
            self.rotation = 0.

        self._get_healpix_list()
        self._load()

    def _get_healpix_list(self):
        self.sky_coverage = self.chip.getSkyCoverageEnlarged(self.chip.img.wcs, margin=0.2)
        ra_min, ra_max, dec_min, dec_max = self.sky_coverage.xmin, self.sky_coverage.xmax, self.sky_coverage.ymin, self.sky_coverage.ymax
        ra = np.deg2rad(np.array([ra_min, ra_max, ra_max, ra_min]))
        dec = np.deg2rad(np.array([dec_max, dec_max, dec_min, dec_min]))
        vertices = spherical_to_cartesian(1., dec, ra)
        self.pix_list = hp.query_polygon(NSIDE, np.array(vertices).T, inclusive=True)
        print("HEALPix List: ", self.pix_list)

    def load_norm_filt(self, obj):
        if obj.type == "star":
            return self.normF_star
        elif obj.type == "galaxy" or obj.type == "quasar":
            return self.normF_galaxy
        else:
            return None

    def _load_SED_lib_star(self):
        self.tempSED_star = h5.File(self.star_SED_path,'r')

    def _load_SED_lib_gals(self):
        self.tempSed_gal, self.tempRed_gal = seds("galaxy.list", seddir=self.galaxy_SED_path)

    def _load_gals(self, gals, pix_id=None):
        ngals = len(gals['galaxyID'])
        self.rng_sedGal = random.Random()
        self.rng_sedGal.seed(pix_id) # Use healpix index as the random seed
        self.ud = galsim.UniformDeviate(pix_id)
        for igals in range(ngals):
            param = self.initialize_param()
            param['ra'] = gals['ra_true'][igals]
            param['dec'] = gals['dec_true'][igals]
            if not self.chip.isContainObj(ra_obj=param['ra'], dec_obj=param['dec'], margin=200):
                continue
            param['mag_use_normal'] = gals['mag_true_g_lsst'][igals]
            if param['mag_use_normal'] >= 26.5:
                continue
            param['z'] = gals['redshift_true'][igals]
            param['model_tag'] = 'None'
            param['gamma1'] = 0
            param['gamma2'] = 0
            param['kappa'] = 0
            param['delta_ra'] = 0
            param['delta_dec'] = 0
            # sersicB = gals['sersic_bulge'][igals]
            hlrMajB = gals['size_bulge_true'][igals]
            hlrMinB = gals['size_minor_bulge_true'][igals]
            # sersicD = gals['sersic_disk'][igals]
            hlrMajD = gals['size_disk_true'][igals]
            hlrMinD = gals['size_minor_disk_true'][igals]
            aGal = gals['size_true'][igals]
            bGal = gals['size_minor_true'][igals]
            param['bfrac'] = gals['bulge_to_total_ratio_i'][igals]
            param['theta'] = gals['position_angle_true'][igals]
            param['hlr_bulge'] = np.sqrt(hlrMajB * hlrMinB)
            param['hlr_disk'] = np.sqrt(hlrMajD * hlrMinD)
            param['ell_bulge'] = (hlrMajB - hlrMinB)/(hlrMajB + hlrMinB)
            param['ell_disk'] = (hlrMajD - hlrMinD)/(hlrMajD + hlrMinD)
            param['ell_tot'] = (aGal - bGal) / (aGal + bGal)

            # Assign each galaxy a template SED
            param['sed_type'] = sed_assign(phz=param['z'], btt=param['bfrac'], rng=self.rng_sedGal)
            param['redden'] = self.tempRed_gal[param['sed_type']]
            param['av'] = self.avGal[int(self.ud()*self.nav)]
            if param['sed_type'] <= 5:
                param['av'] = 0.0
                param['redden'] = 0
            param['star'] = 0   # Galaxy
            if param['sed_type'] >= 29:
                param['av'] = 0.6 * param['av'] / 3.0 # for quasar, av=[0, 0.2], 3.0=av.max-av.im
                param['star'] = 2 # Quasar

            self.ids += 1
            param['id'] = self.ids
            
            if param['star'] == 0:
                obj = Galaxy(param, self.rotation)
                self.objs.append(obj)
            if param['star'] == 2:
                obj = Quasar(param)
                self.objs.append(obj)

    def _load_stars(self, stars, pix_id=None):
        nstars = len(stars['sourceID'])
        # Apply astrometric modeling
        ra_arr = stars["RA"][:]
        dec_arr = stars["Dec"][:]
        # if "astrometric_lib" in self.config["obs_setting"] and self.config["obs_setting"]["enable_astrometric_model"]:
        if self.config["obs_setting"]["enable_astrometric_model"]:
            ra_list = ra_arr.tolist()
            dec_list = dec_arr.tolist()
            pmra_list = np.zeros(nstars).tolist()
            pmdec_list = np.zeros(nstars).tolist()
            rv_list = np.zeros(nstars).tolist()
            parallax_list = [1e-9] * nstars
            # lib_path = os.path.join(self.config["data_dir"], self.config["obs_setting"]["astrometric_lib"])
            dt = datetime.fromtimestamp(self.pointing.timestamp)
            date_str = dt.date().isoformat()
            time_str = dt.time().isoformat()
            ra_arr, dec_arr = on_orbit_obs_position(
                input_ra_list=ra_list,
                input_dec_list=dec_list,
                input_pmra_list=pmra_list,
                input_pmdec_list=pmdec_list,
                input_rv_list=rv_list,
                input_parallax_list=parallax_list,
                input_nstars=nstars,
                input_x=self.pointing.sat_x,
                input_y=self.pointing.sat_y,
                input_z=self.pointing.sat_z,
                input_vx=self.pointing.sat_vx,
                input_vy=self.pointing.sat_vy,
                input_vz=self.pointing.sat_vz,
                input_epoch="J2015.5",
                input_date_str=date_str,
                input_time_str=time_str,
                # lib_path=lib_path
        for istars in range(nstars):
            param = self.initialize_param()
            param['ra'] = ra_arr[istars]
            param['dec'] = dec_arr[istars]
            if not self.chip.isContainObj(ra_obj=param['ra'], dec_obj=param['dec'], margin=200):
                continue
            param['mag_use_normal'] = stars['app_sdss_g'][istars]
            if param['mag_use_normal'] >= 26.5:
                continue
            self.ids += 1
            param['id'] = self.ids
            param['sed_type'] = stars['sourceID'][istars]
            param['model_tag'] = stars['model_tag'][istars]
            param['teff'] = stars['teff'][istars]
            param['logg'] = stars['grav'][istars]
            param['feh'] = stars['feh'][istars]
            param['z'] = 0.0
            param['star'] = 1   # Star
            obj = Star(param)
            self.objs.append(obj)

    def _load(self, **kwargs):
        self.nav = 15005
        self.avGal = extAv(self.nav, seed=self.seed_Av)
        self.objs = []
        self.ids = 0
        if "star_cat" in self.config["input_path"] and self.config["input_path"]["star_cat"] and not self.config["run_option"]["galaxy_only"]:
Fang Yuedong's avatar
Fang Yuedong committed
            star_cat = h5.File(self.star_path, 'r')['catalog']
            for pix in self.pix_list:
                stars = star_cat[str(pix)]
                self._load_stars(stars, pix_id=pix)
                del stars
        if "galaxy_cat" in self.config["input_path"] and self.config["input_path"]["galaxy_cat"] and not self.config["run_option"]["star_only"]:
Fang Yuedong's avatar
Fang Yuedong committed
            gals_cat = h5.File(self.galaxy_path, 'r')['galaxies']
            for pix in self.pix_list:
                gals = gals_cat[str(pix)]
                self._load_gals(gals, pix_id=pix)
                del gals
        print("number of objects in catalog: ", len(self.objs))
        del self.avGal


    def load_sed(self, obj, **kwargs):
        if obj.type == 'star':
            _, wave, flux = tag_sed(
                h5file=self.tempSED_star,
                model_tag=obj.param['model_tag'],
                teff=obj.param['teff'],
                logg=obj.param['logg'],
                feh=obj.param['feh']
            )
        elif obj.type == 'galaxy' or obj.type == 'quasar':
            sed_data = getObservedSED(
                sedCat=self.tempSed_gal[obj.sed_type],
                redshift=obj.z,
                av=obj.param["av"],
                redden=obj.param["redden"]
            )
            wave, flux = sed_data[0], sed_data[1]
        else:
            raise ValueError("Object type not known")
        speci = interpolate.interp1d(wave, flux)
        # lamb = np.arange(2500, 10001 + 0.5, 0.5)
        lamb = np.arange(2400, 11001 + 0.5, 0.5)
        y = speci(lamb)
        # erg/s/cm2/A --> photo/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'))