JWST-transient-finder

DenseNet169 Ensembles — JWST Transient Detection (F150W, F115W)

A 4-model DenseNet169 ensemble that classifies 64×64 pixel cutouts from JWST NIRCam F150W difference imaging as transient or non-transient. Trained as part of a self-learning pipeline that generates its own labeled data by injecting synthetic point sources (via the instrument PSF) into real F150W difference images. Each NIRCam filter has its own independently trained ensemble, hosted under its own subfolder in this repo.

Available Filters

F115W notes. The F115W ensemble was trained on 5 F115W difference images (see F115W/training_info.json). The first version (v1) flagged far too many subtraction residuals as transients on real data, so it was retrained twice with hard negatives: cutouts that v1/v2 called transient on real F115W difference images but that a human vetted as artifacts (each saved in all 8 rotations/flips). On a held-out F115W difference image (Difference_F115W_61034-59942, 5,000 scored peaks), the predicted-transient count dropped from 1,732 (v1) → 541 (v2) → 305 (v3, released). Only v3 is published here.

Quick Start (inference_pipeline.py)

Single-file, resumable inference pipeline. Point it at a folder of FITS difference files — it reads each file's FILTER header, downloads the matching ensemble from this repo's <FILTER>/ subfolder (strictly that filter — no falling back to another filter's weights), finds candidate peaks, and scores them.

pip install torch torchvision timm astropy numpy pandas scipy pillow tqdm huggingface_hub

fits_files/
    new_diff_file_1.fits
    new_diff_file_2.fits

python inference_pipeline.py

Each run only processes files not yet processed (tracked via a .done marker per file), so it's safe to re-run after dropping new files into fits_files/. Use --force to reprocess everything.

Outputs go to output/<fits_stem>/, one folder per input FITS file:

• predictions.csv — every scored cutout (peak location, sigma, per-model and ensemble mean/std probability, prediction, confidence)
• positives/*.png — cutouts predicted as transient

Downloaded model weights are cached under models/<FILTER>/. Useful flags: --sigma-threshold, --confidence-threshold, --repo-id, --input-dir, --output-dir (--help for the full list).

Legacy multi-script pipeline

test_pipeline.py orchestrates find_peaks_above_k_sigma_test.py + testing_script.py against a local test_directory/ → ML_results/, using local DenseNet169_Ensemble_Model{1-4}_best.pth files instead of fetching from Hugging Face. Kept for reference/offline use; inference_pipeline.py above is the actively maintained entry point.

Model Details

• Architecture: DenseNet169 (growth rate 32, block config (6, 12, 32, 32)), trained from scratch (no pretraining)
• Input: 64×64, 3-channel, ZScale-normalized (astropy ZScaleInterval) cutouts from FITS difference images
• Output: single sigmoid unit — probability of being a real transient
• Ensemble size: 4 independently trained models; predictions are averaged, with the standard deviation across models used as a disagreement/uncertainty signal
• Framework: PyTorch
• Model files: <FILTER>/DenseNet169_Ensemble_Model{1-4}_best.pth (currently F150W/, F115W/)

The pipeline was originally designed to also train a 2-model DeiT (vision transformer) ensemble alongside the DenseNet models. Only the DenseNet169 ensemble was trained for this release (DeiT model count was 0 in the training run), so this card and repo cover the DenseNet ensemble only.

Intended Use

Given a candidate source cutout from a JWST NIRCam difference image (science minus reference epoch), predict whether it is a genuine astrophysical transient (e.g. supernova, variable star) versus a subtraction artifact, cosmic ray, or other spurious peak. Intended as a triage/ranking step ahead of human vetting, not a fully autonomous discovery pipeline.

Out of scope: classifying transient type, non-JWST instruments/pixel scales without re-validation, or images not normalized with ZScale the same way as training.

Training Data

Training data is self-generated from a small set of real JWST NIRCam F150W difference images (no external labeled catalog):

• Positives: synthetic point sources injected at random locations using the instrument PSF, at signal-to-noise ratios spanning SNR 3–10 (psf_injection_script.py).
• Negatives: a mix of quiet/uniform background regions (find_non_peaks_64.py) and real detected peaks that are not injected sources (find_peaks_above_k_sigma_training.py), so the model also learns to reject real (non-transient) astrophysical sources.
• Split: 80% train / 10% validation / 10% test (splitfolders, seed 42) → 14,565 train / 1,820 val / 1,823 test cutouts, roughly balanced between positives and negatives.
• Normalization: ZScale, matching standard astronomical display conventions, scaled to [0, 1].

Difference images used for training were JWST NIRCam F150W filter data. The most recent inference/testing run on new science data was also on F150W difference images; an earlier run additionally evaluated the model on F200W data as a cross-filter check (see Limitations — not a training filter).

Training Procedure

• Loss: binary cross-entropy (BCELoss)
• Optimizer: Adam, lr = 1e-4, weight_decay = 1e-4
• Batch size: 16
• Epochs: up to 10 per model, with best-validation-accuracy checkpointing
• Each of the 4 ensemble members is trained independently (same architecture/data pipeline, different initialization) to provide ensemble diversity/uncertainty estimates
• Hardware: single NVIDIA RTX 4090 Laptop GPU

Evaluation Results

Best validation accuracy per ensemble member (10-epoch training run):

These figures are on the held-out validation split of the self-generated (PSF-injection) dataset, which is easier than real-world discovery — see Limitations.

How to Use Directly (without inference_pipeline.py)

import torch
from pathlib import Path

# see testing_script.py in this repo for the DenseNet definition
# (densenet169()) and the ZScale preprocessing used to turn a FITS
# cutout into a normalized 64x64x3 tensor.

models = []
for path in sorted(Path(".").glob("DenseNet169_Ensemble_Model*_best.pth")):
    model = densenet169(num_classes=1)
    checkpoint = torch.load(path, map_location="cpu")
    model.load_state_dict(checkpoint["model_state_dict"])
    model.eval()
    models.append(model)

with torch.no_grad():
    probs = torch.stack([m(image_tensor) for m in models])
    mean_prob = probs.mean(0)
    std_prob = probs.std(0)

Limitations and Biases

• Positive examples are synthetic PSF injections, not confirmed real transients — the model has not been validated against a catalog of known real transients, and may not generalize to transient morphologies unlike a clean PSF (e.g. blended, saturated, or host-contaminated sources).
• Each ensemble is trained only on difference images in its own filter (F150W or F115W) from a limited set of source fields; the pipeline never applies one filter's weights to another. Performance on other filters, instruments, or pixel scales is unverified.
• The F115W hard negatives come from only two human-labelled difference images, so the F115W model may still flag subtraction residuals of kinds not represented there (e.g. bright-galaxy-core residuals in cluster fields).
• The near-100% validation accuracy reflects the relative ease of the self-supervised task (clean injected point sources vs. background/real peaks) and should not be read as real-world discovery precision/recall. Always inspect flagged detections visually before treating them as candidate transients.
• No independent test against human-vetted or spectroscopically confirmed transients has been performed.

Citation / Provenance

Produced by the self-learning-pipeline-transient-detection project. See the parent repository's training pipeline (train_pipeline.py, training_script.py) for full training code and data-generation details.