add module to output PSF

import os

import galsim

import numpy as np

import astropy.constants as cons

from astropy import wcs

from astropy.table import Table

import astropy.io.fits as fitsio

from ObservationSim.MockObject._util import magToFlux, VC_A, convolveGaussXorders, convolveImg

from ObservationSim.MockObject._util import integrate_sed_bandpass, getNormFactorForSpecWithABMAG, getObservedSED, \

        sig_obj = np.std(stamp.array)

        snr_obj = img_flux / sig_obj

        return snr_obj

    def drawObj_PSF(self, tel, pos_img, psf_model, bandpass_list, filt, chip, nphotons_tot=None, g1=0, g2=0,

                          exptime=150., fd_shear=None, chip_output=None):

        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)

            if self.logger:

                self.logger.error(e)

            return 2, None

        # Set Galsim Parameters

        if self.getMagFilter(filt) <= 15:

            folding_threshold = 5.e-4

        else:

            folding_threshold = 5.e-3

        gsp = galsim.GSParams(folding_threshold=folding_threshold)

        # Get real image position of object (deal with chip rotation w.r.t its center)

        self.real_pos = self.getRealPos(chip.img, global_x=self.posImg.x, global_y=self.posImg.y,

                                        img_real_wcs=self.chip_wcs)

        x, y = self.real_pos.x + 0.5, self.real_pos.y + 0.5

        x_nominal = int(np.floor(x + 0.5))

        y_nominal = int(np.floor(y + 0.5))

        dx = x - x_nominal

        dy = y - y_nominal

        offset = galsim.PositionD(dx, dy)

        # Get real local wcs of object (deal with chip rotation w.r.t its center)

        chip_wcs_local = self.chip_wcs.local(self.real_pos)

        is_updated = 0

        # Loop over all sub-bandpasses

        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)

                if self.logger:

                    self.logger.error(e)

                continue

            ratio = sub / full

            if not (ratio == -1 or (ratio != ratio)):

                nphotons = ratio * nphotons_tot

            else:

                continue

            # Get PSF model

            psf, pos_shear = psf_model.get_PSF(chip=chip, pos_img=pos_img, bandpass=bandpass,

                                               folding_threshold=folding_threshold)

            star_temp = psf.withFlux(nphotons)

            if i==0:

                star = star_temp

            else:

                star = star+star_temp

        pixelScale = 0.074

        stamp = star.drawImage(wcs=chip_wcs_local, offset=offset)

        #stamp = star.drawImage(nx=256, ny=256, scale=pixelScale)

        if np.sum(np.isnan(stamp.array)) > 0:

            return None

        fn = chip_output.subdir + "/psfIDW"

        os.makedirs(fn, exist_ok=True)

        fn = fn + "/ccd_{:}".format(chip.chipID)+"_psf_"+str(self.param['id'])+".fits"

        if fn != None:

            if os.path.exists(fn):

                os.remove(fn)

        hdu = fitsio.PrimaryHDU()

        hdu.data = stamp.array

        hdu.header.set('name',      self.type)

        hdu.header.set('pixScale',  pixelScale)

        hdu.header.set('objID',     self.param['id'])

        hdu.writeto(fn)

        del stamp

        return None

                            normFilter=norm_filt,

                            fd_shear=fd_shear)

                    if isUpdated == 1 and self.config["run_option"]["out_psf"]:

                        obj.drawObj_PSF(

                            tel=self.tel,

                            pos_img=pos_img,

                            psf_model=psf_model,

                            bandpass_list=filt.bandpass_sub_list,

                            filt=filt,

                            chip=chip,

                            g1=obj.g1,

                            g2=obj.g2,

                            exptime=pointing.exp_time,

                            fd_shear=fd_shear,

                            chip_output=chip_output)

                    if isUpdated == 1:

                        # TODO: add up stats

                        chip_output.cat_add_obj(obj, pos_img, pos_shear)

# OK to pass either way or both, as long as they are consistent

work_dir: "/share/simudata/CSSOSDataProductsSims/data/CSSTSimImage_C8/"

data_dir: "/share/simudata/CSSOSDataProductsSims/data/"

run_name: "testRun_fits"

run_name: "testRun_FGS"

project_cycle: 6

run_counter: 1

  # If yes, no imaging simulation will run

  out_cat_only: NO

  #output PSF

  out_psf: YES

###############################################

# Catalog setting

###############################################

  galaxy_only: NO

  # With stamp?

  stamp_yes: YES

  stamp_yes: NO

  # rotate galaxy ellipticity

  rotateEll: 0. # [degree]

  # - give a list of indexes of chips: [ip_1, ip_2...]

  # - run all chips: null

  # Note: for all pointings

  run_chips: [1, 2, 3, 7, 8, 9]

  run_chips: [39,40,41,42]

  # Whether to enable astrometric modeling

  enable_astrometric_model: True

  add_bias:       YES  # Whether to add bias-level to images

  bias_16channel: YES  # Whether to add different biases for 16 channels

  gain_16channel: YES  # Whether to make different gains for 16 channels

  shutter_effect: YES  # Whether to add shutter effect

  shutter_effect: NO   # Whether to add shutter effect

  flat_fielding:  YES  # Whether to add flat-fielding effect

  prnu_effect:    YES  # Whether to add PRNU effect

  non_linear:     YES  # Whether to add non-linearity