Galaxy.py

import numpy as np
import galsim
import os, sys
import astropy.constants as cons
from astropy.table import Table
from ._util import eObs, integrate_sed_bandpass, getNormFactorForSpecWithABMAG, getObservedSED, getABMAG
from .SpecDisperser import SpecDisperser
from .MockObject import MockObject
from scipy import interpolate

class Galaxy(MockObject):
    def __init__(self, param, rotation=None):
        super().__init__(param)
        self.thetaR = self.param["theta"]
        self.bfrac = self.param["bfrac"]
        self.hlr_disk = self.param["hlr_disk"]
        self.hlr_bulge = self.param["hlr_bulge"]

        # Extract ellipticity components
        self.e_disk = galsim.Shear(g=self.param["ell_disk"], beta=self.thetaR*galsim.degrees)
        self.e_bulge = galsim.Shear(g=self.param["ell_bulge"], beta=self.thetaR*galsim.degrees)
        self.e_total = galsim.Shear(g=self.param["ell_tot"], beta=self.thetaR*galsim.degrees)
        self.e1_disk, self.e2_disk = self.e_disk.g1, self.e_disk.g2
        self.e1_bulge, self.e2_bulge = self.e_bulge.g1, self.e_bulge.g2
        self.e1_total, self.e2_total = self.e_total.g1, self.e_total.g2

        if rotation is not None:
            self.rotateEllipticity(rotation)

    def load_SED(self, survey_type, sed_path, cosids=None, objtypes=None, sed_templates=None, normFilter=None, target_filt=None):
        if survey_type == "photometric":
            norm_thr_rang_ids = normFilter['SENSITIVITY'] > 0.001
            if sed_templates is None:
                # Read SED data directly
                itype = objtypes[cosids==self.sed_type][0]
                sed_file = os.path.join(sed_path, itype + "_ID%s.sed"%(self.sed_type))
                if not os.path.exists(sed_file):
                    raise ValueError("!!! No SED found.")
                sed_data = Table.read(sed_file, format="ascii")
                wave, flux = sed_data["observedLambda"].data, sed_data["observedFlux"].data
            else:
                # Load SED from templates
                sed_data = sed_templates[self.sed_type]
                # redshift, intrinsic extinction
                sed_data = getObservedSED(
                    sedCat=sed_data, 
                    redshift=self.z, 
                    av=self.param['av'], 
                    redden=self.param['redden'])
                wave, flux = sed_data[0], sed_data[1] 
            flux_photon = flux * (wave / (cons.h.value * cons.c.value)) * 1e-13
            sed_photon = Table(np.array([wave, flux_photon]).T, names=('WAVELENGTH', 'FLUX'))
            # Get scaling factor for SED
            sedNormFactor = getNormFactorForSpecWithABMAG(ABMag=self.param['mag_use_normal'],
                spectrum=sed_photon,
                norm_thr=normFilter,
                sWave=np.floor(normFilter[norm_thr_rang_ids][0][0]),
                eWave=np.ceil(normFilter[norm_thr_rang_ids][-1][0]))
            sed_photon = np.array([sed_photon['WAVELENGTH'], sed_photon['FLUX']*sedNormFactor]).T
            # Convert to galsim.SED object
            spec = galsim.LookupTable(x=np.array(sed_photon[:, 0]), f=np.array(sed_photon[:, 1]), interpolant='nearest')
            self.sed = galsim.SED(spec, wave_type='A', flux_type='1', fast=False)
            # Get magnitude
            interFlux = integrate_sed_bandpass(sed=self.sed, bandpass=target_filt.bandpass_full)
            self.param['mag_%s'%target_filt.filter_type] = getABMAG(
                interFlux=interFlux, 
                bandpass=target_filt.bandpass_full)
            # print('mag_use_normal = ', self.param['mag_use_normal'])
            # print('mag_%s = '%target_filt.filter_type, self.param['mag_%s'%target_filt.filter_type])
            # print('redshift = %.3f'%(self.z))
            # print('sed_type = %d, av = %.2f, redden = %d'%(self.sed_type, self.param['av'], self.param['redden']))

        elif survey_type == "spectroscopic":
            if sed_templates is None:
                self.sedPhotons(sed_path=sed_path, cosids=cosids, objtypes=objtypes)
            else:
                sed_data = sed_templates[self.sed_type]
                sed_data = getObservedSED(
                    sedCat=sed_data, 
                    redshift=self.z, 
                    av=self.param['av'], 
                    redden=self.param['redden'])
                speci = interpolate.interp1d(sed_data[0], sed_data[1])
                lamb = np.arange(2500, 10001 + 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
                self.sed = Table(np.array([lamb, all_sed]).T, names=('WAVELENGTH', 'FLUX'))


    def unload_SED(self):
        """(Test) free up SED memory
        """
        del self.sed

    def sedPhotons(self, sed_path, cosids, objtypes):

        itype = objtypes[cosids == self.sed_type][0]
        sed_file = os.path.join(sed_path, itype + "_ID%s.sed" % (self.sed_type))
        if not os.path.exists(sed_file):
            raise ValueError("!!! No SED found.")
        sed = Table.read(sed_file, format="ascii")
        spec_data = {}
        f_orig = sed["observedFlux"].data
        w_orig = sed["observedLambda"].data

        speci = interpolate.interp1d(w_orig, f_orig)
        lamb = np.arange(2500, 10001 + 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
        self.sed = Table(np.array([lamb, all_sed]).T, names=('WAVELENGTH', 'FLUX'))

    def getGSObj_multiband(self, tel, psf_list, bandpass_list, filt, nphotons_tot=None, g1=0, g2=0, exptime=150.):
        if len(psf_list) != len(bandpass_list):
            raise ValueError("!!!The number of PSF profiles and the number of bandpasses must be equal.")
        objs = []
        if nphotons_tot == None:
            nphotons_tot = self.getElectronFluxFilt(filt, tel, exptime)
        # print("nphotons_tot = ", nphotons_tot)

        try:
            full = integrate_sed_bandpass(sed=self.sed, bandpass=filt.bandpass_full)
        except Exception as e:
            print(e)
            return -1
        for i in range(len(bandpass_list)):
            bandpass = bandpass_list[i]
            try:
                sub = integrate_sed_bandpass(sed=self.sed, bandpass=bandpass)
            except Exception as e:
                print(e)
                return -1
            
            ratio = sub/full
            if not (ratio == -1 or (ratio != ratio)):
                nphotons = ratio * nphotons_tot
            else:
                return -1

            psf = psf_list[i]
            disk = galsim.Sersic(n=1.0, half_light_radius=self.hlr_disk, flux=1.0)
            disk_shape = galsim.Shear(g1=self.e1_disk, g2=self.e2_disk)
            disk = disk.shear(disk_shape)
            bulge = galsim.Sersic(n=4.0, half_light_radius=self.hlr_bulge, flux=1.0)
            bulge_shape = galsim.Shear(g1=self.e1_bulge, g2=self.e2_bulge)
            bulge = bulge.shear(bulge_shape)

            gal = self.bfrac * bulge + (1.0 - self.bfrac) * disk
            gal = gal.withFlux(nphotons)
            gal_shear = galsim.Shear(g1=g1, g2=g2)
            gal = gal.shear(gal_shear)
            gal = galsim.Convolve(psf, gal)
            objs.append(gal)
        final = galsim.Sum(objs)
        return final

    def drawObj_multiband(self, tel, pos_img, psf_model, bandpass_list, filt, chip, nphotons_tot=None, g1=0, g2=0, exptime=150.):
        if nphotons_tot == None:
            nphotons_tot = self.getElectronFluxFilt(filt, tel, exptime)
        # print("nphotons_tot = ", nphotons_tot)

        try:
            full = integrate_sed_bandpass(sed=self.sed, bandpass=filt.bandpass_full)
        except Exception as e:
            print(e)
            return False

        nphotons_sum = 0
        photons_list = []
        xmax, ymax = 0, 0

        # print('hlr_disk = %.4f, hlr_bulge = %.4f'%(self.hlr_disk, self.hlr_bulge))
        big_galaxy = False
        if self.hlr_disk > 3.0: # Very big galaxy
            big_galaxy = True

        # (TEST) Galsim Parameters
        if self.getMagFilter(filt) <= 15 and (not big_galaxy):
            folding_threshold = 5.e-4
        else:
            folding_threshold = 5.e-3
        gsp = galsim.GSParams(folding_threshold=folding_threshold)

        for i in range(len(bandpass_list)):
            bandpass = bandpass_list[i]

            try:
                sub = integrate_sed_bandpass(sed=self.sed, bandpass=bandpass)
            except Exception as e:
                print(e)
                # return False
                continue
            
            ratio = sub/full
            if not (ratio == -1 or (ratio != ratio)):
                nphotons = ratio * nphotons_tot
            else:
                # return False
                continue
            nphotons_sum += nphotons
            # print("nphotons_sub-band_%d = %.2f"%(i, nphotons))

            psf, pos_shear = psf_model.get_PSF(chip=chip, pos_img=pos_img, bandpass=bandpass, folding_threshold=folding_threshold)
            disk = galsim.Sersic(n=1.0, half_light_radius=self.hlr_disk, flux=1.0, gsparams=gsp)
            disk_shape = galsim.Shear(g1=self.e1_disk, g2=self.e2_disk)
            disk = disk.shear(disk_shape)
            bulge = galsim.Sersic(n=4.0, half_light_radius=self.hlr_bulge, flux=1.0, gsparams=gsp)
            bulge_shape = galsim.Shear(g1=self.e1_bulge, g2=self.e2_bulge)
            bulge = bulge.shear(bulge_shape)

            gal = self.bfrac * bulge + (1.0 - self.bfrac) * disk
            gal = gal.withFlux(nphotons)
            gal_shear = galsim.Shear(g1=g1, g2=g2)
            gal = gal.shear(gal_shear)

            if self.hlr_disk < 10.0: # Not apply PSF for very big galaxy
                gal = galsim.Convolve(psf, gal)

            # Use (explicit) stamps to draw
            stamp = gal.drawImage(wcs=self.localWCS, method='phot', offset=self.offset, save_photons=True)
            xmax = max(xmax, stamp.xmax)
            ymax = max(ymax, stamp.ymax)
            photons = stamp.photons
            photons.x += self.x_nominal
            photons.y += self.y_nominal
            photons_list.append(photons)

        # print('xmax = %d, ymax = %d '%(xmax, ymax))

        stamp = galsim.ImageF(int(xmax*1.1), int(ymax*1.1))
        stamp.wcs = self.localWCS
        stamp.setCenter(self.x_nominal, self.y_nominal)
        bounds = stamp.bounds & chip.img.bounds
        stamp[bounds] = chip.img[bounds]

        if not big_galaxy:
            for i in range(len(photons_list)):
                if i == 0:
                    chip.sensor.accumulate(photons_list[i], stamp)
                else:
                    chip.sensor.accumulate(photons_list[i], stamp, resume=True)
        else:
            sensor = galsim.Sensor()
            for i in range(len(photons_list)):
                if i == 0:
                    sensor.accumulate(photons_list[i], stamp)
                else:
                    sensor.accumulate(photons_list[i], stamp, resume=True)

        # print(stamp.array.sum())
        # chip.img[bounds] += stamp[bounds]
        chip.img[bounds] = stamp[bounds]
        # print("nphotons_sum = ", nphotons_sum)
        del photons_list
        del stamp
        return True, pos_shear

    def drawObj_slitless(self, tel, pos_img, psf_model, bandpass_list, filt, chip, nphotons_tot=None, g1=0, g2=0,
                         exptime=150., normFilter=None):

        norm_thr_rang_ids = normFilter['SENSITIVITY'] > 0.001
        sedNormFactor = getNormFactorForSpecWithABMAG(ABMag=self.param['mag_use_normal'], spectrum=self.sed,
                                                      norm_thr=normFilter,
                                                      sWave=np.floor(normFilter[norm_thr_rang_ids][0][0]),
                                                      eWave=np.ceil(normFilter[norm_thr_rang_ids][-1][0]))
        if sedNormFactor == 0:
            return False
        normalSED = Table(np.array([self.sed['WAVELENGTH'], self.sed['FLUX'] * sedNormFactor]).T,
                          names=('WAVELENGTH', 'FLUX'))

        big_galaxy = False
        if self.hlr_disk > 3.0: # Very big galaxy
            big_galaxy = True

        if self.getMagFilter(filt) <= 15 and (not big_galaxy):
            folding_threshold = 5.e-4
        else:
            folding_threshold = 5.e-3
        gsp = galsim.GSParams(folding_threshold=folding_threshold)
        # nphotons_sum = 0

        for i in range(len(bandpass_list)):
            bandpass = bandpass_list[i]

            psf, pos_shear = psf_model.get_PSF(chip=chip, pos_img=pos_img, bandpass=bandpass, folding_threshold=folding_threshold)
            disk = galsim.Sersic(n=1.0, half_light_radius=self.hlr_disk, flux=1.0, gsparams=gsp)
            disk_shape = galsim.Shear(g1=self.e1_disk, g2=self.e2_disk)
            disk = disk.shear(disk_shape)
            bulge = galsim.Sersic(n=4.0, half_light_radius=self.hlr_bulge, flux=1.0, gsparams=gsp)
            bulge_shape = galsim.Shear(g1=self.e1_bulge, g2=self.e2_bulge)
            bulge = bulge.shear(bulge_shape)

            gal = self.bfrac * bulge + (1.0 - self.bfrac) * disk
            gal = gal.withFlux(tel.pupil_area * exptime)
            gal_shear = galsim.Shear(g1=g1, g2=g2)
            gal = gal.shear(gal_shear)
            gal = galsim.Convolve(psf, gal)

            starImg = gal.drawImage(wcs=self.localWCS)

            origin_star = [self.y_nominal - (starImg.center.y - starImg.ymin),
                           self.x_nominal - (starImg.center.x - starImg.xmin)]

            # print(self.y_nominal, starImg.center.y, starImg.ymin)

            sdp = SpecDisperser(orig_img=starImg, xcenter=self.x_nominal-chip.bound.xmin,
                                ycenter=self.y_nominal-chip.bound.ymin, origin=origin_star,
                                tar_spec=normalSED,
                                band_start=bandpass.blue_limit * 10, band_end=bandpass.red_limit * 10,
                                conf=chip.sls_conf[1],
                                isAlongY=0)

            spec_orders = sdp.compute_spec_orders()
            for k, v in spec_orders.items():
                img_s = v[0]
                origin_order_x = v[1]
                origin_order_y = v[2]
                specImg = galsim.ImageF(img_s)
                photons = galsim.PhotonArray.makeFromImage(specImg)
                photons.x += origin_order_x
                photons.y += origin_order_y

                xlen_imf = int(specImg.xmax - specImg.xmin + 1)
                ylen_imf = int(specImg.ymax - specImg.ymin + 1)
                stamp = galsim.ImageF(xlen_imf, ylen_imf)
                stamp.wcs = self.localWCS
                stamp.setOrigin(origin_order_x, origin_order_y)
                bounds = stamp.bounds & chip.img.bounds
                if bounds.area() == 0:
                    continue
                stamp[bounds] = chip.img[bounds]

                if not big_galaxy:
                    chip.sensor.accumulate(photons, stamp)
                else:
                    sensor = galsim.Sensor()
                    sensor.accumulate(photons, stamp)

                # chip.sensor.accumulate(photons, stamp)
                chip.img[bounds] = stamp[bounds]
                del stamp
            del sdp
            del spec_orders
            del psf
        return True, pos_shear

    def getGSObj(self, psf, g1=0, g2=0, flux=None, filt=None, tel=None, exptime=150.):
        if flux == None:
            flux = self.getElectronFluxFilt(filt, tel, exptime)
        disk = galsim.Sersic(n=1.0, half_light_radius=self.hlr_disk, flux=1.0)
        disk_shape = galsim.Shear(g1=self.e1_disk, g2=self.e2_disk)
        disk = disk.shear(disk_shape)

        bulge = galsim.Sersic(n=4.0, half_light_radius=self.hlr_bulge, flux=1.0)
        bulge_shape = galsim.Shear(g1=self.e1_bulge, g2=self.e2_bulge)
        bulge = bulge.shear(bulge_shape)

        gal = self.bfrac * bulge + (1.0 - self.bfrac) * disk
        gal = gal.withFlux(flux)
        gal_shear = galsim.Shear(g1=g1, g2=g2)
        gal = gal.shear(gal_shear)
        final = galsim.Convolve(psf, gal)
        return final

    def rotateEllipticity(self, rotation):
        if rotation == 1:
            self.e1_disk, self.e2_disk, self.e1_bulge, self.e2_bulge, self.e1_total, self.e2_total = -self.e2_disk, self.e1_disk, -self.e2_bulge, self.e1_bulge, -self.e2_total, self.e1_total
        if rotation == 2:
            self.e1_disk, self.e2_disk, self.e1_bulge, self.e2_bulge, self.e1_total, self.e2_total = -self.e1_disk, -self.e2_disk, -self.e1_bulge, -self.e2_bulge, -self.e1_total, -self.e2_total
        if rotation == 3:
            self.e1_disk, self.e2_disk, self.e1_bulge, self.e2_bulge, self.e1_total, self.e2_total = self.e2_disk, -self.e1_disk, self.e2_bulge, -self.e1_bulge, self.e2_total, -self.e1_total

    def drawObject(self, img, final, noise_level=0.0, flux=None, filt=None, tel=None, exptime=150.):
        """ Override the method in parent class 
        Need to constrain the size of image stamp for extended objects
        """
        isUpdated = True
        if flux == None:
            flux = self.getElectronFluxFilt(filt, tel, exptime)
        stamp = final.drawImage(wcs=self.localWCS, offset=self.offset)
        stamp_arr = stamp.array
        mask = (stamp_arr >= 0.001*noise_level) # why 0.001?
        err = int(np.sqrt(mask.sum()))
        if np.mod(err, 2) == 1:
            err += 1
        # if err == 1:
        if err == 0:
            subSize = 16 # why 16?
        else:
            subSize = max([err, 16])
            fluxRatio = flux / stamp_arr[mask].sum()
            final = final.withScaledFlux(fluxRatio)

        imgSub = galsim.ImageF(subSize, subSize)

        # Draw with FFT
        # stamp = final.drawImage(image=imgSub, wcs=self.localWCS, offset=self.offset)

        # Draw with Photon Shoot
        stamp = final.drawImage(image=imgSub, wcs=self.localWCS, method='phot', offset=self.offset)
        
        stamp.setCenter(self.x_nominal, self.y_nominal)
        if np.sum(np.isnan(stamp.array)) >= 1:
            stamp.setZero()
        bounds = stamp.bounds & img.bounds
        if bounds.area() == 0:
            isUpdated = False
        else:
            img[bounds] += stamp[bounds]
        return img, stamp, isUpdated


    def getObservedEll(self, g1=0, g2=0):
        e1_obs, e2_obs, e_obs, theta = eObs(self.e1_total, self.e2_total, g1, g2)
        return self.e1_total, self.e2_total, g1, g2, e1_obs, e2_obs