diff --git a/.gitignore b/.gitignore index 67052527..899e7843 100644 --- a/.gitignore +++ b/.gitignore @@ -75,3 +75,4 @@ gh-pages prof.out downloaded_eigen +AGENTS.md diff --git a/devel/roman/diagnose_bilinear_vs_direct.py b/devel/roman/diagnose_bilinear_vs_direct.py new file mode 100644 index 00000000..68eaa4e1 --- /dev/null +++ b/devel/roman/diagnose_bilinear_vs_direct.py @@ -0,0 +1,335 @@ +#!/usr/bin/env python3 + +"""Diagnose Roman bilinear-corner approximation error against direct getPSF. + +Example: + python devel/roman/diagnose_bilinear_vs_direct.py --grid 5 --plot devel/roman/bilinear_map.png +""" + +import argparse +import os + +import galsim +import numpy as np +import piff +import piff.roman_psf + + +def parse_scas(spec): + if spec is None or spec.strip() == "": + return list(range(1, 19)) + out = [] + for token in spec.split(","): + tok = token.strip() + if "-" in tok: + a, b = tok.split("-", 1) + out.extend(range(int(a), int(b) + 1)) + else: + out.append(int(tok)) + out = sorted(set(out)) + for sca in out: + if sca < 1 or sca > 18: + raise ValueError(f"Invalid SCA {sca}. Must be in 1..18.") + return out + + +def parse_aberrations(spec, nparam): + if spec is None: + return np.zeros(nparam, dtype=float) + vals = [float(v.strip()) for v in spec.split(",") if v.strip() != ""] + if len(vals) != nparam: + raise ValueError(f"--aberrations requires {nparam} values (got {len(vals)})") + return np.array(vals, dtype=float) + + +def point_metrics(model_image, direct_image): + diff = model_image - direct_image + max_abs = float(np.max(np.abs(diff))) + rms = float(np.sqrt(np.mean(diff**2))) + peak = float(np.max(np.abs(direct_image))) + l2 = float(np.sqrt(np.sum(direct_image**2))) + frac_peak = np.nan if peak == 0.0 else max_abs / peak + frac_l2 = np.nan if l2 == 0.0 else float(np.sqrt(np.sum(diff**2)) / l2) + return max_abs, rms, frac_peak, frac_l2 + + +def make_grid(npix, grid, margin): + x = np.linspace(margin, npix - margin, grid) + y = np.linspace(margin, npix - margin, grid) + return x, y + + +def run_scan(model, scas, grid, stamp_size, scale, margin, aberrations): + nparam = model.param_len + if aberrations.shape[0] != nparam: + raise ValueError(f"Expected {nparam} aberration values, got {aberrations.shape[0]}") + + xg, yg = make_grid(model.sca_size, grid, margin) + rows = [] + maps = {} + + for sca in scas: + max_abs_map = np.zeros((grid, grid), dtype=float) + rms_map = np.zeros((grid, grid), dtype=float) + frac_peak_map = np.zeros((grid, grid), dtype=float) + frac_l2_map = np.zeros((grid, grid), dtype=float) + + for iy, y in enumerate(yg): + for ix, x in enumerate(xg): + star = piff.Star.makeTarget( + x=float(x), + y=float(y), + stamp_size=stamp_size, + scale=scale, + properties={"sca": int(sca)}, + ).withFlux(1.0, (0.0, 0.0)) + star = model.initialize(star) + + model_star = model.draw( + piff.Star( + star.data, + star.fit.newParams( + aberrations, params_var=np.zeros_like(aberrations), num=model._num + ), + ) + ) + + prof = galsim.roman.getPSF( + int(sca), + model.filter, + SCA_pos=star.image_pos, + pupil_bin=piff.roman_psf.pupil_bin, + wcs=star.image.wcs, + extra_aberrations=model._make_extra_aberrations(aberrations), + wavelength=None if model.chromatic else model.bandpass.effective_wavelength, + ) + direct_image = star.image.copy() + prof.drawImage( + direct_image, + method=model._method, + center=star.image_pos, + bandpass=model.bandpass if model.chromatic else None, + ) + + max_abs, rms, frac_peak, frac_l2 = point_metrics( + model_star.image.array, direct_image.array + ) + max_abs_map[iy, ix] = max_abs + rms_map[iy, ix] = rms + frac_peak_map[iy, ix] = frac_peak + frac_l2_map[iy, ix] = frac_l2 + + maps[sca] = { + "x_grid": xg.copy(), + "y_grid": yg.copy(), + "max_abs": max_abs_map, + "rms": rms_map, + "frac_peak": frac_peak_map, + "frac_l2": frac_l2_map, + } + + rows.append( + ( + sca, + float(np.nanmedian(max_abs_map)), + float(np.nanpercentile(max_abs_map, 95)), + float(np.nanmax(max_abs_map)), + float(np.nanmedian(frac_peak_map)), + float(np.nanpercentile(frac_peak_map, 95)), + float(np.nanmax(frac_peak_map)), + float(np.nanmedian(frac_l2_map)), + float(np.nanpercentile(frac_l2_map, 95)), + float(np.nanmax(frac_l2_map)), + ) + ) + return rows, maps + + +def print_summary(rows): + print("\nPer-SCA bilinear-vs-direct mismatch summary:") + print( + " SCA med|max| p95|max| max|max| med(max/peak) p95(max/peak)" + " max(max/peak) med(l2frac) p95(l2frac) max(l2frac)" + ) + for row in sorted(rows, key=lambda r: r[0]): + print( + f" {row[0]:3d} {row[1]:10.4e} {row[2]:10.4e} {row[3]:10.4e}" + f" {row[4]:13.4e} {row[5]:13.4e} {row[6]:13.4e}" + f" {row[7]:11.4e} {row[8]:11.4e} {row[9]:11.4e}" + ) + + worst = max(rows, key=lambda r: r[6]) + print( + "\nWorst SCA by peak-fraction max mismatch: " + f"SCA {worst[0]} with max(max/peak)={worst[6]:.4e}" + ) + + +def save_maps_npz(path, maps): + outdir = os.path.dirname(path) + if outdir: + os.makedirs(outdir, exist_ok=True) + payload = {} + for sca, d in maps.items(): + payload[f"sca{sca}_x"] = d["x_grid"] + payload[f"sca{sca}_y"] = d["y_grid"] + payload[f"sca{sca}_max_abs"] = d["max_abs"] + payload[f"sca{sca}_rms"] = d["rms"] + payload[f"sca{sca}_frac_peak"] = d["frac_peak"] + payload[f"sca{sca}_frac_l2"] = d["frac_l2"] + np.savez(path, **payload) + print(f"\nWrote mismatch maps to {path}") + + +def maybe_plot(path, maps, metric): + if path is None: + return + import matplotlib.pyplot as plt + outdir = os.path.dirname(path) + if outdir: + os.makedirs(outdir, exist_ok=True) + + if metric not in ("max_abs", "rms", "frac_peak", "frac_l2"): + raise ValueError(f"Invalid metric {metric}") + + scas = sorted(maps) + ncol = 6 + nrow = int(np.ceil(len(scas) / ncol)) + fig, axes = plt.subplots( + nrow, ncol, figsize=(3.1 * ncol, 2.9 * nrow), squeeze=False, constrained_layout=True + ) + + vmin = min(np.nanmin(maps[sca][metric]) for sca in scas) + vmax = max(np.nanmax(maps[sca][metric]) for sca in scas) + + for ax in axes.ravel(): + ax.set_visible(False) + + im = None + for i, sca in enumerate(scas): + ax = axes.ravel()[i] + ax.set_visible(True) + data = maps[sca][metric] + x = maps[sca]["x_grid"] + y = maps[sca]["y_grid"] + im = ax.imshow( + data, + origin="lower", + extent=[x.min(), x.max(), y.min(), y.max()], + interpolation="nearest", + vmin=vmin, + vmax=vmax, + aspect="auto", + ) + ax.set_title(f"SCA {sca}") + ax.set_xlabel("x") + ax.set_ylabel("y") + + if im is not None: + cbar = fig.colorbar(im, ax=axes.ravel().tolist(), shrink=0.92, pad=0.01) + cbar.set_label(metric) + fig.suptitle(f"Roman bilinear-vs-direct mismatch map: {metric}") + fig.savefig(path, dpi=150) + print(f"Wrote plot to {path}") + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--scas", + default=None, + help="SCA selection, e.g. '1-18' or '5,8,10'. Default: all 1..18.", + ) + parser.add_argument( + "--grid", + type=int, + default=5, + help="Grid points per axis per SCA (default: 5).", + ) + parser.add_argument( + "--margin", + type=float, + default=64.0, + help="Grid margin in pixels from each SCA edge (default: 64).", + ) + parser.add_argument( + "--stamp-size", + type=int, + default=25, + help="Stamp size in pixels (default: 25).", + ) + parser.add_argument( + "--scale", + type=float, + default=0.11, + help="Pixel scale in arcsec/pixel (default: 0.11).", + ) + parser.add_argument( + "--pupil-bin", + type=int, + default=8, + help="Roman pupil_bin for GalSim getPSF draws (default: 8).", + ) + parser.add_argument( + "--aberrations", + default=None, + help="Comma-separated z4..zN values (length max_zernike-3). Default all zeros.", + ) + parser.add_argument( + "--max-zernike", + type=int, + default=6, + help="Roman max_zernike for diagnostic model (default: 6).", + ) + parser.add_argument( + "--plot", + default=None, + help="Optional output image path for 18-SCA metric map.", + ) + parser.add_argument( + "--plot-metric", + default="frac_peak", + choices=["max_abs", "rms", "frac_peak", "frac_l2"], + help="Metric to visualize when --plot is set (default: frac_peak).", + ) + parser.add_argument( + "--save-npz", + default=None, + help="Optional .npz path to save all per-SCA mismatch maps.", + ) + args = parser.parse_args() + + piff.roman_psf.pupil_bin = int(args.pupil_bin) + model = piff.Roman( + filter="H158", + chromatic=False, + max_zernike=args.max_zernike, + aberration_interp="constant", + aberration_prior_sigma=1.0e6, + ) + + scas = parse_scas(args.scas) + aberrations = parse_aberrations(args.aberrations, model.param_len) + print( + "Running bilinear-vs-direct scan with " + f"scas={scas}, grid={args.grid}, pupil_bin={args.pupil_bin}, " + f"max_zernike={args.max_zernike}, aberrations={aberrations}" + ) + + rows, maps = run_scan( + model=model, + scas=scas, + grid=args.grid, + stamp_size=args.stamp_size, + scale=args.scale, + margin=args.margin, + aberrations=aberrations, + ) + print_summary(rows) + if args.save_npz: + save_maps_npz(args.save_npz, maps) + maybe_plot(args.plot, maps, args.plot_metric) + + +if __name__ == "__main__": + main() diff --git a/devel/roman/diagnose_piff.py b/devel/roman/diagnose_piff.py new file mode 100644 index 00000000..60c7d2b6 --- /dev/null +++ b/devel/roman/diagnose_piff.py @@ -0,0 +1,270 @@ +#!/usr/bin/env python3 + +"""Quick diagnostics for a fitted Piff file. + +Usage: + python devel/roman/diagnose_piff.py devel/roman/ffov_13906_1.piff +""" + +import argparse +from collections import defaultdict + +import numpy as np +import piff + + +def _get_sca(star): + props = star.data.properties + if "sca" in props: + return int(props["sca"]) + if "chipnum" in props: + return int(props["chipnum"]) + return -1 + + +def _fmt(v): + if np.isscalar(v): + return f"{v:.4g}" + return str(v) + + +def _collect_roman_sca_means(psf): + if not (hasattr(psf, "components") and len(psf.components) > 0): + return None, None + comp0 = psf.components[0] + if getattr(comp0, "_type_name", None) != "RomanOptics": + return None, None + + model_num = comp0.model._num + params_by_sca = defaultdict(list) + for star in psf.stars: + if star.is_flagged or star.is_reserve: + continue + params_by_sca[_get_sca(star)].append(star.fit.get_params(model_num)) + + if not params_by_sca: + return comp0, None + + sca_mean = {sca: np.mean(vals, axis=0) for sca, vals in params_by_sca.items()} + return comp0, sca_mean + + +def _print_roman_aberration_diagnostics(sca_mean, focus_sca): + if not sca_mean: + return + + scas = np.array(sorted(sca_mean), dtype=int) + params = np.array([sca_mean[s] for s in scas], dtype=float) + nterms = params.shape[1] + + med = np.median(params, axis=0) + mad = np.median(np.abs(params - med), axis=0) + robust_sigma = 1.4826 * mad + robust_sigma_safe = np.where(robust_sigma > 0, robust_sigma, np.nan) + + print("\nRomanOptics per-SCA aberration summary:") + print(" (term 0 -> z4, term 1 -> z5, etc.)") + for sca, p in zip(scas, params): + max_i = int(np.argmax(np.abs(p))) + l2 = np.linalg.norm(p) + l1 = np.sum(np.abs(p)) + print( + " SCA {:2d}: ||p||2={:8.4f} ||p||1={:8.4f} max|term|=t{:02d}:{:8.4f}".format( + int(sca), float(l2), float(l1), max_i, float(p[max_i]) + ) + ) + + mean_abs = np.mean(np.abs(params), axis=0) + print("\nRomanOptics per-term amplitudes across SCAs:") + for j in range(nterms): + print( + " term {:2d} (z{:2d}): mean|a|={:8.4f} med={:8.4f} robust_sigma={:8.4f}".format( + j, j + 4, float(mean_abs[j]), float(med[j]), + float(robust_sigma[j]) if np.isfinite(robust_sigma[j]) else np.nan + ) + ) + + if focus_sca is None: + return + elif focus_sca not in sca_mean: + print(f"\nFocused SCA {focus_sca} not present in fitted usable stars.") + return + + p = sca_mean[focus_sca] + zscore = (p - med) / robust_sigma_safe + print(f"\nFocused SCA {focus_sca} term-by-term offsets from focal-plane median:") + for j, (val, m, dz, rs) in enumerate(zip(p, med, zscore, robust_sigma)): + ztxt = f"{dz:+7.2f}" if np.isfinite(dz) else " n/a " + rstxt = f"{rs:8.4f}" if np.isfinite(rs) else " n/a " + print( + " term {:2d} (z{:2d}): val={:8.4f} med={:8.4f} delta={:+8.4f}" + " robust_sigma={} robust_z={}".format( + j, j + 4, float(val), float(m), float(val - m), rstxt, ztxt + ) + ) + + sca_norm = {int(s): float(np.linalg.norm(v)) for s, v in sca_mean.items()} + worst = sorted(sca_norm.items(), key=lambda kv: -kv[1])[:5] + print("\nTop 5 SCAs by RomanOptics ||p||2:") + for sca, norm in worst: + tag = " <-- focused" if sca == focus_sca else "" + print(f" SCA {sca:2d}: {norm:8.4f}{tag}") + + +def summarize(psf, center_thresh, top, focus_sca): + stars = psf.stars + print(f"File summary: nstars={len(stars)}") + if getattr(psf, "dof", 0): + print( + " fit: chisq={:.3f} dof={} chisq/dof={:.6f} niter={} nremoved={}".format( + float(psf.chisq), + int(psf.dof), + float(psf.chisq) / float(psf.dof), + int(getattr(psf, "niter", 0)), + int(getattr(psf, "nremoved", 0)), + ) + ) + else: + print(" fit metadata dof is 0 or unavailable") + + rows = [] + for i, star in enumerate(stars): + chisq = np.nan if star.fit.chisq is None else float(star.fit.chisq) + dof = -1 if star.fit.dof is None else int(star.fit.dof) + flux = np.nan if star.fit.flux is None else float(star.fit.flux) + cx, cy = star.fit.center + rows.append( + ( + i, + _get_sca(star), + chisq, + dof, + flux, + float(cx), + float(cy), + bool(star.is_flagged), + bool(star.is_reserve), + float(star.x), + float(star.y), + ) + ) + arr = np.array( + rows, + dtype=[ + ("idx", "i4"), + ("sca", "i4"), + ("chisq", "f8"), + ("dof", "i4"), + ("flux", "f8"), + ("cx", "f8"), + ("cy", "f8"), + ("flagged", "?"), + ("reserve", "?"), + ("x", "f8"), + ("y", "f8"), + ], + ) + good = (~arr["flagged"]) & (~arr["reserve"]) & (arr["dof"] > 0) + print(f" usable stars (not flagged/reserve): {int(np.sum(good))}") + + if np.any(good): + chi = arr["chisq"][good] / arr["dof"][good] + print( + " usable chi2/dof: median={:.4f} p90={:.4f} max={:.4f}".format( + float(np.median(chi)), + float(np.percentile(chi, 90)), + float(np.max(chi)), + ) + ) + + center = np.hypot(arr["cx"], arr["cy"]) + large_center = center > center_thresh + print( + f" stars with |center| > {center_thresh:g}: {int(np.sum(large_center))}" + f" ({int(np.sum(large_center & good))} usable)" + ) + print(f" stars with negative flux: {int(np.sum(arr['flux'] < 0))}") + + print("\nPer-SCA (usable stars):") + for sca in sorted(np.unique(arr["sca"])): + m = good & (arr["sca"] == sca) + if not np.any(m): + continue + c = arr["chisq"][m] / arr["dof"][m] + cen = np.hypot(arr["cx"][m], arr["cy"][m]) + print( + f" SCA {sca:2d}: n={int(np.sum(m)):3d}" + f" med(chi2/dof)={np.median(c):7.3f}" + f" p90={np.percentile(c, 90):7.3f}" + f" max={np.max(c):7.3f}" + f" med|cen|={np.median(cen):6.3f}" + f" max|cen|={np.max(cen):6.3f}" + ) + + print(f"\nWorst {top} stars by chi2/dof:") + denom = np.maximum(arr["dof"], 1) + order = np.argsort(-(arr["chisq"] / denom)) + for j in order[:top]: + cdr = arr["chisq"][j] / max(arr["dof"][j], 1) + cen = np.hypot(arr["cx"][j], arr["cy"][j]) + print( + " idx={} sca={} chi2/dof={:.3f} |cen|={:.3f} flux={} flagged={} reserve={}" + " x={:.1f} y={:.1f}".format( + int(arr["idx"][j]), + int(arr["sca"][j]), + float(cdr), + float(cen), + _fmt(arr["flux"][j]), + bool(arr["flagged"][j]), + bool(arr["reserve"][j]), + float(arr["x"][j]), + float(arr["y"][j]), + ) + ) + + comp0, sca_mean = _collect_roman_sca_means(psf) + if comp0 is not None and sca_mean: + print("\nComponent 0 (RomanOptics) mean params by SCA:") + for sca in sorted(sca_mean): + pstr = np.array2string( + sca_mean[sca], + formatter={"float_kind": lambda x: f"{x:.4f}"}, + ) + print(f" SCA {sca:2d}: {pstr}") + _print_roman_aberration_diagnostics(sca_mean, focus_sca=focus_sca) + + +def main(): + parser = argparse.ArgumentParser(description="Summarize diagnostics from a .piff file") + parser.add_argument("piff_file", help="Path to .piff file") + parser.add_argument( + "--center-thresh", + type=float, + default=0.5, + help="Threshold in pixels for reporting large fitted centers", + ) + parser.add_argument( + "--top", + type=int, + default=10, + help="How many worst stars (by chi2/dof) to print", + ) + parser.add_argument( + "--focus-sca", + type=int, + default=None, + help="SCA to print detailed aberration diagnostics for", + ) + args = parser.parse_args() + + psf = piff.read(args.piff_file) + summarize( + psf, + center_thresh=args.center_thresh, + top=args.top, + focus_sca=args.focus_sca, + ) + + +if __name__ == "__main__": + main() diff --git a/devel/roman/piff.yaml b/devel/roman/piff.yaml index 25f7625e..fb545ed4 100644 --- a/devel/roman/piff.yaml +++ b/devel/roman/piff.yaml @@ -23,19 +23,19 @@ select: min_snr: 50 # reject very faint stars max_snr_weight: 500 # don't over-weight very bright stars - #max_mask_pixels: 10 # how many pixels in stamp are allowed to be masked. hsm_size_reject: 5 # number of inter-quartile ranges from median to reject size estimate. max_pixel_cut: 50000 # reject stars with a pixel value higher than this. (sort of) min_sep: 1.0 # reject stars with a neighbor closer than 1 arcsec away. #nstars: 100 output: - file_name: 'ffov_13906_1.piff' + dir: 'output' + file_name: 'ffov_13906_11.piff' stats: - type: StarImages - file_name: stars.png + file_name: 'stars_11.png' nplot: 0 # all stars psf: @@ -49,8 +49,11 @@ psf: filter: H158 chromatic: False max_zernike: 12 # Fit Zernike coefficients 4-12 inclusive. - aberration_prior_sigma: 0.05 # Gaussian prior on all Zernike coefficients in absolute units. - nproc: 4 + aberration_prior_sigma: 3.e-3 # Gaussian prior on all Zernike coefficients in absolute units + aberration_interp: constant # linear is possible but doesn't seem to work better. + nominal_interp: bilinear # five_points is possible but doesn't work better on + # this sim. May be different on real data. + nproc: 6 - type: Simple @@ -65,8 +68,8 @@ psf: outliers: - type: Chisq - nsigma: 5 - max_remove: 0.01 + nsigma: 10 + max_remove: 0.03 - type: Centroid max_offset: 0.2 # arcsec diff --git a/piff/roman/roman_psf.py b/piff/roman/roman_psf.py index 62cbcb85..834dc18f 100644 --- a/piff/roman/roman_psf.py +++ b/piff/roman/roman_psf.py @@ -130,10 +130,17 @@ def _finish_read(self, reader): class RomanOpticalModel(Model): """Model a Roman PSF using GalSim's built-in Roman optical model. - The expensive GalSim Roman PSF construction is approximated as a bilinear function across - each SCA. For each ``(sca, params)`` state we build corner PSFs at the four SCA corners and - interpolate between them at each star position. This keeps the optical model much faster when - many stars on a chip share similar parameter vectors. + The expensive GalSim Roman PSF construction is approximated across each SCA by evaluating + PSFs at a small set of fixed sample points and interpolating between them at each star + position. ``nominal_interp='bilinear'`` uses the four SCA corners; ``'five_point'`` adds + the SCA center and uses a five-term basis ``[1, x, y, x*y, (x^2+y^2)]`` in normalized + coordinates. + + The ``aberration_interp`` option controls how the fitted extra aberrations vary across an + SCA. ``'global'`` and ``'constant'`` both use one fitted aberration vector per SCA in the + model, with ``'global'`` only changing how the outer interpolator shares those vectors across + SCAs. ``'linear'`` fits a planar extra-aberration field on each SCA, modeled as + ``a + b x + c y`` in normalized SCA coordinates. """ _type_name = 'Roman' @@ -147,6 +154,7 @@ def __init__( chromatic=True, max_zernike=22, aberration_interp='constant', + nominal_interp='bilinear', aberration_prior_sigma=0.05, nproc=1, logger=None, @@ -156,13 +164,16 @@ def __init__( self.chromatic = chromatic self.max_zernike = int(max_zernike) self.aberration_interp = str(aberration_interp) + self.nominal_interp = str(nominal_interp) self.nproc = int(nproc) self.set_num(None) if self.max_zernike < 4 or self.max_zernike > 22: raise ValueError("max_zernike must be in the range 4..22") - if self.aberration_interp not in ('global', 'constant'): - raise ValueError("aberration_interp must be one of 'global', 'constant'") + if self.aberration_interp not in ('global', 'constant', 'linear'): + raise ValueError("aberration_interp must be one of 'global', 'constant', 'linear'") + if self.nominal_interp not in ('bilinear', 'five_point'): + raise ValueError("nominal_interp must be one of 'bilinear', 'five_point'") # Notation: filter is the string name of the filter. # bandpass is the galsim.Bandpass object with the transmission function. @@ -170,7 +181,9 @@ def __init__( if self.filter not in bandpasses: raise ValueError("Roman filter %r is not a valid GalSim Roman bandpass" % self.filter) self.bandpass = bandpasses[self.filter] - self.prior_sigma = self._parse_prior_sigma(aberration_prior_sigma) + # Save the original for serialization. + self.orig_prior_sigma = self._parse_prior_sigma(aberration_prior_sigma) + self.prior_sigma = self._expand_prior_sigma(self.orig_prior_sigma) self.prior_invsigsq = 1.0/self.prior_sigma**2 if self.prior_sigma is not None else None self.kwargs = { @@ -178,10 +191,12 @@ def __init__( 'chromatic': self.chromatic, 'max_zernike': self.max_zernike, 'aberration_interp': self.aberration_interp, - 'aberration_prior_sigma': self.prior_sigma, + 'nominal_interp': self.nominal_interp, + 'aberration_prior_sigma': self.orig_prior_sigma, 'nproc': self.nproc, } self.sca_size = float(galsim.roman.n_pix) + self._roman_five_point_data = {} self.clear_cache() def __getstate__(self): @@ -192,6 +207,13 @@ def __getstate__(self): @property def param_len(self): + if self.aberration_interp == 'linear': + return 3 * self._single_point_param_len + else: + return self._single_point_param_len + + @property + def _single_point_param_len(self): return self.max_zernike - 3 def initialize_iteration(self): @@ -201,6 +223,61 @@ def initialize_iteration(self): def clear_cache(self): self._corner_cache = {} + def _get_roman_five_point_data(self, sca): + if sca in self._roman_five_point_data: + return self._roman_five_point_data[sca] + + aberrations, x_pos, y_pos = galsim.roman.roman_psfs._read_aberrations(sca) + center = galsim.PositionD(x_pos[0], y_pos[0]) + + # The project aberrations data has aberration values at 5 points on each SCA: + # the four corners plus the center. The center is always at position 0, but + # the others aren't necessarily in the same order each time. So figure out + # which one is which. + # cf. _interp_aberrations_bilinear in galsim.roman.roman_psfs.py + ll = ul = lr = ur = None + for i in range(1, 5): + if x_pos[i] < x_pos[0] and y_pos[i] < y_pos[0]: + ll = i + if x_pos[i] < x_pos[0] and y_pos[i] > y_pos[0]: + ul = i + if x_pos[i] > x_pos[0] and y_pos[i] < y_pos[0]: + lr = i + if x_pos[i] > x_pos[0] and y_pos[i] > y_pos[0]: + ur = i + assert None not in (ll, ul, lr, ur) + + # Order as (ll, lr, ul, ur, center) to match the rest of this class. + idx = (ll, lr, ul, ur, 0) + points = tuple(galsim.PositionD(x_pos[i], y_pos[i]) for i in idx) + + # Build the solver matrix for basis [1, x, y, xy, x^2+y^2]. + basis = np.empty((5, 5), dtype=float) + px = np.array([x_pos[i] for i in idx], dtype=float) + py = np.array([y_pos[i] for i in idx], dtype=float) + basis[0,:] = 1.0 + basis[1,:] = px + basis[2,:] = py + basis[3,:] = px * py + basis[4,:] = px * px + py * py + # If f(x,y) = c·phi(x,y), where phi(x,y) = [1, x, y, xy, x^2+y^2], + # then sampled values satisfy s = B c, where B[:,k] = phi(x_k,y_k). + # For any target point, weights w that interpolate from + # sample values are w = B^{-1} phi(target), so that f(target) = w·s. + solver = np.linalg.inv(basis) + + # GalSim uses bilinear nominal interpolation, so correct center by adding this delta. + bilinear_center = galsim.roman.roman_psfs._interp_aberrations_bilinear( + aberrations, x_pos, y_pos, center + ) + delta = np.array(aberrations[0] - bilinear_center, dtype=float) + center_delta = np.zeros(self.max_zernike + 1, dtype=float) + center_delta[4:] = delta[4:self.max_zernike + 1] + + out = {'points': points, 'solver': solver, 'center_delta': center_delta} + self._roman_five_point_data[sca] = out + return out + def _make_extra_aberrations(self, params): # GalSim expects extra_aberrations indexed by Zernike number. Our parameter vector # corresponds to z4..z_max, so pad indices 0..3. @@ -236,7 +313,7 @@ def _parse_prior_sigma(self, aberration_prior_sigma): if aberration_prior_sigma is None: return None sigma = np.array(aberration_prior_sigma, dtype=float).ravel() - nprior = self.param_len + nprior = self._single_point_param_len if sigma.size == 1: sigma = np.full(nprior, sigma[0], dtype=float) elif sigma.size != nprior: @@ -247,6 +324,23 @@ def _parse_prior_sigma(self, aberration_prior_sigma): raise ValueError("aberration_prior_sigma values must all be > 0") return sigma + def _expand_prior_sigma(self, prior_sigma): + # Possibly tile the prior_sigma array for the (a, b, c) coefficient blocks. + if self.aberration_interp == 'linear' and prior_sigma is not None: + return np.tile(prior_sigma, 3) + else: + return prior_sigma + + def _evaluate_linear_params(self, params, points): + coeffs = params.reshape(3, self._single_point_param_len) + a, b, c = coeffs + sample_params = np.empty((len(points), self._single_point_param_len), dtype=float) + for i, pt in enumerate(points): + xh = pt.x / self.sca_size + yh = pt.y / self.sca_size + sample_params[i] = a + b * xh + c * yh + return sample_params + def _apply_prior(self, ata, atb, params): if self.prior_invsigsq is None: return @@ -262,7 +356,7 @@ def initialize(self, star, logger=None, default_init=None): def fit(self, star, logger=None, convert_func=None, draw_method=None): # We usually batch the stars in groups by SCA and do the fitting calculation - # together for efficiency of building the adjusted corner profiles. So this + # together for efficiency of building the adjusted sample-point profiles. So this # function is rarely used. However, it's part of the Interp API, so just # farm this out to the fit_many function with a list of one star. return self.fit_many( @@ -274,14 +368,14 @@ def fit(self, star, logger=None, convert_func=None, draw_method=None): @staticmethod def _fit_sca_group_worker(model, sca, stars, convert_funcs, draw_method, logger): - fit_group, mean_params, corner_profiles = model._fit_sca_group( + fit_group, mean_params, sample_profiles = model._fit_sca_group( sca, stars, convert_funcs=convert_funcs, logger=logger, draw_method=draw_method, ) - return sca, fit_group, mean_params, corner_profiles + return sca, fit_group, mean_params, sample_profiles def fit_many(self, stars, logger=None, convert_funcs=None, draw_method=None): logger = LoggerWrapper(logger) @@ -327,11 +421,11 @@ def fit_many(self, stars, logger=None, convert_funcs=None, draw_method=None): out = [None] * len(stars) sca_mean = {} for fit_result in fit_results: - sca, fit_group, mean_params, corner_profiles = fit_result + sca, fit_group, mean_params, sample_profiles = fit_result indices = grouped[sca]['indices'] sca_mean[sca] = mean_params wcs = stars[indices[0]].image.wcs - self._corner_cache[sca] = (mean_params, wcs, corner_profiles) + self._corner_cache[sca] = (mean_params, wcs, sample_profiles) # Output the stars in the same order as input, so they stay matched with # convert_funcs array if there is one. for i, star in zip(indices, fit_group): @@ -343,9 +437,9 @@ def _fit_sca_group(self, sca, stars, convert_funcs, logger=None, draw_method=Non ref = stars[0] params = ref.fit.get_params(self._num) - # We're about to change params, so if the corner profiles aren't yet in the cache, + # We're about to change params, so if the sample profiles aren't yet in the cache, # there is not reason now to add them. - corner_base_profiles = self._get_corner_profiles(ref, params, cache=False, sca=sca) + sample_base_profiles = self._get_sample_profiles(ref, params, cache=False, sca=sca) # First draw the base image for each star. # Also set up an empty Jacobian array for each to be filled in below. @@ -354,9 +448,9 @@ def _fit_sca_group(self, sca, stars, convert_funcs, logger=None, draw_method=Non for star, convert_func in zip(stars, convert_funcs): image = star.image.array weight = star.weight.array - base_image = self._draw_model_image_from_corners( + base_image = self._draw_model_image_from_samples( star, - corner_base_profiles, + sample_base_profiles, convert_func=convert_func, ) resid = image - base_image @@ -370,13 +464,13 @@ def _fit_sca_group(self, sca, stars, convert_funcs, logger=None, draw_method=Non for i, step in enumerate(steps): p1 = params.copy() p1[i] += step - corner_plus_profiles = self._get_corner_profiles(ref, p1, cache=False, sca=sca) + sample_plus_profiles = self._get_sample_profiles(ref, p1, cache=False, sca=sca) scale = 1.0 / step - # Use these adjusted corner profiles for all stars on the SCA. + # Use these adjusted sample profiles for all stars on the SCA. for star, convert_func, _, _, _, jac, base_image in group_data: - im1 = self._draw_model_image_from_corners( + im1 = self._draw_model_image_from_samples( star, - corner_plus_profiles, + sample_plus_profiles, convert_func=convert_func, ) jac[:, i] = ((im1 - base_image) * scale).ravel() @@ -393,13 +487,13 @@ def _fit_sca_group(self, sca, stars, convert_funcs, logger=None, draw_method=Non # Solve once for a single SCA parameter vector. sca_params, var = self._solve_normal_equations(ata_sum, atb_sum, params) - corner_sca_profiles = self._get_corner_profiles(ref, sca_params, cache=False, sca=sca) + sample_sca_profiles = self._get_sample_profiles(ref, sca_params, cache=False, sca=sca) # Compute the final model and chisq for each star. out = [] for star, convert_func, weight in solved: - model_image = self._draw_model_image_from_corners( - star, corner_sca_profiles, convert_func=convert_func + model_image = self._draw_model_image_from_samples( + star, sample_sca_profiles, convert_func=convert_func ) chisq = np.sum(weight * (star.image.array - model_image) ** 2) dof = np.count_nonzero(weight) - 3 @@ -411,7 +505,7 @@ def _fit_sca_group(self, sca, stars, convert_funcs, logger=None, draw_method=Non ), ) ) - return out, sca_params, corner_sca_profiles + return out, sca_params, sample_sca_profiles def draw(self, star, copy_image=True): params = star.fit.get_params(self._num) @@ -431,11 +525,11 @@ def getProfile(self, params=None, star=None, cache=True): raise ValueError( "Roman params must have length %d (got %d)" % (self.param_len, params.size) ) - corner_profiles = self._get_corner_profiles(star, params, cache=cache) - return self._interpolate_corners(star, corner_profiles) + sample_profiles = self._get_sample_profiles(star, params, cache=cache) + return self._interpolate_samples(star, sample_profiles) - def _draw_model_image_from_corners(self, star, corner_profiles, convert_func=None): - prof = self._interpolate_corners(star, corner_profiles) + def _draw_model_image_from_samples(self, star, sample_profiles, convert_func=None): + prof = self._interpolate_samples(star, sample_profiles) prof = prof.shift(star.fit.center) * star.fit.flux if convert_func is not None: prof = convert_func(prof) @@ -443,7 +537,7 @@ def _draw_model_image_from_corners(self, star, corner_profiles, convert_func=Non self._draw_profile_to_image(prof, image, star.image_pos) return image.array - def _get_corner_profiles(self, star, params, cache=True, sca=None): + def _get_sample_profiles(self, star, params, cache=True, sca=None): if sca is None: sca = _get_sca(star) wcs = star.image.wcs @@ -456,33 +550,51 @@ def _get_corner_profiles(self, star, params, cache=True, sca=None): return cached_profiles wavelength = None if self.chromatic else self.bandpass.effective_wavelength - corners = ( - galsim.PositionD(0.0, 0.0), - galsim.PositionD(self.sca_size, 0.0), - galsim.PositionD(0.0, self.sca_size), - galsim.PositionD(self.sca_size, self.sca_size), - ) - corner_params = np.tile(params, 4).reshape(4, self.param_len) - profiles = tuple( - galsim.roman.getPSF( - sca, - self.filter, - SCA_pos=corner, - pupil_bin=pupil_bin, - wcs=wcs, - extra_aberrations=self._make_extra_aberrations(p), - wavelength=wavelength, + if self.nominal_interp == 'five_point': + points = self._get_roman_five_point_data(sca)['points'] + else: + points = ( + galsim.PositionD(0.0, 0.0), + galsim.PositionD(self.sca_size, 0.0), + galsim.PositionD(0.0, self.sca_size), + galsim.PositionD(self.sca_size, self.sca_size), ) - for corner, p in zip(corners, corner_params) + if self.aberration_interp == 'linear': + sample_params = self._evaluate_linear_params(params, points) + else: + npts = 5 if self.nominal_interp == 'five_point' else 4 + # In constant/global mode, the same per-SCA vector applies everywhere. + sample_params = np.tile(params, npts).reshape(npts, self._single_point_param_len) + profiles = [] + for i, (pt, p) in enumerate(zip(points, sample_params)): + extra = self._make_extra_aberrations(p) + if self.nominal_interp == 'five_point' and i == 4: + extra = extra + self._get_roman_five_point_data(sca)['center_delta'] + profiles.append( + galsim.roman.getPSF( + sca, + self.filter, + SCA_pos=pt, + pupil_bin=pupil_bin, + wcs=wcs, + extra_aberrations=extra, + wavelength=wavelength, + ) ) + profiles = tuple(profiles) if cache: self._corner_cache[sca] = (params, wcs, profiles) return profiles - def _interpolate_corners(self, star, corner_profiles): - w_ll, w_lr, w_ul, w_ur = self._corner_weights(star) - ll, lr, ul, ur = corner_profiles - return w_ll * ll + w_lr * lr + w_ul * ul + w_ur * ur + def _interpolate_samples(self, star, sample_profiles): + if self.nominal_interp == 'five_point': + w_ll, w_lr, w_ul, w_ur, w_c = self._five_point_weights(star) + ll, lr, ul, ur, c = sample_profiles + return w_ll * ll + w_lr * lr + w_ul * ul + w_ur * ur + w_c * c + else: + w_ll, w_lr, w_ul, w_ur = self._corner_weights(star) + ll, lr, ul, ur = sample_profiles + return w_ll * ll + w_lr * lr + w_ul * ul + w_ur * ur def _corner_weights(self, star): x = np.clip(star.image_pos.x, 0.0, self.sca_size) @@ -496,6 +608,16 @@ def _corner_weights(self, star): fx * fy, ) + def _five_point_weights(self, star): + sca = _get_sca(star) + solver = self._get_roman_five_point_data(sca)['solver'] + x = star.image_pos.x + y = star.image_pos.y + phi = np.array([1.0, x, y, x * y, x * x + y * y], dtype=float) + # Convert basis values at this star position into interpolation weights on the + # 5 sample locations (LL, LR, UL, UR, C). + return tuple(solver.dot(phi)) + class RomanOpticsPSF(PSF): """A PSF wrapper for Roman optical fits with configurable aberration interpolation.""" @@ -569,7 +691,7 @@ def parseKwargs(cls, config_psf, logger): def initialize_params(self, stars, logger=None, default_init=None): nremoved = 0 - # Start each new fit with a fresh corner-profile cache. + # Start each new fit with a fresh sample-profile cache. self.model.clear_cache() logger.debug("Initializing models") diff --git a/tests/test_roman.py b/tests/test_roman.py index 8af0db43..e45651bb 100644 --- a/tests/test_roman.py +++ b/tests/test_roman.py @@ -58,6 +58,7 @@ def test_roman_optics(): assert psf.fit_center is False assert psf.include_model_centroid is False assert psf.model.aberration_interp == 'constant' + assert psf.model.nominal_interp == 'bilinear' assert psf.interp.per_sca is True # Global mode uses one aberration vector for full focal plane. @@ -196,8 +197,12 @@ def traced_get_psf(*args, **kwargs): RomanOpticalModel(filter='H158', chromatic=False, max_zernike=6, aberration_interp='bad') assert "must be one of" in str(err.value) + with pytest.raises(ValueError) as err: + RomanOpticalModel(filter='H158', chromatic=False, max_zernike=6, nominal_interp='bad') + assert "must be one of" in str(err.value) + @timer -def test_roman_corner_cache(): +def test_corner_cache(): """Verify corner-profile caching reuses one 4-corner set for same SCA and params. """ with fast_pupil_bin(): @@ -226,16 +231,90 @@ def test_roman_corner_cache(): psf.drawStar(stars[1]) params = stars[0].fit.params - profiles1 = psf.model._get_corner_profiles(stars[0], params, cache=True) - profiles2 = psf.model._get_corner_profiles(stars[1], params, cache=True) + profiles1 = psf.model._get_sample_profiles(stars[0], params, cache=True) + profiles2 = psf.model._get_sample_profiles(stars[1], params, cache=True) assert profiles1 is profiles2 assert len(psf.model._corner_cache) == 1 assert 5 in psf.model._corner_cache assert len(psf.model._corner_cache[5][2]) == 4 + # In five-point mode, the cache stores four corners plus center. + psf5 = piff.PSF.process( + { + 'type': 'RomanOptics', + 'filter': 'H158', + 'chromatic': False, + 'max_zernike': 6, + 'nominal_interp': 'five_point', + } + ) + stars5, _ = psf5.initialize_params(stars, logger=logger) + psf5.drawStar(stars5[0]) + params5 = stars5[0].fit.params + profiles5 = psf5.model._get_sample_profiles(stars5[0], params5, cache=True) + assert len(profiles5) == 5 + assert len(psf5.model._corner_cache[5][2]) == 5 + + +@timer +def test_five_point_weights(): + """Check five-point interpolation weights for corner/center and quadratic basis terms. + """ + model = RomanOpticalModel( + filter='H158', + chromatic=False, + max_zernike=6, + nominal_interp='five_point', + ) + sca = 7 + size = model.sca_size + + def make_star(x, y): + return piff.Star.makeTarget( + x=x, + y=y, + stamp_size=25, + scale=0.11, + properties={'sca': sca}, + ).withFlux(1.0, (0.0, 0.0)) + + fp = model._get_roman_five_point_data(sca) + points = fp['points'] + + # Sampling locations used by the interpolation should have a single w=1. + w_ll = model._five_point_weights(make_star(points[0].x, points[0].y)) + w_lr = model._five_point_weights(make_star(points[1].x, points[1].y)) + w_ul = model._five_point_weights(make_star(points[2].x, points[2].y)) + w_ur = model._five_point_weights(make_star(points[3].x, points[3].y)) + w_c = model._five_point_weights(make_star(points[4].x, points[4].y)) + np.testing.assert_allclose(w_ll, [1, 0, 0, 0, 0], atol=1.e-11, rtol=0.0) + np.testing.assert_allclose(w_lr, [0, 1, 0, 0, 0], atol=1.e-11, rtol=0.0) + np.testing.assert_allclose(w_ul, [0, 0, 1, 0, 0], atol=1.e-11, rtol=0.0) + np.testing.assert_allclose(w_ur, [0, 0, 0, 1, 0], atol=1.e-11, rtol=0.0) + np.testing.assert_allclose(w_c, [0, 0, 0, 0, 1], atol=1.e-11, rtol=0.0) + + # Verify exact reproduction of [1, x, y, xy, x^2+y^2] basis under this interpolation. + sample_xy = np.array([(p.x, p.y) for p in points], dtype=float) + + def basis_values(x, y): + return np.array([1.0, x, y, x * y, x * x + y * y], dtype=float) + + sample_basis = np.array([basis_values(x, y) for x, y in sample_xy]) + test_positions = [ + (0.13 * size, 0.27 * size), + (0.82 * size, 0.39 * size), + (0.61 * size, 0.91 * size), + ] + for x, y in test_positions: + star = make_star(x, y) + w = np.array(model._five_point_weights(star)) + expected = basis_values(x, y) + interp = w.dot(sample_basis) + np.testing.assert_allclose(interp, expected, atol=1.e-8, rtol=0.0) + @timer -def test_roman_fit(): +def test_fit(): """Check local convergence of single-star Roman fits in constant mode. """ with fast_pupil_bin(): @@ -253,25 +332,14 @@ def test_roman_fit(): properties={'sca': 5}, ).withFlux(1.0, (0.0, 0.0)) - init_star = model.initialize(star) + star = model.initialize(star) # Keep injected extra aberrations modest relative to baseline Roman optics. # Typical built-in |z4..z22| values are around 1e-3 to a few e-2, so 4-5e-3 is # realistic while still providing measurable signal in this unit test. truth_params = np.array([0.004, -0.003, 0.005]) - truth_fit = init_star.fit.newParams( - truth_params, - params_var=np.zeros_like(truth_params), - ) - truth_star = model.draw(piff.Star(init_star.data, truth_fit)) - - fit_star = piff.Star( - truth_star.data, - init_star.fit.newParams( - np.zeros_like(truth_params), - params_var=np.zeros_like(truth_params) - ), - ) - fitted = model.fit(fit_star) + prof = model.getProfile(truth_params, star=star) + model._draw_profile_to_image(prof * star.fit.flux, star.image, star.image_pos) + fitted = model.fit(star) print('fit = ',fitted.fit.params) # 1 pass isn't great, but after 2 passes, the agreement is sub percent. fitted = model.fit(fitted) @@ -290,10 +358,10 @@ def test_roman_fit(): assert np.all(fitted_var >= 0) # Sanity tests about getting profile and drawing it. - prof = model.getProfile(star=fit_star) + prof = model.getProfile(star=star) assert prof is not None - drawn_star = model.draw(fit_star) - assert drawn_star.image.array.shape == fit_star.image.array.shape + drawn_star = model.draw(star) + assert drawn_star.image.array.shape == star.image.array.shape # Error if no star argument in getProfile. (Allowed by some other classes.) with pytest.raises(ValueError) as err: @@ -302,17 +370,17 @@ def test_roman_fit(): # Error if params is given and has wrong length. with pytest.raises(ValueError) as err: - model.getProfile(star=fit_star, params=np.zeros(model.param_len+2)) + model.getProfile(star=star, params=np.zeros(model.param_len+2)) assert "params must have length 3" in str(err.value) # Error if convert_funcs is given but has different length than stars. with pytest.raises(ValueError) as err: - model.fit_many([fit_star], convert_funcs=[]) + model.fit_many([star], convert_funcs=[]) assert "len(convert_funcs) must match len(stars)" in str(err.value) @timer -def test_roman_aberration_prior(): +def test_aberration_prior(): """Validate the use of priors on the aberration values. """ # Check that fitted aberrations stay closer to 0 when strong prior is applied. @@ -379,6 +447,33 @@ def test_roman_aberration_prior(): aberration_prior_sigma=[]) assert "scalar or length 3" in str(err.value) + # Same for linear mode (in particular the allowed length is 3 here, not 9, + # which is the full param_len). + with pytest.raises(ValueError) as err: + piff.roman.RomanOpticalModel(filter='H158', chromatic=False, max_zernike=6, + aberration_interp='linear', + aberration_prior_sigma=[1.0] * 9) + assert "scalar or length 3" in str(err.value) + + # In linear mode, scalar and one-point vectors tile to the (a, b, c) coefficient blocks. + linear_scalar = piff.roman.RomanOpticalModel( + filter='H158', + chromatic=False, + max_zernike=6, + aberration_interp='linear', + aberration_prior_sigma=0.05, + ) + np.testing.assert_allclose(linear_scalar.prior_sigma, np.full(9, 0.05)) + + linear_vec = piff.roman.RomanOpticalModel( + filter='H158', + chromatic=False, + max_zernike=6, + aberration_interp='linear', + aberration_prior_sigma=[0.1, 0.2, 0.3], + ) + np.testing.assert_allclose(linear_vec.prior_sigma, np.tile([0.1, 0.2, 0.3], 3)) + # priors cannot be <= 0 with pytest.raises(ValueError) as err: RomanOpticalModel(filter='H158', chromatic=False, max_zernike=6, @@ -392,7 +487,73 @@ def test_roman_aberration_prior(): @timer -def test_roman_fit_many(): +def test_linear_prior_io(): + """Ensure linear-mode aberration priors serialize compactly and round-trip correctly. + """ + psf = piff.PSF.process( + { + 'type': 'RomanOptics', + 'filter': 'H158', + 'chromatic': False, + 'max_zernike': 6, + 'aberration_interp': 'linear', + 'aberration_prior_sigma': [0.11, 0.22, 0.33], + } + ) + + fn = os.path.join('output', 'roman_linear_prior_io.piff') + psf.write(fn) + psf2 = piff.read(fn) + + model = psf2.model + np.testing.assert_allclose(model.orig_prior_sigma, + [0.11, 0.22, 0.33], atol=0.0, rtol=0.0) + np.testing.assert_allclose( + model.prior_sigma, + np.tile([0.11, 0.22, 0.33], 3), + atol=0.0, + rtol=0.0, + ) + np.testing.assert_allclose( + model.kwargs['aberration_prior_sigma'], + [0.11, 0.22, 0.33], + atol=0.0, + rtol=0.0, + ) + + # Also check scalar prior in input config. + psf_scalar = piff.PSF.process( + { + 'type': 'RomanOptics', + 'filter': 'H158', + 'chromatic': False, + 'max_zernike': 6, + 'aberration_interp': 'linear', + 'aberration_prior_sigma': 0.05, + } + ) + fn_scalar = os.path.join('output', 'roman_linear_prior_io_scalar.piff') + psf_scalar.write(fn_scalar) + psf_scalar2 = piff.read(fn_scalar) + model_scalar = psf_scalar2.model + np.testing.assert_allclose(model_scalar.orig_prior_sigma, + [0.05, 0.05, 0.05], atol=0.0, rtol=0.0) + np.testing.assert_allclose( + model_scalar.prior_sigma, + np.full(9, 0.05), + atol=0.0, + rtol=0.0, + ) + np.testing.assert_allclose( + model_scalar.kwargs['aberration_prior_sigma'], + [0.05, 0.05, 0.05], + atol=0.0, + rtol=0.0, + ) + + +@timer +def test_fit_many(): """Check accuracy of fitting multiple stars using fit_many. """ with fast_pupil_bin(): @@ -419,39 +580,30 @@ def test_roman_fit_many(): ).withFlux(1.0, (0.0, 0.0)), ] stars = [model.initialize(s) for s in stars] - # Use the same realistic small-amplitude extra-aberration vector as test_roman_fit. + # Use the same realistic small-amplitude extra-aberration vector as test_fit. truth_params = np.array([0.004, -0.003, 0.005]) - truth = [ - model.draw( - piff.Star( - s.data, - s.fit.newParams( - truth_params, - params_var=np.zeros_like(truth_params), - ), - ) - ) - for s in stars - ] - fit_stars = [ + stars = [ piff.Star( s.data, - stars[i].fit.newParams( + s.fit.newParams( np.zeros_like(truth_params), params_var=np.zeros_like(truth_params), ), ) - for i, s in enumerate(truth) + for s in stars ] + for star in stars: + prof = model.getProfile(truth_params, star=star) + model._draw_profile_to_image(prof * star.fit.flux, star.image, star.image_pos) # 1 pass isn't great, but after 2 passes, the agreement is sub percent. # And 3 is within 0.1% agreement. for _ in range(3): - fit_stars = model.fit_many(fit_stars) - print('fit[0] => ',fit_stars[0].fit.params) + stars = model.fit_many(stars) + print('fit[0] => ',stars[0].fit.params) - p0 = fit_stars[0].fit.params - p1 = fit_stars[1].fit.params + p0 = stars[0].fit.params + p1 = stars[1].fit.params print('final p0 = ',p0) print('final p1 = ',p1) print(' truth = ',truth_params) @@ -460,7 +612,7 @@ def test_roman_fit_many(): @timer -def test_roman_fit_many_nproc(): +def test_fit_many_nproc(): """Check `fit_many` multiprocessing path and accuracy with `nproc > 1`. """ with fast_pupil_bin(): @@ -489,40 +641,359 @@ def test_roman_fit_many_nproc(): ] stars = [model.initialize(s) for s in stars] truth_params = np.array([0.004, -0.003, 0.005]) - truth = [ - model.draw( - piff.Star( - s.data, - s.fit.newParams( - truth_params, - params_var=np.zeros_like(truth_params), - ), - ) - ) - for s in stars - ] - fit_stars = [ + stars = [ piff.Star( s.data, - stars[i].fit.newParams( + s.fit.newParams( np.zeros_like(truth_params), params_var=np.zeros_like(truth_params), ), ) - for i, s in enumerate(truth) + for s in stars ] + for star in stars: + prof = model.getProfile(truth_params, star=star) + model._draw_profile_to_image(prof * star.fit.flux, star.image, star.image_pos) for _ in range(3): - fit_stars = model.fit_many(fit_stars) + stars = model.fit_many(stars) - assert [int(s['sca']) for s in fit_stars] == [5, 5] - for s in fit_stars: + assert [int(s['sca']) for s in stars] == [5, 5] + for s in stars: np.testing.assert_allclose(s.fit.params, truth_params, atol=0.0, rtol=2.e-3) assert list(model._corner_cache.keys()) == [5] @timer -def test_roman_optics_convert_funcs(): +def test_bilinear_vs_five_point(): + """Exercise practical fit behavior and quantify bilinear vs five-point differences. + """ + with fast_pupil_bin(): + # Make 5 stars at various points around an SCA using GalSim. + filt = 'H158' + sca = 8 + size = galsim.roman.n_pix + + def make_star(x, y): + return piff.Star.makeTarget( + x=x, + y=y, + stamp_size=25, + scale=0.11, + properties={'sca': sca}, + ).withFlux(1.0, (0.0, 0.0)) + + positions = [ + (0.12 * size, 0.18 * size), + (0.33 * size, 0.74 * size), + (0.71 * size, 0.21 * size), + (0.83 * size, 0.66 * size), + (0.49 * size, 0.52 * size), + ] + stars = [make_star(x, y) for x, y in positions] + truth_params = np.array([0.004, -0.003, 0.005]) + extra_truth = np.concatenate(([0,0,0,0], truth_params)) + wavelength = galsim.roman.getBandpasses()[filt].effective_wavelength + for st in stars: + prof = galsim.roman.getPSF( + int(st['sca']), filt, + SCA_pos=st.image_pos, + pupil_bin=piff.roman.roman_psf.pupil_bin, + wcs=st.image.wcs, + extra_aberrations=extra_truth, + wavelength=wavelength, + ) + prof.drawImage(st.image, center=st.image_pos) + + # Fit these stars with both models in Piff. Bilinear and 5-point. + # Note: GalSim internally does bilinear interpolation of the nominal aberrations. + # This is not the same bilinear that we do in Piff -- we do either bilinear or + # 5-point interpolation of the profiles. These are not equivalent. + model_bilinear = RomanOpticalModel( + filter=filt, + chromatic=False, + max_zernike=6, + nominal_interp='bilinear', + aberration_prior_sigma=1.0e6, + ) + model_five = RomanOpticalModel( + filter=filt, + chromatic=False, + max_zernike=6, + nominal_interp='five_point', + aberration_prior_sigma=1.0e6, + ) + fit_bilinear = [model_bilinear.initialize(st) for st in stars] + fit_five = [model_five.initialize(st) for st in stars] + for it in range(3): + fit_bilinear = model_bilinear.fit_many(fit_bilinear) + fit_five = model_five.fit_many(fit_five) + print(f'iter {it+1} bilinear fit[0] = ', fit_bilinear[0].fit.params) + print(f'iter {it+1} five-point fit[0] = ', fit_five[0].fit.params) + print( + f'iter {it+1} chisq (bilinear, five-point) = ', + sum(st.fit.chisq for st in fit_bilinear), + sum(st.fit.chisq for st in fit_five), + ) + + # Compare the results. + chisq_bilinear = sum(st.fit.chisq for st in fit_bilinear) + chisq_five = sum(st.fit.chisq for st in fit_five) + params_bilinear = fit_bilinear[0].fit.params + params_five = fit_five[0].fit.params + print('bilinear chisq = ', chisq_bilinear) + print('five-point chisq = ', chisq_five) + print('final bilinear fit[0] = ', params_bilinear) + print('final five-point fit[0] = ', params_five) + print('truth params = ', truth_params) + + # With direct GalSim truth generation, both nominal interpolation options should converge + # to reasonable fits, but neither is expected to recover truth_params exactly + # for the reason mentioned above. (Profile interpolation != aberration interpolation.) + # Interesting that it turns out that bilinear gets the model parameters somewhat + # closer to the truth, but the chisq for five_point is smaller. + np.testing.assert_allclose(params_bilinear, truth_params, atol=1.5e-3) + np.testing.assert_allclose(params_five, truth_params, atol=4.e-3) + assert chisq_bilinear < 1.0e-6 + assert chisq_five < 4.0e-7 + # Quantify that the two nominal interpolation choices are observably different. + assert abs(chisq_five - chisq_bilinear) > 1.e-7 + + # Quantify direct model-image difference at the same position and params. + sample = make_star(0.63 * size, 0.37 * size) + image_b = model_bilinear.draw(model_bilinear.initialize(sample)).image.array + image_f = model_five.draw(model_five.initialize(sample)).image.array + mean_abs_diff = np.mean(np.abs(image_b - image_f)) + print('mean abs image difference = ', mean_abs_diff) + assert mean_abs_diff > 1.e-6 + + +@timer +def test_fit_linear(): + """Check linear-mode convergence with stars spanning the SCA geometry. + """ + with fast_pupil_bin(): + # Use stars spread across the SCA so the constant and slope terms are constrained. + pos = [ + (512.0, 512.0), + (2044.0, 512.0), + (3576.0, 512.0), + (512.0, 2044.0), + (2044.0, 2044.0), + (3576.0, 2044.0), + (512.0, 3576.0), + (2044.0, 3576.0), + (3576.0, 3576.0), + ] + truth_params = np.array([ + 0.004, -0.003, 0.005, + 0.0, 0.0, 0.0, + 0.0, 0.0, 0.0, + ]) + for nominal_interp in ['bilinear', 'five_point']: + model = piff.roman.RomanOpticalModel( + filter='H158', + chromatic=False, + max_zernike=6, + aberration_interp='linear', + nominal_interp=nominal_interp, + aberration_prior_sigma=1.0e6, + ) + stars = [ + piff.Star.makeTarget( + x=x, + y=y, + stamp_size=25, + scale=0.11, + properties={'sca': 5}, + ).withFlux(1.0, (0.0, 0.0)) + for x, y in pos + ] + stars = [model.initialize(star) for star in stars] + for star in stars: + prof = model.getProfile(truth_params, star=star) + model._draw_profile_to_image(prof * star.fit.flux, star.image, star.image_pos) + + for _ in range(4): + stars = model.fit_many(stars) + + fitted_params = stars[0].fit.params + fitted_var = stars[0].fit.params_var + + for star in stars: + assert star.fit.chisq >= 0 + assert star.fit.dof > 0 + np.testing.assert_allclose(star.fit.params, fitted_params, atol=1.e-12, rtol=0.0) + print(f'linear {nominal_interp} fit = ', fitted_params) + print(' truth = ', truth_params) + np.testing.assert_allclose(fitted_params[:3], truth_params[:3], atol=0.0, rtol=1.e-3) + np.testing.assert_allclose(fitted_params[3:], truth_params[3:], atol=5.e-9, rtol=0.0) + assert np.all(fitted_var >= 0) + + for star in stars: + model_star = model.draw(star) + np.testing.assert_allclose( + model_star.image.array, star.image.array, atol=5.e-5, rtol=0.0 + ) + + +@timer +def test_fit_linear_gradient(): + """Check linear-mode recovery when aberrations vary as a + b x + c y across the SCA. + """ + with fast_pupil_bin(): + model = piff.roman.RomanOpticalModel( + filter='H158', + chromatic=False, + max_zernike=6, + aberration_interp='linear', + aberration_prior_sigma=1.0e6, + ) + pos = [ + (512.0, 512.0), + (2044.0, 512.0), + (3576.0, 512.0), + (512.0, 2044.0), + (2044.0, 2044.0), + (3576.0, 2044.0), + (512.0, 3576.0), + (2044.0, 3576.0), + (3576.0, 3576.0), + ] + stars = [ + piff.Star.makeTarget( + x=x, + y=y, + stamp_size=25, + scale=0.11, + properties={'sca': 5}, + ).withFlux(1.0, (0.0, 0.0)) + for x, y in pos + ] + stars = [model.initialize(star) for star in stars] + + # Define each Zernike truth as a + b*(x/4088) + c*(y/4088). + abc = np.array( + [ + [0.004, 2.e-5, -7.e-5], + [-0.003, 6.e-5, -2.e-5], + [0.005, -3.e-5, -5.e-5], + ] + ) + + def eval_truth(x, y): + xh = x / model.sca_size + yh = y / model.sca_size + return abc[:, 0] + abc[:, 1] * xh + abc[:, 2] * yh + + truth_params = np.concatenate([abc[:, 0], abc[:, 1], abc[:, 2]]) + corner_truth = np.array([ + eval_truth(0.0, 0.0), + eval_truth(model.sca_size, 0.0), + eval_truth(0.0, model.sca_size), + eval_truth(model.sca_size, model.sca_size), + ]) + assert not np.any(np.isclose(corner_truth[0], corner_truth[1])) + assert not np.any(np.isclose(corner_truth[0], corner_truth[2])) + assert not np.any(np.isclose(corner_truth[0], corner_truth[3])) + + # Verify independently that bilinear interpolation of corners reproduces the + # per-star truth function for a purely planar field. + ll, lr, ul, ur = corner_truth + for x, y in pos: + fx = x / model.sca_size + fy = y / model.sca_size + wll = (1.0 - fx) * (1.0 - fy) + wlr = fx * (1.0 - fy) + wul = (1.0 - fx) * fy + wur = fx * fy + local_from_corners = wll * ll + wlr * lr + wul * ul + wur * ur + np.testing.assert_allclose( + local_from_corners, eval_truth(x, y), atol=1.e-12, rtol=0.0 + ) + + # Draw the stars to use for fitting using direct GalSim Roman PSFs at each star + # position, with the local extra-aberration vector from eval_truth. + for star in stars: + aber = eval_truth(star.image_pos.x, star.image_pos.y) + prof = galsim.roman.getPSF(5, 'H158', star.image_pos, + pupil_bin=piff.roman.roman_psf.pupil_bin, + wcs=star.image.wcs, + extra_aberrations=model._make_extra_aberrations(aber), + wavelength=model.bandpass.effective_wavelength) + prof.drawImage(star.image, method='auto', center=star.image_pos) + + for _ in range(5): + stars = model.fit_many(stars) + + # Note: these are not expected to match exactly. The real images include the + # natural variation within the SCA from the roman aberration pattern, fully separate + # from the extra_aberrations we're fitting for. This variation is not quite linear, + # so the linear approximation is a small model mismatch. However, it's relatively + # close in the fitted extra aberrations, and the drawn images are very close. + fitted_params = stars[0].fit.params + print('linear gradient fit = ', fitted_params) + print(' truth = ', truth_params) + np.testing.assert_allclose(fitted_params, truth_params, atol=3.e-4, rtol=0.3) + for star in stars: + np.testing.assert_allclose(star.fit.params, fitted_params) + for star in stars: + model_star = model.draw(star) + np.testing.assert_allclose( + model_star.image.array, star.image.array, atol=1.e-4, rtol=0.002) + + +@timer +def test_fit_linear_nproc(): + """Check linear-mode fit_many behavior with multiprocessing enabled. + """ + with fast_pupil_bin(): + model = piff.roman.RomanOpticalModel( + filter='H158', + chromatic=False, + max_zernike=6, + aberration_interp='linear', + aberration_prior_sigma=1.0e6, + nproc=2, + ) + pos = [ + (512.0, 512.0), + (2044.0, 512.0), + (3576.0, 512.0), + (512.0, 3576.0), + (2044.0, 3576.0), + (3576.0, 3576.0), + ] + stars = [ + piff.Star.makeTarget( + x=x, + y=y, + stamp_size=25, + scale=0.11, + properties={'sca': 5}, + ).withFlux(1.0, (0.0, 0.0)) + for x, y in pos + ] + stars = [model.initialize(star) for star in stars] + truth_params = np.array([ + 0.004, -0.003, 0.005, + 0.0, 0.0, 0.0, + 0.0, 0.0, 0.0, + ]) + for star in stars: + prof = model.getProfile(truth_params, star=star) + model._draw_profile_to_image(prof * star.fit.flux, star.image, star.image_pos) + + for _ in range(4): + stars = model.fit_many(stars) + for star in stars: + np.testing.assert_allclose(star.fit.params[:3], truth_params[:3], atol=0.0, rtol=1.e-3) + np.testing.assert_allclose(star.fit.params[3:], truth_params[3:], atol=5.e-9, rtol=0.0) + + + +@timer +def test_optics_convert_funcs(): """Check aberration recovery when fitting with a nontrivial convert_func (profile shear). """ with fast_pupil_bin(): @@ -560,38 +1031,30 @@ def apply_shear(prof): return prof.shear(g1=0.01, g2=-0.005) convert_funcs = [apply_shear] * len(stars) - fit_stars = [] for star in stars: - truth_fit = star.fit.newParams( - truth_params, - params_var=np.zeros_like(truth_params), - ) - truth_star = piff.Star(star.data, truth_fit) # True profile is sheared version of optical PSF. - prof = psf.model.getProfile(truth_params, star=truth_star).shear(g1=0.01, g2=-0.005) - image = star.image.copy() - psf.model._draw_profile_to_image(prof, image, star.image_pos) - fit_stars.append(piff.Star(star.data.withNew(image=image), star.fit)) + prof = psf.model.getProfile(truth_params, star=star).shear(g1=0.01, g2=-0.005) + psf.model._draw_profile_to_image(prof, star.image, star.image_pos) for _ in range(3): - fit_stars, nremoved = psf.single_iteration( - fit_stars, + stars, nremoved = psf.single_iteration( + stars, logger=logger, convert_funcs=convert_funcs, draw_method=None, ) assert nremoved == 0 - print('params[0] => ',fit_stars[0].fit.params) + print('params[0] => ',stars[0].fit.params) - assert len(fit_stars) == len(stars) + assert len(stars) == len(convert_funcs) print('truth = ',truth_params) - for i, star in enumerate(fit_stars): + for i, star in enumerate(stars): print(f'star {i} params = ',star.fit.params) np.testing.assert_allclose(star.fit.params, truth_params, atol=0.0, rtol=1.e-3) @timer -def test_roman_sca_interp(): +def test_sca_interp(): """Test per-SCA/global interpolation behavior and RomanSCAInterp serialization round-trip. """ with fast_pupil_bin(): @@ -690,9 +1153,15 @@ def test_roman_sca_interp(): if __name__ == '__main__': test_roman_optics() - test_roman_corner_cache() - test_roman_fit() - test_roman_aberration_prior() - test_roman_fit_many() - test_roman_optics_convert_funcs() - test_roman_sca_interp() + test_corner_cache() + test_five_point_weights() + test_fit() + test_aberration_prior() + test_fit_many() + test_fit_many_nproc() + test_bilinear_vs_five_point() + test_fit_linear() + test_fit_linear_gradient() + test_fit_linear_nproc() + test_optics_convert_funcs() + test_sca_interp()