Last active
November 1, 2023 00:04
-
-
Save sidneymau/59f42482430340a01f8ae3e5f6b2185d to your computer and use it in GitHub Desktop.
Chromatic Galaxy Image Simulations
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
As of rev. 17,
test.py
fails for ~0.4% of magnitudes at the specified tolerance: