End-to-End PSF Dataset Generation Pipeline¶

%matplotlib inline
import psfcraft
from psfcraft.utils import generate_coefficients, build_wfe_budget
import numpy as np
import matplotlib.pyplot as plt

# --- Telescope configuration ---
PRIMARY_RADIUS   = 0.5    # metres
SECONDARY_RADIUS = 0.1    # metres
PIXEL_SCALE      = 0.298  # arcsec/pixel
FOV_PIXELS       = 32     # image size (keep small for speed)
WAVELENGTH       = 1e-6   # 1 µm

# --- WFE budget: [0, tip, tilt, defocus, astig×2, coma×2, trefoil×2, spherical, …]
# Radial budget: amplitude in nm per radial order (0=piston, 1=tip/tilt, …)
radial_budget_nm = [0, 100, 50, 36, 18, 9, 5]
wfe_budget = build_wfe_budget(radial_budget_nm)

print(f"Budget has {len(wfe_budget)} Zernike coefficients.")
print("Per-mode limits (nm):", wfe_budget[:12], "…")

# Quick bar chart
fig, ax = plt.subplots(figsize=(10, 3))
ax.bar(range(1, len(wfe_budget) + 1), wfe_budget, color="steelblue", edgecolor="white")
ax.set_xlabel("Noll index")
ax.set_ylabel("Max amplitude (nm)")
ax.set_title("WFE budget per Zernike mode")
ax.set_xlim(0.5, len(wfe_budget) + 0.5)
ax.grid(True, axis="y", alpha=0.3)
plt.tight_layout()
plt.show()

N_PSFS = 12   # small for tutorial speed
RNG_SEED = 42

rng = np.random.default_rng(RNG_SEED)  # reproducible random state

# Pre-build the telescope object (reused across draws)
tel = psfcraft.NewtonianTelescope(
    version="1_3",
    primary_radius=PRIMARY_RADIUS,
    secondary_radius=SECONDARY_RADIUS,
)
tel.pixelscale = PIXEL_SCALE

psf_stack   = []   # oversampled PSF images
coef_stack  = []   # corresponding Zernike coefficients

for i in range(N_PSFS):
    # Draw random coefficients within the per-mode budget
    coefs = generate_coefficients(wfe_budget)

    tel.wfe_coefficients = coefs
    psf = tel.calc_psf(monochromatic=WAVELENGTH, fov_pixels=FOV_PIXELS)

    psf_stack.append(psf[0].data)   # oversampled extension
    coef_stack.append(coefs)
    print(f"  PSF {i+1:02d}/{N_PSFS}  RMS WFE = {np.sqrt(np.sum(coefs**2))*1e9:.1f} nm", end="\r")

psf_array  = np.array(psf_stack)    # shape: (N_PSFS, H, W)
coef_array = np.array(coef_stack)   # shape: (N_PSFS, n_modes)

print(f"\nDataset ready: {psf_array.shape}  (N × H × W)")

from matplotlib.colors import LogNorm

n_cols = 6
n_rows = int(np.ceil(N_PSFS / n_cols))

fig, axes = plt.subplots(n_rows, n_cols, figsize=(n_cols * 2.5, n_rows * 2.5))
axes = axes.flatten()

for i, (psf_img, ax) in enumerate(zip(psf_array, axes)):
    rms_nm = np.sqrt(np.sum(coef_array[i] ** 2)) * 1e9
    im = ax.imshow(np.maximum(psf_img, 1e-10), cmap="inferno",
                   norm=LogNorm(), origin="lower")
    ax.set_title(f"#{i+1}  {rms_nm:.0f} nm", fontsize=9)
    ax.axis("off")

# Hide unused axes
for ax in axes[N_PSFS:]:
    ax.axis("off")

plt.suptitle("Random PSF dataset  (log scale, labelled by total RMS WFE)", fontsize=12)
plt.tight_layout()
plt.show()

from astropy.io import fits as pyfits

# We need HDUList objects for psfcraft.measure_ee — rebuild them
rms_wfe_nm = []
ee50_arcsec = []

for i, coefs in enumerate(coef_array):
    tel.wfe_coefficients = coefs
    psf_hdul = tel.calc_psf(monochromatic=WAVELENGTH, fov_pixels=FOV_PIXELS)

    rms_wfe_nm.append(np.sqrt(np.sum(coefs**2)) * 1e9)

    # Find radius at EE = 0.5
    ee_fn = psfcraft.measure_ee(psf_hdul)
    # Binary search for EE50 radius
    radii = np.linspace(0.01, 1.0, 200)
    ee_vals = np.array([ee_fn(r) for r in radii])
    idx = np.searchsorted(ee_vals, 0.5)
    ee50_arcsec.append(radii[min(idx, len(radii) - 1)])

rms_wfe_nm  = np.array(rms_wfe_nm)
ee50_arcsec = np.array(ee50_arcsec)

fig, axes = plt.subplots(1, 2, figsize=(11, 4))

axes[0].hist(rms_wfe_nm, bins=8, color="steelblue", edgecolor="white")
axes[0].set_xlabel("Total RMS WFE (nm)")
axes[0].set_ylabel("Count")
axes[0].set_title("Distribution of RMS WFE")

axes[1].scatter(rms_wfe_nm, ee50_arcsec, c="steelblue", edgecolors="white", s=60)
axes[1].set_xlabel("Total RMS WFE (nm)")
axes[1].set_ylabel("EE50 radius (arcsec)")
axes[1].set_title("EE50 vs. WFE — larger WFE → broader PSF")
axes[1].grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

output_dir = "/tmp/psfcraft_dataset"
import os
os.makedirs(output_dir, exist_ok=True)

# --- NumPy format (lightweight, fast) ---
np.save(os.path.join(output_dir, "psf_images.npy"),   psf_array)
np.save(os.path.join(output_dir, "wfe_coefficients.npy"), coef_array)

print(f"Saved {N_PSFS} PSFs to {output_dir}/")
print(f"  psf_images.npy      : {psf_array.shape}  dtype={psf_array.dtype}")
print(f"  wfe_coefficients.npy: {coef_array.shape}  dtype={coef_array.dtype}")

# --- Reload and verify ---
psf_reload  = np.load(os.path.join(output_dir, "psf_images.npy"))
coef_reload = np.load(os.path.join(output_dir, "wfe_coefficients.npy"))
assert np.allclose(psf_reload, psf_array), "Round-trip check failed!"
print("\nRound-trip verification: ✅ OK")