Chromatic Galaxy Image Simulations

main.py

import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import Divider, Size
import numpy as np
from scipy import stats

import galsim

import dsps
from dsps.data_loaders import load_ssp_templates
from lsstdesc_diffsky import read_diffskypop_params
from lsstdesc_diffsky.defaults import OUTER_RIM_COSMO_PARAMS
from lsstdesc_diffsky.io_utils import load_healpixel
from lsstdesc_diffsky.io_utils import load_diffsky_params
from lsstdesc_diffsky.sed import calc_rest_sed_disk_bulge_knot_galpop


def make_gal(
    igal,
    mock,
    ssp_data,
    disk_bulge_sed_info,
    cosmo_params,
    n_knots=0,
    morphology="chromatic",
):
    bulge_ellipticity = galsim.Shear(
        g1=mock["spheroidEllipticity1"][igal],
        g2=mock["spheroidEllipticity2"][igal],
    )
    bulge = galsim.DeVaucouleurs(
        half_light_radius=mock["spheroidHalfLightRadiusArcsec"][igal],
    ).shear(bulge_ellipticity)

    disk_ellipticity = galsim.Shear(
        g1=mock["diskEllipticity1"][igal],
        g2=mock["diskEllipticity2"][igal],
    )
    disk = galsim.Exponential(
        half_light_radius=mock["diskHalfLightRadiusArcsec"][igal],
    ).shear(disk_ellipticity)

    if n_knots > 0:
        knots = galsim.RandomKnots(
            n_knots,
            profile=disk,
        )

    redshift = mock["redshift"][igal]
    luminosity_distance = dsps.cosmology.luminosity_distance_to_z(
        redshift,
        cosmo_params.Om0,
        cosmo_params.w0,
        cosmo_params.wa,
        cosmo_params.h,
    )
    luminosity_area = 4 * np.pi * luminosity_distance**2

    wave_type = "angstrom"
    flux_type = "fnu"
    # DSPS units
    # import astropy.units as u
    # _wave_type = u.angstrom
    # _flux_type = u.Lsun / u.Hz / u.Mpc**2
    # flux_factor = (1 * _flux_type).to(galsim.SED._fnu).value
    flux_factor = 4.0204145742268754e-16

    bulge_sed_table = galsim.LookupTable(
        ssp_data.ssp_wave,
        disk_bulge_sed_info.rest_sed_bulge[igal] / luminosity_area * flux_factor,
    )
    bulge_sed = galsim.SED(bulge_sed_table, wave_type=wave_type, flux_type=flux_type)

    disk_sed_table = galsim.LookupTable(
        ssp_data.ssp_wave,
        disk_bulge_sed_info.rest_sed_diffuse_disk[igal] / luminosity_area * flux_factor,
    )
    disk_sed = galsim.SED(disk_sed_table, wave_type=wave_type, flux_type=flux_type)

    knot_sed_table = galsim.LookupTable(
        ssp_data.ssp_wave,
        disk_bulge_sed_info.rest_sed_knot[igal] / luminosity_area * flux_factor,
    )
    knot_sed = galsim.SED(knot_sed_table, wave_type=wave_type, flux_type=flux_type)

    match morphology:
        case "chromatic":
            if n_knots > 0:
                gal = bulge * bulge_sed + disk * disk_sed + knots * knot_sed
            else:
                # if no knots are drawn, it is important to reweight the disk
                # to preserve relative magnitudes
                gal = bulge * bulge_sed + disk * (disk_sed + knot_sed)
        case "achromatic":
            # decouple the morphology from the spectra
            m_total = disk_bulge_sed_info.mstar_total[igal]
            m_bulge = disk_bulge_sed_info.mstar_bulge[igal]
            m_disk = disk_bulge_sed_info.mstar_diffuse_disk[igal]
            m_knot = disk_bulge_sed_info.mstar_knot[igal]

            if n_knots > 0:
                gal = (
                    (bulge * m_bulge + disk * m_disk + knots * m_knot) / m_total \
                    * (bulge_sed + disk_sed + knot_sed)
                )
            else:
                gal = (
                    (bulge * m_bulge + disk * (m_disk + m_knot)) / m_total \
                    * (bulge_sed + disk_sed + knot_sed)
                )
        case _:
            raise ValueError("Unrecognized morphology: %s" % morphology)

    gal = gal.atRedshift(redshift)

    return gal


if __name__ == "__main__":

    rng = np.random.default_rng(13720)

    pixel_scale = 0.2

    filters = {"u", "g", "r", "i", "z", "y"}
    bps = {
        f: galsim.Bandpass(f"LSST_{f}.dat", "nm").withZeropoint("AB")
        for f in filters
    }

    fn = "/opt/diffsky.testdata.hdf5"
    patlist = ('LSST', 'disk', 'spheroid', 'knot')
    mock, metadata = load_healpixel(fn, patlist)
    diffsky_param_data = load_diffsky_params(mock)

    all_diffskypop_params = read_diffskypop_params("roman_rubin_2023")
    print(all_diffskypop_params._fields)

    ssp_data = load_ssp_templates(fn='/opt/dsps_ssp_data.h5')
    print(ssp_data._fields)

    print('\nssp_lgmet.shape = {}'.format(ssp_data.ssp_lgmet.shape))
    print('ssp_lg_age_gyr.shape = {}'.format(ssp_data.ssp_lg_age_gyr.shape))
    print('ssp_wave.shape = {}'.format(ssp_data.ssp_wave.shape))
    print('ssp_flux.shape = {}'.format(ssp_data.ssp_flux.shape))

    args = (
        mock['redshift'],
        diffsky_param_data.mah_params,
        diffsky_param_data.ms_params,
        diffsky_param_data.q_params,
        diffsky_param_data.fbulge_params,
        diffsky_param_data.fknot,
        *ssp_data,
        *all_diffskypop_params,
        OUTER_RIM_COSMO_PARAMS
    )

    disk_bulge_sed_info = calc_rest_sed_disk_bulge_knot_galpop(*args)

    gi_mock = []
    gi_gal = []
    for igal in range(len(mock["redshift"])):

        morphology = "achromatic"

        # poisson mean of 20 suggested by Mike Jarvis
        # n_knots = np.clip(rng.poisson(20), 1, None)
        n_knots = 0

        gal = make_gal(
            igal,
            mock,
            ssp_data,
            disk_bulge_sed_info,
            OUTER_RIM_COSMO_PARAMS,
            n_knots=n_knots,
            morphology=morphology,
        )
        gal = gal.rotate(rng.uniform(0, 180) * galsim.degrees)

        print("---      gal: %d" % igal)
        print("mock     %f" % mock["LSST_obs_g_nodust"][igal])
        print("galsim   %f [bulge + disk + knots]" % gal.calculateMagnitude(bps["g"]))

        base_psf = galsim.Gaussian(fwhm=0.7)
        psf = galsim.ChromaticAtmosphere(
            base_psf,
            500,
            alpha=-0.3,
            zenith_angle=0 * galsim.degrees,
            parallactic_angle=0 * galsim.degrees,
        )

        mag_u = gal.calculateMagnitude(bps["u"])
        mag_g = gal.calculateMagnitude(bps["g"])
        mag_r = gal.calculateMagnitude(bps["r"])
        mag_i = gal.calculateMagnitude(bps["i"])
        mag_z = gal.calculateMagnitude(bps["z"])
        mag_y = gal.calculateMagnitude(bps["y"])

        gi_mock.append(mock["LSST_obs_g"][igal] - mock["LSST_obs_i"][igal])
        gi_gal.append(mag_g - mag_i)

        g = galsim.Convolve(gal, psf).drawImage(nx=32, ny=32, scale=pixel_scale, bandpass=bps["g"])
        r = galsim.Convolve(gal, psf).drawImage(nx=32, ny=32, scale=pixel_scale, bandpass=bps["r"])
        i = galsim.Convolve(gal, psf).drawImage(nx=32, ny=32, scale=pixel_scale, bandpass=bps["i"])
        norm = np.sqrt(np.square(np.max(g.array)) + np.square(np.max(r.array)) + np.square(np.max(i.array)))

        fig = plt.figure(figsize=(8.5, 5.5))
        h = [Size.Fixed(1), Size.Fixed(2), Size.Fixed(0.5), Size.Fixed(2), Size.Fixed(0.5), Size.Fixed(2), Size.Fixed(0.5)]
        v = [Size.Fixed(0.5), Size.Fixed(2), Size.Fixed(0.5), Size.Fixed(2), Size.Fixed(0.5)]
        divider = Divider(fig, (0, 0, 1, 1), h, v, aspect=False)

        ax = fig.add_axes(
            divider.get_position(),
            axes_locator=divider.new_locator(nx=1, ny=3)
        )
        ax.imshow(g.array / norm, origin="lower", vmin=0, vmax=1)
        ax.text(0.05, 0.95, f"{mag_g:.2f} mag", verticalalignment='top', horizontalalignment='left', transform=ax.transAxes)
        ax.set_title("$g$")

        ax = fig.add_axes(
            divider.get_position(),
            axes_locator=divider.new_locator(nx=3, ny=3)
        )
        ax.imshow(r.array / norm, origin="lower", vmin=0, vmax=1)
        ax.text(0.05, 0.95, f"{mag_r:.2f} mag", verticalalignment='top', horizontalalignment='left', transform=ax.transAxes)
        ax.set_title("$r$")

        ax = fig.add_axes(
            divider.get_position(),
            axes_locator=divider.new_locator(nx=5, ny=3)
        )
        ax.imshow(i.array / norm, origin="lower", vmin=0, vmax=1)
        ax.text(0.05, 0.95, f"{mag_i:.2f} mag", verticalalignment='top', horizontalalignment='left', transform=ax.transAxes)
        ax.set_title("$i$")

        ax = fig.add_axes(
            divider.get_position(),
            axes_locator=divider.new_locator(nx=1, nx1=6, ny=1)
        )
        xs = ssp_data.ssp_wave / 10
        xs = xs[(xs > gal.sed.blue_limit) & (xs < gal.sed.red_limit)]
        ax.set_xlim(200, 1200)
        ax.set_yscale("log")
        ax.plot(xs, bps["u"](xs) * np.trapz(gal.sed(xs), x=xs), c="lightgrey")
        ax.plot(xs, bps["g"](xs) * np.trapz(gal.sed(xs), x=xs), c="lightgrey")
        ax.plot(xs, bps["r"](xs) * np.trapz(gal.sed(xs), x=xs), c="lightgrey")
        ax.plot(xs, bps["i"](xs) * np.trapz(gal.sed(xs), x=xs), c="lightgrey")
        ax.plot(xs, bps["z"](xs) * np.trapz(gal.sed(xs), x=xs), c="lightgrey")
        ax.plot(xs, bps["y"](xs) * np.trapz(gal.sed(xs), x=xs), c="lightgrey")
        ax.plot(xs, gal.sed(xs), c="k")
        ax.text(0.05, 0.95, f"$z = {{{gal.redshift:.2f}}}$", verticalalignment='top', horizontalalignment='left', transform=ax.transAxes)
        ax.set_xlabel(r"$\lambda$ [nm]")
        ax.set_ylabel(r"$f_{\rm photons}$ [photons/nm/cm$^2$/s]")

        plt.savefig(f"img-{igal}.png", dpi=200)
        plt.close()

    fig = plt.figure(figsize=(6, 3))
    h = [Size.Fixed(0.5), Size.Fixed(2), Size.Fixed(1), Size.Fixed(2), Size.Fixed(0.5)]
    v = [Size.Fixed(0.5), Size.Fixed(2), Size.Fixed(0.5)]
    divider = Divider(fig, (0, 0, 1, 1), h, v, aspect=False)

    ax = fig.add_axes(
        divider.get_position(),
        axes_locator=divider.new_locator(nx=1, ny=1)
    )

    ksp = stats.kstest(gi_gal, gi_mock).pvalue
    bins = np.linspace(0, 2, 10)
    ax.hist(gi_mock, bins=bins, histtype="step", ec="k", label="mock")
    ax.hist(gi_gal, bins=bins, histtype="step", ec="b", label="galsim")
    ax.legend(loc='upper right')
    ax.set_xlabel("$g - i$")
    ax.text(0.05, 0.95, f"$p = {{{ksp:.2f}}}$", verticalalignment='top', horizontalalignment='left', transform=ax.transAxes)

    ax = fig.add_axes(
        divider.get_position(),
        axes_locator=divider.new_locator(nx=3, ny=1)
    )

    pearsonr = stats.pearsonr(gi_mock, gi_gal).statistic
    ax.plot([0, 2], [0, 2], c="lightgrey", linestyle="--", zorder=0)
    ax.scatter(gi_mock, gi_gal, c="k")
    ax.text(0.05, 0.95, f"$r = {{{pearsonr:.2f}}}$", verticalalignment='top', horizontalalignment='left', transform=ax.transAxes)
    ax.set_xlabel("$g - i$ [mock]")
    ax.set_ylabel("$g - i$ [galsim]")
    ax.set_xlim(0, 2)
    ax.set_ylim(0, 2)

    plt.savefig(f"hist.png", dpi=200)
    plt.close()

test.py

import numpy as np

import galsim

import dsps
from dsps.data_loaders import load_ssp_templates
from lsstdesc_diffsky import read_diffskypop_params
from lsstdesc_diffsky.defaults import OUTER_RIM_COSMO_PARAMS
from lsstdesc_diffsky.io_utils import load_healpixel
from lsstdesc_diffsky.io_utils import load_diffsky_params
from lsstdesc_diffsky.sed import calc_rest_sed_disk_bulge_knot_galpop


if __name__ == "__main__":

    rng = np.random.default_rng(13720)

    pixel_scale = 0.2

    filters = {"u", "g", "r", "i", "z", "y"}
    bps = {
        f: galsim.Bandpass(f"LSST_{f}.dat", "nm").withZeropoint("AB")
        for f in filters
    }

    fn = "/opt/diffsky.testdata.hdf5"
    patlist = ('LSST', 'disk', 'spheroid', 'knot')
    mock, metadata = load_healpixel(fn, patlist)
    diffsky_param_data = load_diffsky_params(mock)

    all_diffskypop_params = read_diffskypop_params("roman_rubin_2023")
    print(all_diffskypop_params._fields)

    ssp_data = load_ssp_templates(fn='/opt/dsps_ssp_data.h5')
    print(ssp_data._fields)

    print('\nssp_lgmet.shape = {}'.format(ssp_data.ssp_lgmet.shape))
    print('ssp_lg_age_gyr.shape = {}'.format(ssp_data.ssp_lg_age_gyr.shape))
    print('ssp_wave.shape = {}'.format(ssp_data.ssp_wave.shape))
    print('ssp_flux.shape = {}'.format(ssp_data.ssp_flux.shape))

    args = (
        mock['redshift'],
        diffsky_param_data.mah_params,
        diffsky_param_data.ms_params,
        diffsky_param_data.q_params,
        diffsky_param_data.fbulge_params,
        diffsky_param_data.fknot,
        *ssp_data,
        *all_diffskypop_params,
        OUTER_RIM_COSMO_PARAMS
    )

    disk_bulge_sed_info = calc_rest_sed_disk_bulge_knot_galpop(*args)

    galsim_mags = []
    dsps_mags = []

    for igal in range(len(mock["redshift"])):

        redshift = mock["redshift"][igal]
        luminosity_distance = dsps.cosmology.luminosity_distance_to_z(
            redshift,
            OUTER_RIM_COSMO_PARAMS.Om0,
            OUTER_RIM_COSMO_PARAMS.w0,
            OUTER_RIM_COSMO_PARAMS.wa,
            OUTER_RIM_COSMO_PARAMS.h,
        )
        luminosity_area = 4 * np.pi * luminosity_distance**2

        wave_type = "angstrom"
        flux_type = "fnu"
        # DSPS units
        # import astropy.units as u
        # _wave_type = u.angstrom
        # _flux_type = u.Lsun / u.Hz / u.Mpc**2
        # flux_factor = (1 * _flux_type).to(galsim.SED._fnu).value
        flux_factor = 4.0204145742268754e-16

        bulge_sed_table = galsim.LookupTable(
            ssp_data.ssp_wave,
            disk_bulge_sed_info.rest_sed_bulge[igal] / luminosity_area * flux_factor,
        )
        bulge_sed = galsim.SED(
            bulge_sed_table,
            wave_type=wave_type,
            flux_type=flux_type,
            redshift=redshift,
        )

        disk_sed_table = galsim.LookupTable(
            ssp_data.ssp_wave,
            disk_bulge_sed_info.rest_sed_diffuse_disk[igal] / luminosity_area * flux_factor,
        )
        disk_sed = galsim.SED(
            disk_sed_table,
            wave_type=wave_type,
            flux_type=flux_type,
            redshift=redshift,
        )

        knot_sed_table = galsim.LookupTable(
            ssp_data.ssp_wave,
            disk_bulge_sed_info.rest_sed_knot[igal] / luminosity_area * flux_factor,
        )
        knot_sed = galsim.SED(
            knot_sed_table,
            wave_type=wave_type,
            flux_type=flux_type,
            redshift=redshift,
        )

        _galsim_mags = []
        _dsps_mags = []

        for _f, bp in bps.items():
            galsim_disk_mag = disk_sed.calculateMagnitude(bp)
            galsim_bulge_mag = bulge_sed.calculateMagnitude(bp)
            galsim_knot_mag = knot_sed.calculateMagnitude(bp)

            dsps_disk_mag = dsps.calc_obs_mag(
                ssp_data.ssp_wave,
                disk_bulge_sed_info.rest_sed_diffuse_disk[igal],
                ssp_data.ssp_wave,
                bp(ssp_data.ssp_wave / 10.),
                redshift,
                OUTER_RIM_COSMO_PARAMS.Om0,
                OUTER_RIM_COSMO_PARAMS.w0,
                OUTER_RIM_COSMO_PARAMS.wa,
                OUTER_RIM_COSMO_PARAMS.h,
            )
            dsps_bulge_mag = dsps.calc_obs_mag(
                ssp_data.ssp_wave,
                disk_bulge_sed_info.rest_sed_bulge[igal],
                ssp_data.ssp_wave,
                bp(ssp_data.ssp_wave / 10.),
                redshift,
                OUTER_RIM_COSMO_PARAMS.Om0,
                OUTER_RIM_COSMO_PARAMS.w0,
                OUTER_RIM_COSMO_PARAMS.wa,
                OUTER_RIM_COSMO_PARAMS.h,
            )
            dsps_knot_mag = dsps.calc_obs_mag(
                ssp_data.ssp_wave,
                disk_bulge_sed_info.rest_sed_knot[igal],
                ssp_data.ssp_wave,
                bp(ssp_data.ssp_wave / 10.),
                redshift,
                OUTER_RIM_COSMO_PARAMS.Om0,
                OUTER_RIM_COSMO_PARAMS.w0,
                OUTER_RIM_COSMO_PARAMS.wa,
                OUTER_RIM_COSMO_PARAMS.h,
            )

            _galsim_mags.append([galsim_disk_mag, galsim_bulge_mag, galsim_knot_mag])
            _dsps_mags.append([dsps_disk_mag, dsps_bulge_mag, dsps_knot_mag])

        galsim_mags.append(np.array(_galsim_mags))
        dsps_mags.append(np.array(_dsps_mags))

    galsim_mags = np.array(galsim_mags)
    dsps_mags = np.array(dsps_mags)

    rtol=1e-2
    atol=0
    np.testing.assert_allclose(galsim_mags, dsps_mags, rtol=rtol, atol=atol)

As of rev. 17, test.py fails for ~0.4% of magnitudes at the specified tolerance:

Not equal to tolerance rtol=0.01, atol=0

Mismatched elements: 3 / 864 (0.347%)
Max absolute difference: 0.29969147
Max relative difference: 0.01303367
 x: array([[17.845955, 17.442211, 17.31659 , 18.507389, 17.540592, 17.622269,
        19.08969 , 18.69437 , 18.541994, 19.782619, 18.781718, 18.832596,
        23.556089, 23.666928, 23.581794, 24.016879, 23.580543, 23.606393,...
 y: array([[17.845878, 17.442189, 17.316638, 18.507395, 17.541002, 17.622273,
        19.089554, 18.694361, 18.542069, 19.782604, 18.781878, 18.832567,
        23.556017, 23.666881, 23.58201 , 24.016766, 23.584639, 23.606465,...