made compatible with csst simultion v3.0

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

# Instrumental effects setting

# (NOTE) Here only used to construct

# ObservationSim.Instrument.Chip object

# observation_sim.instruments.chip object

# (TODO) Should readout from header

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

ins_effects:

import matplotlib.pyplot as plt

from astropy.io import ascii, fits

from cross_match_catalogs import read_catalog, match_catalogs_img

from ObservationSim.Instrument import Telescope, Filter, FilterParam

from observation_sim.instruments import Telescope, Filter, FilterParam

VC_A = 2.99792458e+18  # speed of light: A/s

#                         default='./workspace', help='output path')

#     return parser

def define_options():

    parser = argparse.ArgumentParser()

    parser.add_argument('--TU_catalog_list', dest='TU_catalog_list', type=str, required=True,

                        default='./workspace', help='output path')

    return parser

def getChipFilter(chipID, filter_layout=None):

    """Return the filter index and type for a given chip #(chipID)

    """

        filter_type_list = ["nuv","u", "g", "r", "i","z","y","GU", "GV", "GI", "FGS"]

    filter_type_list = ["nuv", "u", "g", "r",

                        "i", "z", "y", "GU", "GV", "GI", "FGS"]

    if filter_layout is not None:

        return filter_layout[chipID][0], filter_layout[chipID][1]

    # updated configurations

        if chipID>42 or chipID<1: raise ValueError("!!! Chip ID: [1,42]")

        if chipID in [6, 15, 16, 25]:  filter_type = "y"

        if chipID in [11, 20]:         filter_type = "z"

        if chipID in [7, 24]:          filter_type = "i"

        if chipID in [14, 17]:         filter_type = "u"

        if chipID in [9, 22]:          filter_type = "r"

        if chipID in [12, 13, 18, 19]: filter_type = "nuv"

        if chipID in [8, 23]:          filter_type = "g"

        if chipID in [1, 10, 21, 30]:  filter_type = "GI"

        if chipID in [2, 5, 26, 29]:   filter_type = "GV"

        if chipID in [3, 4, 27, 28]:   filter_type = "GU"

        if chipID in range(31, 43):    filter_type = 'FGS'

    if chipID > 42 or chipID < 1:

        raise ValueError("!!! Chip ID: [1,42]")

    if chipID in [6, 15, 16, 25]:

        filter_type = "y"

    if chipID in [11, 20]:

        filter_type = "z"

    if chipID in [7, 24]:

        filter_type = "i"

    if chipID in [14, 17]:

        filter_type = "u"

    if chipID in [9, 22]:

        filter_type = "r"

    if chipID in [12, 13, 18, 19]:

        filter_type = "nuv"

    if chipID in [8, 23]:

        filter_type = "g"

    if chipID in [1, 10, 21, 30]:

        filter_type = "GI"

    if chipID in [2, 5, 26, 29]:

        filter_type = "GV"

    if chipID in [3, 4, 27, 28]:

        filter_type = "GU"

    if chipID in range(31, 43):

        filter_type = 'FGS'

    filter_id = filter_type_list.index(filter_type)

    return filter_id, filter_type

def magToFlux(mag):

    """

    flux of a given AB magnitude

    flux = 10**(-0.4*(mag+48.6))

    return flux

def getElectronFluxFilt(mag, filt, tel, exptime=150.):

    photonEnergy = filt.getPhotonE()

    flux = magToFlux(mag)

    factor = 1.0e4 * flux/photonEnergy * VC_A * (1.0/filt.blue_limit - 1.0/filt.red_limit)

    factor = 1.0e4 * flux/photonEnergy * VC_A * \

        (1.0/filt.blue_limit - 1.0/filt.red_limit)

    return factor * filt.efficiency * tel.pupil_area * exptime

def convert_catalog(catname):

    data_dir = os.path.dirname(catname)

    base_name = os.path.basename(catname)

    fits_filename = os.path.join(data_dir, base_name + '.fits')

    text_file.write(fits_filename, overwrite=True)

def validation_hist(val, idx, name="val", nbins=10, bins=None, fig_name='detected_counts.png', output_dir='./', create_figure=True):

    if bins is None:

        counts, bins = np.histogram(val, bins=nbins)

    else:

        counts_detected, _ = np.histogram(val[~is_empty], bins=bins)

    if create_figure:

        create_hist_figure(counts, counts_detected, bins, name, output_dir, fig_name)

        create_hist_figure(counts, counts_detected, bins,

                           name, output_dir, fig_name)

    return counts, counts_detected, bins

def create_hist_figure(counts, counts_detected, bins, name="val", output_dir='./', fig_name='detected_counts.png'):

    plt.figure()

    plt.stairs(counts, bins, color='r', label='TU objects')

    fig_name = os.path.join(output_dir, fig_name)

    plt.savefig(fig_name)

def hist_fraction(val, idx, name='val', nbins=10, bins=None, output_dir='./', fig_name="completeness_fraction.png"):

    if bins is None:

        counts, bins = np.histogram(val, bins=nbins)

    plt.savefig(fig_name)

    return fraction

def create_fraction_figure(counts, counts_detected, bins, name='val', output_dir='./', fig_name="completeness_fraction.png"):

    fraction = counts_detected / counts

    fraction[np.where(np.isnan(fraction))[0]] = 0.

    plt.savefig(fig_name)

    return fraction

def calculate_fraction(TU_catalog, source_catalog, output_dir, nbins=10):

    convert_catalog(TU_catalog)

    x_TU, y_TU, col_list = read_catalog(TU_catalog + '.fits', ext_num=1, ra_name="xImage", dec_name="yImage", col_list=["mag"])

    x_TU, y_TU, col_list = read_catalog(

        TU_catalog + '.fits', ext_num=1, ra_name="xImage", dec_name="yImage", col_list=["mag"])

    mag_TU = col_list[0]

    x_source, y_source, _ = read_catalog(source_catalog, ext_num=1, ra_name="X_IMAGE", dec_name="Y_IMAGE")

    idx1, idx2, = match_catalogs_img(x1=x_TU, y1=y_TU, x2=x_source, y2=y_source)

    counts, bins = validation_hist(val=mag_TU, idx=idx1, name="mag_injected", output_dir=output_dir)

    fraction = hist_fraction(val=mag_TU, idx=idx1, name="mag_injected", nbins=10, output_dir=output_dir)

    x_source, y_source, _ = read_catalog(

        source_catalog, ext_num=1, ra_name="X_IMAGE", dec_name="Y_IMAGE")

    idx1, idx2, = match_catalogs_img(

        x1=x_TU, y1=y_TU, x2=x_source, y2=y_source)

    counts, bins = validation_hist(

        val=mag_TU, idx=idx1, name="mag_injected", output_dir=output_dir)

    fraction = hist_fraction(

        val=mag_TU, idx=idx1, name="mag_injected", nbins=10, output_dir=output_dir)

    return counts, bins, fraction

def calculate_fraction_multi_cats(TU_catalog_list, source_catalog_list, output_dir, nbins=10):

    counts = np.zeros(nbins)

    counts_detected = np.zeros(nbins)

        TU_catalog = TU_catalog_list[i]

        source_catalog = source_catalog_list[i]

        convert_catalog(TU_catalog)

        x_TU_temp, y_TU_temp, col_list = read_catalog(TU_catalog + '.fits', ext_num=1, ra_name="xImage", dec_name="yImage", col_list=["mag"])

        x_TU_temp, y_TU_temp, col_list = read_catalog(

            TU_catalog + '.fits', ext_num=1, ra_name="xImage", dec_name="yImage", col_list=["mag"])

        mag_TU_temp = col_list[0]

        x_source_temp, y_source_temp, _ = read_catalog(source_catalog, ext_num=1, ra_name="X_IMAGE", dec_name="Y_IMAGE")

        idx1, _, = match_catalogs_img(x1=x_TU_temp, y1=y_TU_temp, x2=x_source_temp, y2=y_source_temp)

        counts_temp, counts_detected_temp, _ = validation_hist(val=mag_TU_temp, idx=idx1, name="mag_injected", bins=bins, output_dir=output_dir, create_figure=False)

        x_source_temp, y_source_temp, _ = read_catalog(

            source_catalog, ext_num=1, ra_name="X_IMAGE", dec_name="Y_IMAGE")

        idx1, _, = match_catalogs_img(

            x1=x_TU_temp, y1=y_TU_temp, x2=x_source_temp, y2=y_source_temp)

        counts_temp, counts_detected_temp, _ = validation_hist(

            val=mag_TU_temp, idx=idx1, name="mag_injected", bins=bins, output_dir=output_dir, create_figure=False)

        counts += counts_temp

        counts_detected += counts_detected_temp

    create_hist_figure(counts, counts_detected, bins, "mag_injected", output_dir)

    fraction = create_fraction_figure(counts, counts_detected, bins, 'mag_injected', output_dir)

    create_hist_figure(counts, counts_detected, bins,

                       "mag_injected", output_dir)

    fraction = create_fraction_figure(

        counts, counts_detected, bins, 'mag_injected', output_dir)

    return counts, counts_detected, bins, fraction

def calculate_undetected_flux(orig_cat, mag_bins, fraction, mag_low=20.0, mag_high=26.0, image=None,  output_dir='./'):

    # Get info from original image

    hdu = fits.open(image)

    hdu.close()

    # Get info from original catalog

    ra_orig, dec_orig, col_list_orig = read_catalog(orig_cat, ra_name='RA', dec_name='DEC', col_list=['Mag_Kron'])

    ra_orig, dec_orig, col_list_orig = read_catalog(

        orig_cat, ra_name='RA', dec_name='DEC', col_list=['Mag_Kron'])

    mag_orig = col_list_orig[0]

    nbins = len(mag_bins) - 1

    counts, _ = np.histogram(mag_orig, bins=nbins)

    for i in range(len(mags)):

        if mags[i] < mag_low or mags[i] > mag_high:

            continue

        flux_electrons = counts_missing[i] * getElectronFluxFilt(mag=mags[i], filt=filt, tel=tel)

        flux_electrons = counts_missing[i] * \

            getElectronFluxFilt(mag=mags[i], filt=filt, tel=tel)

        undetected_flux += flux_electrons

    undetected_flux /= (float(nx_pix) * float(ny_pix))

#     )

#     print(undetected_flux)

if __name__ == "__main__":

    args = define_options().parse_args()

    with open(args.TU_catalog_list) as file:

from astropy.wcs import WCS

from astropy.io import ascii, fits

from matplotlib.colors import LogNorm

from scipy.stats import binned_statistic

from astropy.visualization import ZScaleInterval

import os

import numpy as np

import matplotlib.pyplot as plt

from cross_match_catalogs import match_catalogs_img

def plot_injection_comparison(orig_img, injected_img, flg_img=None, save_fig_dir=None, fig_prefix="", figsize=(12, 8)):

    z = ZScaleInterval()

    plt.figure(figsize=figsize, dpi=100)

    hdu_orig = fits.open(orig_img)[1]

    data_orig = hdu_orig.data

    if flg_img is not None:

        flg_data = fits.getdata(flg_img)

        data_orig[flg_data > 0] = 0.

    wcs = WCS(hdu_orig.header)

    plt.subplot(projection=wcs)

    z1, z2 = z.get_limits(data_orig)

    plt.imshow(data_orig, origin='lower', cmap='gray', vmin=z1, vmax=z2)

    plt.grid(color='white', ls='solid')

    if save_fig_dir is not None:

        output_filename = fig_prefix + "original_img.png"

        output_img_path = os.path.join(save_fig_dir, output_filename)

        plt.savefig(output_img_path)

    plt.show()

    plt.figure(figsize=figsize, dpi=100)

    hdu_inj = fits.open(injected_img)[1]

    data_inj = hdu_inj.data

    if flg_img is not None:

        data_inj[flg_data > 0] = 0.

    wcs = WCS(hdu_inj.header)

    plt.subplot(projection=wcs)

    z1, z2 = z.get_limits(data_inj)

    plt.imshow(data_inj, origin='lower', cmap='gray', vmin=z1, vmax=z2)

    plt.grid(color='white', ls='solid')

    if save_fig_dir is not None:

        output_filename = fig_prefix + "injected_img.png"

        output_img_path = os.path.join(save_fig_dir, output_filename)

        plt.savefig(output_img_path)

    plt.show()

    plt.figure(figsize=figsize, dpi=100)

    plt.subplot(projection=wcs)

    img_diff = data_inj - data_orig

    z1 = 0.

    z2 = 0.001

    plt.imshow(img_diff, origin='lower', cmap='gray', vmin=z1, vmax=z2)

    plt.grid(color='white', ls='solid')

    if save_fig_dir is not None:

        output_filename = fig_prefix + "diff_img.png"

        output_img_path = os.path.join(save_fig_dir, output_filename)

        plt.savefig(output_img_path)

    plt.show()

def plot_ensemble_hist(cat_path_list, column_name="Mag_Kron", column_unit="mag", title="Total KRON MAG distribution", save_fig_dir=None, fig_prefix="",

                       nbins=50, low=16., high=28., density=False):

    values = []

    bins = np.linspace(low, high, nbins+1)

    for cat_path in cat_path_list:

        if cat_path.endswith(".fits"):

            hdu = fits.open(cat_path)

            value_temp = hdu[1].data[column_name]

        elif cat_path.endswith(".cat"):

            data = ascii.read(cat_path)

            value_temp = data[column_name]

        print("number of objects in %s: %d" %

              (os.path.basename(cat_path), len(value_temp)))

        values = np.append(values, value_temp)

    plt.figure()

    plt.hist(values, bins=bins, density=density)

    plt.xlabel(column_name + '/' + column_unit, size='x-large')

    if density is False:

        plt.ylabel("Counts", size='x-large')

    plt.title(title, size='x-large')

    if save_fig_dir is not None:

        output_filename = fig_prefix + "%s_ensemble_hist.png" % (column_name)

        output_img_path = os.path.join(save_fig_dir, output_filename)

        plt.savefig(output_img_path)

    plt.show()

def create_hist_figure(counts, counts_detected, bins, name="val", output_dir='./', fig_name='detected_counts.png', save_figure=False, title=None):

    fig = plt.figure()

    ax = fig.add_subplot(111)

    ax.set_xlabel(name, size='x-large')

    ax.set_ylabel("Counts", size='x-large')

    if title is not None:

        ax.set_title(title, size='x-large')

    ax.stairs(counts, bins, color='r', label='TU objects')

    ax.stairs(counts_detected, bins, color='g', label='Detected')

    ax.legend(loc='upper right', fancybox=True)

    if save_figure:

        fig_name = os.path.join(output_dir, fig_name)

        fig.savefig(fig_name)

    return fig, ax

def create_fraction_figure(counts, counts_detected, bins, name='val', output_dir='./', fig_name="completeness_fraction.png",

                           save_figure=False, title=None, figure=None, color='r', label='patch_1', show_legend=False):

    fraction = counts_detected / counts

    fraction[np.where(np.isnan(fraction))[0]] = 0.

    if figure is not None:

        fig = figure

        ax = fig.axes[0]

        ax.stairs(fraction, bins, color=color, label=label)

        if title is not None:

            ax.set_title(title, size='x-large')

    else:

        fig = plt.figure()

        ax = fig.add_subplot(111)

        ax.stairs(fraction, bins, color=color, label=label)

        ax.set_xlabel(name, size='x-large')

        if title is not None:

            ax.set_title(title, size='x-large')

        else:

            ax.set_title("Completeness Fraction")

    if show_legend:

        ax.legend(loc='upper right', fancybox=True)

    if save_figure:

        fig_name = os.path.join(output_dir, fig_name)

        fig.savefig(fig_name)

    return fig, ax, fraction

def validation_hist(val, idx, name="val", nbins=10, bins=None, fig_name='detected_counts.png', output_dir='./', create_figure=True):

    if bins is None:

        counts, bins = np.histogram(val, bins=nbins)

    else:

        counts, bins = np.histogram(val, bins=bins)

    is_empty = np.full(len(val), False)

    for i in range(len(idx)):

        if idx[i].size == 0:

            is_empty[i] = True

    if bins is None:

        counts_detected, _ = np.histogram(val[~is_empty], bins=nbins)

    else:

        counts_detected, _ = np.histogram(val[~is_empty], bins=bins)

    if create_figure:

        create_hist_figure(counts, counts_detected, bins,

                           name, output_dir, fig_name)

    return counts, counts_detected, bins

def plot_mag_comparison(truth_cat_list, measured_cat_root_dir, mag1_name="mag", mag2_name="Mag_Kron", save_fig_dir=None, fig_prefix="",

                        nbins=20, low=18., high=26., ylim=[-1., 1.], title=None):

    diff_list = []

    truth_list = []

    bins = np.linspace(low, high, nbins+1)

    for cat_path_truth in truth_cat_list:

        print("Injected truth catalog: ",

              os.path.basename(cat_path_truth))

        obs_id = cat_path_truth.split('/')[-2]

        # Read truth catalog

        data = ascii.read(cat_path_truth)

        x_truth = data["xImage"]

        y_truth = data["yImage"]

        mag_truth = data[mag1_name]

        # Read measured catalog

        cat_path_measured = os.path.join(measured_cat_root_dir,

                                         obs_id,

                                         os.path.basename(cat_path_truth).replace("img", "cat").replace(".cat", ".fits"))

        print("L1 processed photometry catalog: ",

              os.path.basename(cat_path_truth))

        hdu = fits.open(cat_path_measured)

        x_measure = hdu[1].data["X"]

        y_measure = hdu[1].data["Y"]

        mag_measure = hdu[1].data["Mag_Kron"]

        # Match measured objects vs truth

        idx1, _, = match_catalogs_img(

            x1=x_truth, y1=y_truth, x2=x_measure, y2=y_measure)

        for i in range(len(idx1)):

            if idx1[i].size == 0:

                continue

            else:

                diff_list.append(mag_measure[idx1[i][0]] - mag_truth[i])

                truth_list.append(mag_truth[i])

    bin_means, bin_edges, binnumber = binned_statistic(truth_list, diff_list, 'mean',

                                                       bins=nbins, range=[low, high])

    bin_median, bin_edges, binnumber = binned_statistic(truth_list, diff_list, 'median',

                                                        bins=nbins, range=[low, high])

    bin_std, bin_edges, binnumber = binned_statistic(truth_list, diff_list, 'std',

                                                     bins=nbins, range=[low, high])

    bin_width = (bin_edges[1] - bin_edges[0])

    bin_centers = bin_edges[1:] - bin_width/2

    plt.figure()

    plt.plot(truth_list, diff_list, 'ro', alpha=0.1)

    plt.axhline(y=0., color='k', alpha=0.6)

    plt.plot(bin_centers, bin_median, '--',

             label=r'$\rm{median}\ \Delta mag$', alpha=0.6)

    plt.errorbar(bin_centers, bin_means, yerr=bin_std, fmt='bo', capsize=2,

                 label=r'$\rm{mean}\ \Delta mag$', alpha=0.6)

    plt.xlim([low, high])

    plt.ylim(ylim)

    plt.xlabel("True mag", size='x-large')

    plt.ylabel("Measured (Kron) - True mag", size='x-large')

    plt.legend(loc='upper left', fancybox=True)

    if title is not None:

        plt.title(title, size='x-large')

    if save_fig_dir is not None:

        output_filename = fig_prefix + "measured-true_mag.png"

        output_img_path = os.path.join(save_fig_dir, output_filename)

        plt.savefig(output_img_path)

    plt.show()

from scipy import interpolate

from datetime import datetime

from ObservationSim.MockObject import CatalogBase, Star, Galaxy, Quasar

from ObservationSim.MockObject._util import tag_sed, getObservedSED, getABMAG, integrate_sed_bandpass, comoving_dist

from ObservationSim.Astrometry.Astrometry_util import on_orbit_obs_position

from observation_sim.mock_objects import CatalogBase, Star, Galaxy, Quasar

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

# (TEST)

from astropy.cosmology import FlatLambdaCDM

        # Load how mnay objects?

        max_ngals = len(gals['ra'])

        remain = nobjects

        for igals in range(max_ngals):

        for i in range(max_ngals):

            if remain == 0:

                break

            igals = random.randint(0, len(gals['ra']))

            param = self.initialize_param()

            param['ra'] = gals['ra'][igals]

            param['dec'] = gals['dec'][igals]

            param['e1'] = gals['ellipticity_true'][igals][0]

            param['e2'] = gals['ellipticity_true'][igals][1]

            # For shape calculation

            # For shape calculation

            param['e1'], param['e2'], param['ell_total'] = self.rotate_ellipticity(

                e1=gals['ellipticity_true'][igals][0],