ePSF modelling of multiple images

ePSF modelling of multiple images#

As in the Getting Started, we will replicate results from Calissendorff et al. (2023), mainly because this is the dataset I have in hand as of writing this. This time, instead of fitting a single image, we will demonstrate how epsf can fit multiple images of the same source simultaneously.

Loading images#

from pathlib import Path
from astroquery.mast import Observations

img_dir = Path("./data")
img_dir.mkdir(exist_ok=True)
img_paths = []
for i in range(1, 6):
    filename = f"jw02473053001_03101_0000{i}_nrcblong_cal.fits"
    img_path = img_dir / filename
    status, msg, _ = Observations.download_file(
        f"mast:JWST/product/{filename}", local_path=img_path
    )
    img_paths.append(img_path)
nimg = len(img_paths)
INFO: Found cached file data/jw02473053001_03101_00001_nrcblong_cal.fits with expected size 117576000. [astroquery.query]
INFO: Found cached file data/jw02473053001_03101_00002_nrcblong_cal.fits with expected size 117576000. [astroquery.query]
INFO: Found cached file data/jw02473053001_03101_00003_nrcblong_cal.fits with expected size 117576000. [astroquery.query]
INFO: Found cached file data/jw02473053001_03101_00004_nrcblong_cal.fits with expected size 117576000. [astroquery.query]
INFO: Found cached file data/jw02473053001_03101_00005_nrcblong_cal.fits with expected size 117576000. [astroquery.query]
import numpy as np
import matplotlib.pyplot as plt

import epsf.utils as ut
from epsf.plot import plot_mosaic, plot_with_diff, plot_flat_samples

# The dithers could have different positions so let us make this generic
pos_arr = np.array([[1281, 825]] * nimg)
half_size = 34
final_size = 11

img_arr = np.empty((nimg, final_size, final_size))
err_arr = np.empty((nimg, final_size, final_size))
for i, img_path in enumerate(img_paths):
    img_arr[i], err_arr[i] = ut.open_jwst_image(
        img_path,
        pos_cut=pos_arr[i],
        half_size_cut=half_size,
        final_size=final_size,
        recenter=True,
    )
titles = [f"Dither {i}" for i in range(1, nimg + 1)]
plot_mosaic(img_arr, nrows=1, ncols=nimg, titles=titles)
plt.show()
../_images/fea3764994b7f25b519f5d79bd31ac47d296f74ea2d3eb11fa9bcc36408b9c37.png

Note

JWST1PASS ePSFs must first be downloaded. See the installation instructions for more detail.

from epsf.grid import EPSFGrid

# There is no PSF for module B so we use module A
grid = EPSFGrid.from_params("NIRCAM", "F480M", detector="NRCAL")
psf_list = [grid(*pos) for pos in pos_arr]

Single#

import simpple.distributions as sdist
from epsf.bayes import SingleModelMulti

x_params = {f"x0{i}": sdist.Uniform(-1.5, 1.5) for i in range(nimg)}
y_params = {f"y0{i}": sdist.Uniform(-1.5, 1.5) for i in range(nimg)}
parameters = (
    {
        "flux": sdist.LogUniform(1e-5, 5000),
        "bkg": sdist.Uniform(-100, 100),
        "sigma": sdist.LogUniform(1e-5, 100),
    }
    | x_params
    | y_params
)
model = SingleModelMulti(parameters, psf_list)
from ultranest import ReactiveNestedSampler

sampler = ReactiveNestedSampler(
    model.keys(),
    lambda p: model.log_likelihood(p, img_arr, err_arr),
    model.prior_transform,
)
sampler.run(show_status=False);
[ultranest] Sampling 400 live points from prior ...
[ultranest] Explored until L=-1e+03  
[ultranest] Likelihood function evaluations: 233044
[ultranest]   logZ = -1361 +- 0.2023
[ultranest] Effective samples strategy satisfied (ESS = 3677.8, need >400)
[ultranest] Posterior uncertainty strategy is satisfied (KL: 0.45+-0.07 nat, need <0.50 nat)
[ultranest] Evidency uncertainty strategy wants 398 minimum live points (dlogz from 0.15 to 0.51, need <0.5)
[ultranest]   logZ error budget: single: 0.39 bs:0.20 tail:0.01 total:0.20 required:<0.50
[ultranest] done iterating.
sampler.plot()
plt.show()
../_images/7814cd6fa9b2efb4789359cc25322b34c4bead9053e9f1dde5fbeec6f96ca053.png ../_images/109e92c50927c41eb0cdddda05495b3c0e4470dd660a8f192e7a9752ed19ba11.png ../_images/f22253a94be9a9cc51044515d234629c0b73b57936ae2379176dee7815db7875.png

We can then use the posterior median to compare the model with the data.

med_p = dict(
    zip(
        model.keys(),
        sampler.results["posterior"]["median"],
    )
)
med_mod = model.forward(med_p)

for i in range(nimg):
    fig, axs = plot_with_diff(img_arr[i], med_mod[i], scale="linear")
    axs[0].set_title("Data")
    axs[1].set_title("Median model")
    axs[2].set_title("Difference")
    fig.suptitle(f"Dither {i + 1}")
    plt.show()
../_images/629e20f6f004c6a1aee5694b336d0df77f6ab08be99fcbcb5e508853c130e7ed.png ../_images/9e88874f2be7f79062f8dad71fd1d512c60bc62c2b6db71cb9b685ba6ea9c96c.png ../_images/b83fdd2a083f1d8f7f591b84c47600404b787904146f207c5634d1b23fd838cf.png ../_images/e16ba62f19820e6ce13177e359884ff4036ad0b22a68bd4fbe8cbf906ee80a8d.png ../_images/65b1e55441e34f00fdbfe3c54a7c70c0f34d368e5438543ba80a31e526632d3e.png
img_samples = model.get_posterior_pred(sampler.results["samples"].T, 100)
for i in range(nimg):
    fig, axs = plot_flat_samples(img_arr[i], err_arr[i], img_samples[:, i])
    axs[0].set_title(f"Dither {i + 1}")
plt.show()
../_images/f9856ea930a26746603dfe5da0ec3c88c63830880ff9d5eb23386d79db46e394.png ../_images/3ccb171876ddca15e9d7f5a4cfef7424d2893618aa01ccbbb765bbd7bac31287.png ../_images/3ef3f2502528ae16c1b67aedf75f9ebf6cbbbd1cb2e478d4d3401086acf0457a.png ../_images/e488f396d53ce7cc82cc1d6ee21826b1cac99c10bbbd3b762614849cf0e3472c.png ../_images/c4fa9e549534a970604e6af7a776fc998b98c6deaa1d3e8a56dad12fb6812c04.png

Binary model#

We will now repeat the same steps as above, but with a binary model.

from epsf.bayes import BinaryModelMulti

x_params = {f"x0{i}": sdist.Uniform(-1.5, 1.5) for i in range(nimg)}
y_params = {f"y0{i}": sdist.Uniform(-1.5, 1.5) for i in range(nimg)}
parameters = (
    {
        "flux": sdist.LogUniform(1e-5, 5000),
        "sep": sdist.Uniform(0.5, 5),
        "pa": sdist.Uniform(0, 360),
        "cr": sdist.LogUniform(1e-4, 1.0),
        "bkg": sdist.Uniform(-100, 100),
        "sigma": sdist.LogUniform(1e-5, 100),
    }
    | x_params
    | y_params
)
model_binary = BinaryModelMulti(parameters, psf_list)

Since we have more parameters, UltraNest will struggle to generate valid samples in a high-dimensional space. Instead of region sampling, we add a stepsampler to our nested sampler. This enables a more efficient exploration of parameter space.

from ultranest.stepsampler import SliceSampler, generate_mixture_random_direction

nsteps = len(model.keys()) * 2
wrapped_params = [pname == "pa" for pname in model_binary.keys()]
sampler_binary = ReactiveNestedSampler(
    model_binary.keys(),
    lambda p: model_binary.log_likelihood(p, img_arr, err_arr),
    model_binary.prior_transform,
    wrapped_params=wrapped_params,
)
sampler_binary.stepsampler = SliceSampler(
    nsteps=nsteps,
    generate_direction=generate_mixture_random_direction,
)

sampler_binary.run(show_status=False);
[ultranest] Sampling 400 live points from prior ...
[ultranest] Explored until L=-4e+02  
[ultranest] Likelihood function evaluations: 4142492
[ultranest]   logZ = -537.8 +- 0.2651
[ultranest] Effective samples strategy satisfied (ESS = 4150.9, need >400)
[ultranest] Posterior uncertainty strategy is satisfied (KL: 0.46+-0.06 nat, need <0.50 nat)
[ultranest] Evidency uncertainty strategy wants 398 minimum live points (dlogz from 0.22 to 0.66, need <0.5)
[ultranest]   logZ error budget: single: 0.47 bs:0.27 tail:0.01 total:0.27 required:<0.50
[ultranest] done iterating.
sampler_binary.plot()
plt.show()
../_images/a5a071d65c4a58a6ecfcad4199344baa2596418d32a150ddecced459fab9eff8.png ../_images/355440548599b52f4d6b65d2ee4192a607fe4d5f9fe801ed777aa49949231ecd.png ../_images/211e966bd5482a96e35a6adb6d43f1a6e52f60cf33e0522dd92073822c8c6564.png
med_p = dict(
    zip(
        model_binary.keys(),
        sampler_binary.results["posterior"]["median"],
    )
)
med_mod = model_binary.forward(med_p)

for i in range(nimg):
    fig, axs = plot_with_diff(img_arr[i], med_mod[i], scale="linear")
    axs[0].set_title("Data")
    axs[1].set_title("Median model")
    axs[2].set_title("Difference")
    plt.show()
../_images/1931c952be9f7462940fdd0f708fa2601cf04aa26eb06b514dc4cd3a298b1f28.png ../_images/f5c35495d44a2783e3ad728eccda7821fa074f73f6b104298b835acf8c61ab45.png ../_images/df0fe5cfcf19113ef2088deb00b139b938c303cd4da27086c154c0722a80112a.png ../_images/79208e2c252a7ab7cac631622e4c778410f06e2d2a1f5916fc3f97f5eab6581b.png ../_images/29cb62ba0b8248f5cce1186c5de77f3d54b70b1741dbee8b7ee89475dd99516e.png

In addition to comparing the median binary model, we can subtract only the primary to see the companion pop up.

med_mod_primary = model.forward(med_p)
for i in range(nimg):
    fig, axs = plot_with_diff(img_arr[i], med_mod_primary[i], scale="linear")
    axs[0].set_title("Data")
    axs[1].set_title("Median model")
    axs[2].set_title("Difference")
    plt.show()
../_images/2a3b0359f26bb5135ccb9baec3957fe6e142ed1600d7f3d10aa3e1c5fd44a54e.png ../_images/dffebbb952109425c4dfb305f67a3eb7154757b6593b5516dd7d7bae4d63ab0b.png ../_images/edf089598c431b6bc2024dd241dc4e7c0ef852994d9d196bd29d08384e26448d.png ../_images/a8fdd97915245417ddb9a0626e4b8ef101f7f642e871917593dc290f0722075f.png ../_images/80c9dc26942c37248c807332fad5e72212e9b70a3bf3857568b0ee1879521208.png
img_samples = model_binary.get_posterior_pred(sampler_binary.results["samples"].T, 100)
for i in range(nimg):
    fig, axs = plot_flat_samples(img_arr[i], err_arr[i], img_samples[:, i])
    axs[0].set_title(f"Dither {i + 1}")
plt.show()
../_images/d517b254bec623ac1a23d30fe4017c5e06a5f68055e05803db601cc316b94ecd.png ../_images/493d2e70a5cc08e01371f41809aaab37f690dc34dbf6edf034c190a82e773fd5.png ../_images/139410cc9784cbc0b81c036cce9ae5dda6b1adc7614f3de36c966b836173b23b.png ../_images/bef749c301c2c9bd157ab17f2cdd3282b947e345237370d5c90015e39b3ccea1.png ../_images/6ecb0b59282412f504920caa371b1db5f1b775ec56291f40b55987de7655374e.png
img_samples = model_binary.get_posterior_pred(sampler_binary.results["samples"].T, 100)
median_sigma = sampler_binary.results["posterior"]["median"][
    model_binary.keys().index("sigma")
]
inflated_err = np.sqrt(err_arr**2 + median_sigma**2)
for i in range(nimg):
    fig, axs = plot_flat_samples(img_arr[i], inflated_err[i], img_samples[:, i])
    axs[0].set_title(f"Dither {i + 1}, inflated errors")
plt.show()
../_images/46adedb5a732a694d105e7483612c7692685381c411eca3354bcf0fb433a0db8.png ../_images/b2b339e63558c608ede35cf7d179d7d3caa7ca92db661e993eaf1c13420a2fe8.png ../_images/6e9cdf4958be98f8f78ed55f5fd5650a8d892be5902ae5265963f6feb47909db.png ../_images/bda3e0004a144e3033d28e30d2fc7954a7f6c71e8fdc38a9e394574e05859434.png ../_images/acf24b8c4a66ad2fa05c2ad0160de5410316ab55ff7210a24bc2f46912b67071.png

Model comparison#

Let us now compare the single and binary models using the Bayes factor.

print(f"lnZ single {sampler.results['logz']:.2f} +/- {sampler.results['logzerr']:.2f}")
print(
    f"lnZ binary {sampler_binary.results['logz']:.2f} +/- {sampler_binary.results['logzerr']:.2f}"
)
lnK = sampler_binary.results["logz"] - sampler.results["logz"]
lnK_err = np.sqrt(
    sampler.results["logzerr"] ** 2 + sampler_binary.results["logzerr"] ** 2
)
print(f"lnK binary - single {lnK:.2f} +/- {lnK_err:.2f}")
lnZ single -1360.90 +/- 0.41
lnZ binary -537.86 +/- 0.66
lnK binary - single 823.04 +/- 0.78

As expected, the binary model is overwhelmingly favored by the data.