#!/usr/bin/env python """Run Schlafly crowdsource on the level-3 NRCB mosaic for one filter. Writes the per-filter catalog to benchmark/crowdsource/catalogs/_crowdsource.fits. """ from __future__ import annotations import argparse import os import sys import time from pathlib import Path import numpy as np from astropy.io import fits from astropy.table import Table from astropy.wcs import WCS sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "shared")) import config # noqa: E402 import crowdsource # noqa: E402 from crowdsource import crowdsource_base # noqa: E402 # stpsf grid PSFs pre-computed in sickle root: nircam_nrcbN__*.fits import stpsf # noqa: E402 def load_gridded_psf(filt: str, detector: str = "nrcb5"): """Build a photutils GriddedPSFModel using stpsf, or load cached copy. Uses NIRCam short-wavelength detectors (nrcb1-4) for F187N/F210M and nrcb5 for LW. For the mosaic (multiple detectors resampled), we pick the representative one. """ # Use the pre-made grid PSFs sitting at /orange/adamginsburg/jwst/sickle/ candidates = list((config.SICKLE_ROOT).glob(f"nircam_{detector}_{filt.lower()}_*.fits")) if not candidates: raise FileNotFoundError( f"no gridded PSF found for {filt} {detector} under {config.SICKLE_ROOT}" ) grid_path = candidates[0] print(f"Loading PSF grid: {grid_path}") from stpsf.gridded_library import display_psf_grid # noqa from photutils.psf import GriddedPSFModel from astropy.nddata import NDData hdul = fits.open(grid_path) data = hdul[0].data hdr = hdul[0].header grid_xypos = [] for i in range(hdr.get("NPSF", data.shape[0])): x = hdr[f"DET_YX{i}"] if f"DET_YX{i}" in hdr else None if x is None: break # DET_YX is stored as "(y, x)" string y_, x_ = eval(hdr[f"DET_YX{i}"]) grid_xypos.append((x_, y_)) if len(grid_xypos) != data.shape[0]: raise ValueError("DET_YX keywords missing in PSF grid header") ndd = NDData(data, meta={ "oversampling": hdr.get("OVERSAMP", 4), "grid_xypos": grid_xypos, }) return GriddedPSFModel(ndd) class WrappedPSFModel(crowdsource.psf.SimplePSF): """photutils GriddedPSFModel wrapped for crowdsource's SimplePSF interface. Mirrors the wrapper in brick2221/analysis/crowdsource_catalogs_long.py. """ def __init__(self, psfgridmodel, stampsz=19): self.psfgridmodel = psfgridmodel self.default_stampsz = stampsz def __call__(self, col, row, stampsz=None, deriv=False): if stampsz is None: stampsz = self.default_stampsz parshape = np.broadcast(col, row).shape tparshape = parshape if len(parshape) > 0 else (1,) rows, cols = np.indices((stampsz, stampsz)) - (np.array([stampsz, stampsz]) - 1)[:, None, None] / 2.0 col = np.atleast_1d(col) row = np.atleast_1d(row) stamps = [] for i in range(len(col)): stamps.append(self.psfgridmodel.evaluate(cols + col[i], rows + row[i], 1, col[i], row[i])) stamps = np.array(stamps) stamps /= stamps.sum(axis=(1, 2))[:, None, None] if deriv: dpsfdrow, dpsfdcol = np.gradient(stamps, axis=(1, 2)) ret = stamps if parshape != tparshape: ret = ret.reshape(stampsz, stampsz) if deriv: dpsfdrow = dpsfdrow.reshape(stampsz, stampsz) dpsfdcol = dpsfdcol.reshape(stampsz, stampsz) if deriv: ret = (ret, dpsfdcol, dpsfdrow) return ret def render_model(self, col, row, stampsz=None): if stampsz is not None: self.stampsz = stampsz rows, cols = np.indices(self.stampsz, dtype=float) - (np.array(self.stampsz) - 1)[:, None, None] / 2.0 return self.psfgridmodel.evaluate(cols, rows, 1, col, row).T.squeeze() def run_filter(filt: str, nskyx: int = 3, nskyy: int = 3): mosaic_path = config.MOSAICS[filt] if not mosaic_path.exists(): raise FileNotFoundError(mosaic_path) print(f"[crowdsource] Running on {mosaic_path}") hdul = fits.open(mosaic_path) sci = hdul["SCI"].data.astype(np.float32) # Build weights from ERR extension if present, otherwise 1/std if "ERR" in hdul: err = hdul["ERR"].data.astype(np.float32) w = np.zeros_like(err) good = err > 0 w[good] = 1.0 / err[good] else: w = np.ones_like(sci) # Mask invalid pixels bad = ~np.isfinite(sci) | (sci == 0) | (w == 0) w[bad] = 0.0 sci_clean = np.nan_to_num(sci, nan=0.0, posinf=0.0, neginf=0.0) wcs = WCS(hdul["SCI"].header) # NRCB5 (LW) has its own PSF; SW uses nrcb1 representative detector = "nrcb5" if filt in config.LW_FILTERS else "nrcb1" grid = load_gridded_psf(filt, detector=detector) psf = WrappedPSFModel(grid, stampsz=19) t0 = time.time() results = crowdsource_base.fit_im( sci_clean, psf, weights=w, nskyx=nskyx, nskyy=nskyy, refit_psf=False, verbose=True, threshold=5.0, ) dt = time.time() - t0 # Current crowdsource returns a dict: stars, model[B], sky[B], psfs, iter_history. cat = results["stars"] modelim = results["model"][0] skyim = results["sky"][0] print(f"[crowdsource] Done in {dt:.1f}s. Detected {len(cat)} sources.") # Attach world coordinates tab = Table(cat) ra, dec = wcs.all_pix2world(tab["x"], tab["y"], 0) tab["ra"] = ra tab["dec"] = dec tab.meta["FILTER"] = filt tab.meta["TOOL"] = "crowdsource" tab.meta["MOSAIC"] = str(mosaic_path) tab.meta["RUNTIME_S"] = dt outdir = config.TOOL_DIRS["crowdsource"] / "catalogs" outdir.mkdir(parents=True, exist_ok=True) outpath = outdir / f"{filt.lower()}_crowdsource_nrcb.fits" tab.write(outpath, overwrite=True) print(f"[crowdsource] Wrote {outpath}") # Save residual for diagnostic resid_path = outdir / f"{filt.lower()}_crowdsource_nrcb_residual.fits" fits.writeto(resid_path, sci - modelim - skyim, header=hdul["SCI"].header, overwrite=True) print(f"[crowdsource] Wrote residual {resid_path}") def main(): ap = argparse.ArgumentParser() ap.add_argument("filter", choices=config.FILTERS) # nskyx/nskyy = 0 disables the sky-parameter columns. We need this # because crowdsource_base.build_sparse_matrix has a shape bug at # line 501 when nskypar > 0 (see benchmark/README.md). JWST i2d # mosaics are already sky-subtracted so this is a valid workaround. ap.add_argument("--nskyx", type=int, default=0) ap.add_argument("--nskyy", type=int, default=0) args = ap.parse_args() run_filter(args.filter, nskyx=args.nskyx, nskyy=args.nskyy) if __name__ == "__main__": main()