diff --git a/.github/workflows/pr_quality.yml b/.github/workflows/pr_quality.yml
index a6d47c5..2342494 100644
--- a/.github/workflows/pr_quality.yml
+++ b/.github/workflows/pr_quality.yml
@@ -149,6 +149,7 @@ jobs:
- name: Launcher contracts
run: |
pytest -q \
+ tests/test_launch_aurora.py \
tests/test_launch_aurora_web_vla.py \
tests/platform/test_perlmutter_launch_contract.py \
tests/platform/test_hpc_bridge_contract.py \
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0e5b497..3b79aef 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -30,6 +30,10 @@ will make up v0.1.0.
per-modality support-status table
- `tools/ci/check_doc_links.py`, a link checker for relative Markdown links, run
as a required `doc_links` CI job
+- `docs/tutorials/`, an index of end-to-end tutorials grouped by application
+ area. The materials tutorial moved from `docs/materials_prism_tutorial.md` to
+ `docs/tutorials/material_science/materials_prism_tutorial.md`; notebooks stay
+ next to their scripts in `applications//`
### Changed
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 6fcce00..312dabb 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -274,6 +274,7 @@ descriptions and return types where they are not obvious from the signature.
| `modalities/` | Encoders, projectors, output decoders |
| `models/` | Backbone integrations |
| `evaluation/` | Universal evaluator, vLLM inference |
+ | `tutorials/` | End-to-end walkthroughs, one folder per application area; notebooks stay in the top-level `applications//`, not `docs/applications/` |
| `applications/` | Scientific applications (BioReason, KEGG, materials) |
| `api/` | API reference for the public `src.*` surface |
| `development/` | CI and contributor tooling |
diff --git a/Makefile b/Makefile
index 1aad646..642267f 100644
--- a/Makefile
+++ b/Makefile
@@ -89,6 +89,7 @@ multimodal: ## Run the contract suite as CI's `multimodal_contract` job does
# none of them -- this target is the only make path that does.
launcher: ## Run the launcher contracts as CI's `launcher_contract` job does
$(PYTHON) -m pytest -q \
+ tests/test_launch_aurora.py \
tests/test_launch_aurora_web_vla.py \
tests/platform/test_perlmutter_launch_contract.py \
tests/platform/test_hpc_bridge_contract.py \
diff --git a/README.md b/README.md
index 8d839a4..bddd816 100644
--- a/README.md
+++ b/README.md
@@ -176,6 +176,7 @@ Per-platform detail, including lockfile regeneration:
| **Text** | [OLMo-1B / OLMo-3 7B](https://huggingface.co/allenai/OLMo-7B-0724-hf) | 2048 / 4096 |
| **Image** | [SigLIP2](https://huggingface.co/google/siglip2-base-patch16-224) | 768 |
| **Time series** | [Moirai 2.0](https://huggingface.co/Salesforce/moirai-2.0-R-small) / [Intern-S2 preview](https://huggingface.co/internlm/Intern-S2-Preview-397B) / [TimeOmni](https://arxiv.org/abs/2510.03255) | Variable by backend |
+| **Spectrum** | [AION-1 Base](https://huggingface.co/polymathic-ai/aion-base) | 768 |
| **DNA** | [Nucleotide Transformer v2 500M](https://huggingface.co/InstaDeepAI/nucleotide-transformer-v2-500m-multi-species) | 1024 |
| **Geometry** | [Walrus](https://huggingface.co/polymathic-ai/walrus) | 512 |
| **Graph** | GraphMAE2 (PyG) | 768 |
@@ -193,10 +194,11 @@ actually exercised by shipped configs and CI, not what the architecture permits.
| Modality | Status | Shipped model configs | Notes |
|---|---|---|---|
-| **Text** | Supported | 26 of 26 | Every config enables it. With an HF backbone, the backbone's own embeddings are used and `TextEncoder` is bypassed. |
-| **Image** | Supported | 9 of 26 | Best-validated multimodal path; the basis of the scaling results below. |
-| **Time series** | Supported | 14 of 26 | Five encoder backends (`linear`, `moirai`, `intern_s2`, `intern_s2_397b`, `timeomni`). One of two modalities with a vLLM inference processor (the other is image). |
-| **DNA** | Supported | 2 of 26 | Nucleotide Transformer encoder (`src/encoders/dna.py`), used by the BioReason/KEGG path — see [docs/applications/bioreason_grpo.md](docs/applications/bioreason_grpo.md). Input-only. |
+| **Text** | Supported | 29 of 29 | Every config enables it. With an HF backbone, the backbone's own embeddings are used and `TextEncoder` is bypassed. |
+| **Image** | Supported | 11 of 29 | Best-validated multimodal path; the basis of the scaling results below. |
+| **Time series** | Supported | 14 of 29 | Five encoder backends (`linear`, `moirai`, `intern_s2`, `intern_s2_397b`, `timeomni`). One of two modalities with a vLLM inference processor (the other is image). |
+| **DNA** | Supported | 2 of 29 | Nucleotide Transformer encoder (`src/encoders/dna.py`), used by the BioReason/KEGG path — see [docs/applications/bioreason_grpo.md](docs/applications/bioreason_grpo.md). Input-only. |
+| **Spectrum** | Supported | 2 of 29 | AION-1 Base codec + shared encoder (`src/encoders/aion_spectrum.py`), used by the AstroPRISM path — see [docs/tutorials/astroprism/astroprism_tutorial.md](docs/tutorials/astroprism/astroprism_tutorial.md). Input-only. |
| **Graph** | Experimental | 0 | Requires `torch_geometric`. Runs in the per-modality sweep (172.7 samp/s) but no shipped model config enables it, and the handler is labelled untested in its own smoke config. |
| **Table** | Experimental | 0 | TAPAS-based; runs in the sweep (65.5 samp/s). Input-only — no output decoder. No shipped model config enables it. |
| **Geometry** | Experimental | 0 | Requires `walrus` and `the_well`. Runs in the sweep (69.4 samp/s) **on synthetic data** — no real PDEBench corpus is staged, and `GeometryEncoder` hardcodes `input_dim=6`. A `WALRUS_FALLBACK=1` stub exists for smoke runs and produces meaningless features. |
@@ -375,6 +377,7 @@ PRISM/
[docs/index.md](docs/index.md) indexes every page. Most-used entry points:
- [docs/getting-started.md](docs/getting-started.md) — clone → environment → first run
+- [docs/tutorials/README.md](docs/tutorials/README.md) — end-to-end tutorials by application area (BioReason, materials, …)
- [docs/platforms/site_paths.md](docs/platforms/site_paths.md) — configuring the `PRISM_*` roots
- [docs/platforms/aurora_operations.md](docs/platforms/aurora_operations.md) — Aurora ops, env, CCL tuning, troubleshooting
- [docs/training/training.md](docs/training/training.md) — training stages
diff --git a/docs/applications/astroprism.md b/docs/applications/astroprism.md
new file mode 100644
index 0000000..029a40a
--- /dev/null
+++ b/docs/applications/astroprism.md
@@ -0,0 +1,27 @@
+# AstroPrism: the 646-candidate pilot
+
+Start with the **646 positional matches between released paired AION galaxies and AION-Search image captions**. Recover and validate these objects on Aurora, train image-to-language caption alignment, prepare grounded spectral captions in parallel, then extend to image + spectrum + language, instruction SFT, and a later verifier-based RL experiment.
+
+The earlier goal of **10,000 training galaxies plus separate validation/test galaxies** remains a later expansion after this small pipeline works.
+
+## Contents
+
+| File | Purpose |
+|---|---|
+| [astroprism_training_plan.md](astroprism_training_plan.md) | AION-1/Qwen architecture, phases, splits, training settings and implementation work |
+| [astroprism_synthetic_data_prompts.md](astroprism_synthetic_data_prompts.md) | Image verification, spectral facts/captions, joint captions, instruction conversations and validation prompts |
+| [astroprism_benchmark_landscape.md](astroprism_benchmark_landscape.md) | AstroSpecQA, image benchmarks, paired/interleaved evaluation and RL evaluation boundaries |
+| [astroprism_data.md](astroprism_data.md) | Dataset pins, roster checksum, field definitions and remote recovery procedure |
+| [astroprism_data/overlap_candidates_646.csv.gz](astroprism_data/overlap_candidates_646.csv.gz) | Exact 646-row candidate roster with IDs, coordinates, separations and source shard/row locators |
+
+This page and the four it links to are the whole package. It is a **plan**: no AION inference, no astronomy-specific PRISM integration, no caption synthesis, no SFT and no RL is implemented here. The analysis scripts, audit tables and diagnostic arrays behind the roster are not part of it.
+
+## What is ready
+
+The positional-match roster is committed here. Its 646 paired catalog IDs all map uniquely into the official tutorial raw-file index, with matching DESI IDs for these candidates. Physical association, raw observation quality and compatibility of each caption with AION's exact image crop still need checking. Caption text has not yet been extracted for this roster.
+
+The proposal uses AION-1 Base's modality-specific codecs and shared frozen encoder, separate image/spectrum projectors, and Qwen3-4B. Nothing here establishes model readiness; the numbers below are proposals, not measurements.
+
+For a 646-object accepted pool, the proposed split is approximately **516 train / 65 validation / 65 test**, grouped by physical source before training or instruction generation. Rejections or grouping can change those counts. The 646 rows are the complete candidate pool, not 646 training objects plus additional holdouts.
+
+Remote data recovery, model loading, job submission and training are all subsequent work. [astroprism_data.md](astroprism_data.md) documents how to recover the captions and raw observations from the pinned public sources.
diff --git a/docs/applications/astroprism_benchmark_landscape.md b/docs/applications/astroprism_benchmark_landscape.md
new file mode 100644
index 0000000..0de9797
--- /dev/null
+++ b/docs/applications/astroprism_benchmark_landscape.md
@@ -0,0 +1,141 @@
+# Astronomy benchmark landscape for the PRISM galaxy pilot
+
+Updated: 2026-09-16. This is an evaluation plan, not an executed benchmark. Start from the [646 positional candidates](astroprism_data/overlap_candidates_646.csv.gz), confirm eligible physical pairs using the [data contract](astroprism_data.md), and reserve source-disjoint validation/test sets before any training. A provisional split is **516 train / 65 validation / 65 test**, conditional on all 646 being eligible and group constraints permitting it. The [training plan](astroprism_training_plan.md) covers image-caption alignment, spectrum/joint alignment, instruction SFT and later RL; [generation prompts](astroprism_synthetic_data_prompts.md) specify evidence permissions. The 10,000-training-galaxy design is a later expansion.
+
+## What to use
+
+Use three complementary evaluations: **scientific prediction from real observations**, **answers grounded in the available observation**, and **modality-dependence controls**. Existing resources cover parts of this combination. I did not locate a ready-to-run public benchmark that combines matched galaxy images, numerical 1D spectra, and independently validated open-ended joint answers on the same sources. This is a bounded search finding, not a claim that no such resource can exist.
+
+For this galaxy-only pilot, prioritize held-out redshift and spectral-feature tasks, human-supported morphology labels, and a small expert-reviewed joint QA set. Use paper-figure QA and text knowledge benchmarks as separate optional diagnostics. Do not report all-galaxy object classification or stellar subtype/temperature tasks as meaningful core pilot metrics.
+
+To stay close to the requested AION examples, its image evaluation already crossmatches Galaxy10 with Legacy DR10, yielding approximately 8,000 galaxies with four-band cutouts. This is a useful morphology protocol to adapt, but it does not establish how many of our paired pilot galaxies carry those labels. [AION §7.2.2](https://arxiv.org/html/2510.17960)
+
+## AstroSpecQA: the exact published reference
+
+AstroSpecQA uses numerical DESI spectra, held apart from DESI-SpecInstruct. Its **11,698 questions** comprise:
+
+| Task | Questions | Published metric |
+|---|---:|---|
+| Object class | 5,000 | Accuracy |
+| Redshift, galaxies/QSOs | 3,247 | MAE, R² |
+| Stellar subtype | 1,753 | Accuracy |
+| Stellar effective temperature | 1,698 | MAE, R² |
+
+Appendix A.4 gives object counts: STAR 1,753; GALAXY 1,648; QSO 1,599. Stellar subtypes: K 620; G 386; M 322; F 274; WD 94; A 53; B 3; CV 1. Thus subtype accuracy can hide rare-class failure. Questions must not be counted as distinct physical objects. These details come from [§3.4, §5.1 and Appendix A.4](https://arxiv.org/html/2609.07102).
+
+The paper scores explanations with a fact-conditioned judge and a 100-output human check. The original benchmark files have not been obtained in this project, so its task design is a reference rather than an executable reproduction. [Explanation evaluation](https://arxiv.org/html/2609.07102)
+
+## Existing resources and their actual input modalities
+
+| Resource | What the model receives / does | Verified scope and availability | Role for PRISM |
+|---|---|---|---|
+| **Galaxy Zoo / Galaxy10** | Real galaxy images → morphology labels/vote fractions | Public catalogues/images. Galaxy Zoo DECaLS includes volunteer classifications and machine predictions; Galaxy Zoo DESI's 8.67M catalogue is automated morphology inference. [Official data](https://data.galaxyzoo.org/), [MMU Galaxy10](https://huggingface.co/datasets/UniverseTBD/mmu_gz10) | Strong image-grounding evaluation when the gold labels are human-supported. A machine prediction is a weaker reference, not independent human truth. |
+| **AION galaxy-property evaluation** | Calibrated photometry, images and numerical spectra → physical-property estimates | Published PROVABGS evaluation, R², with photometry / photometry+image / photometry+image+spectrum conditions. [Paper §7.2.1](https://arxiv.org/html/2510.17960), [code](https://github.com/PolymathicAI/AION) | Closest existing multimodal scientific-prediction protocol. Keep input conditions explicit; the paper's strongest joint condition includes photometry. It is not an open-ended QA benchmark. |
+| **AstroCLIP** | Real Legacy images and numerical DESI spectra → embeddings, cross-modal retrieval and downstream predictions | Public model/data workflow. [Official repository](https://github.com/PolymathicAI/AstroCLIP) | Encoder/probe and retrieval baselines. Retrieval performance alone does not establish grounded conversation. |
+| **AION-Search** | Galaxy image / text query → retrieval ranking | Evaluates spirals, mergers and lenses using nDCG@10; also reports Galaxy10 classification. [Paper §§II.4–II.6](https://arxiv.org/html/2512.11982v2) | Optional retrieval diagnostic. Its synthetic captions are training supervision, not independent evaluation answers. |
+| **AstroMMBench** | A rasterized astronomy-paper figure plus a four-option question | 621 MCQs across six subfields; generated from arXiv figure/context pairs and reviewed by experts. Accuracy overall/by subfield. [Paper](https://arxiv.org/html/2510.00063), [public release](https://huggingface.co/datasets/sjh11/AstroMMBench) | Optional figure-understanding/knowledge diagnostic. A paper figure, including a plotted spectrum, is not a numerical 1D DESI input or necessarily a raw sky cutout. |
+| **AstroVLBench / ASIB** | Separate optical-image and DESI-spectrum tasks; spectrum plots and numerical CSVs are distributed | Public anonymous release; ASIB provides line-presence, broad-line and BPT tasks. [Dataset and task definitions](https://huggingface.co/datasets/anonymous4ai/AstroVLBench) | Possible later external check. Inspect the original inference path, spectrum frame and labels first; a native 1D adaptation is a new protocol. Its separate task populations do not establish matched image+spectrum joint QA. |
+
+Galaxy Zoo morphology is the closest ready-made image-only scientific check for this pilot. Match human labels to the exact student view; machine-generated catalogue predictions are weaker references. AstroMMBench is useful for figure QA, but calibrated four-band AION image inputs are not automatically interchangeable with arbitrary RGB paper figures. Neither resource supplies the native paired-encoder evaluation we need.
+
+## Proposed evaluation for the 646-candidate pilot
+
+Everything below is a **new PRISM evaluation proposal**, not an existing published benchmark or a reported result.
+
+### Source and split control
+
+Use the proposed 65 validation galaxies for model selection, inference-prompt/scorer development and threshold selection. Training-data teacher prompts and their demonstrations use training sources only. Lock the proposed 65 test galaxies before new caption generation or instruction synthesis. These are targets, not an existing split: resolve candidate matches first and record the actual numbers after quality cuts and grouping. Group all repeat observations of each physical object; prevent near-duplicate/overlapping image leakage, and report sky/field and observing-condition slices. Exclude test identities from every training stage, the RL reward pool, all augmentation descendants, and demonstrations supplied in prompts.
+
+Document possible AION pretraining exposure separately. A downstream split does not prove that the pretrained encoder never saw a source or an associated catalog label. Record object identity, survey/reduction, raw-data references, processed student views, input metadata, question, target provenance, evidence locations, scoring version, and allowed modalities in a frozen test manifest.
+
+### Core task panel
+
+| Panel | Proposed population | Inputs / questions | Scoring |
+|---|---|---|---|
+| **Galaxy redshift** | Up to 65 eligible test galaxies; report quality exclusions explicitly | Same numerical prediction task under image-only, spectrum-only and joint inputs; no target redshift supplied | MAE and R² for comparability; additionally bias, NMAD and outlier rate for normalized residual Δz/(1+z). Fix the outlier threshold on validation; a possible pilot choice is 0.05, explicitly a design choice. Report parse failures and numeric-answer coverage. |
+| **Image morphology** | Human-labeled crossmatch subset `N_morph ≤ 65`, or a newly expert-labeled subset | Visible structure, orientation and feature presence. Use the exact crop/bands seen by the student | Macro-F1/balanced accuracy for declared discrete classes; proper probability scores for retained vote distributions. Keep machine-only labels in a separate analysis. Do not invent a guaranteed crossmatch count. |
+| **Spectral evidence** | Valid, covered line/continuum items across the test sources; denominator varies by feature | Presence/strength of measured features, line ratios when valid, quality/coverage limitations | Precision/recall and macro-F1 on feature states; numeric errors with units/tolerances; separately score upper limits, missing coverage and masked regions. Compare against versioned measurements, not teacher prose. |
+| **Joint answers** | All eligible test galaxies, up to 65, with expert review | One expert-reviewed joint question per source, evaluated with joint, image-only and spectrum-only views | Claim correctness, evidence attribution, unsupported-claim rate, and correct partial answers/abstention. Each view needs an admissible reference answer. Two independent reviewers adjudicate disagreements; automated judging is supplementary. |
+| **Observation interventions** | A declared subset of the same test sources | Remove relevant lines or image regions, remove a modality, or swap one input with controlled alternatives | Direction/consistency of scientifically justified answer changes. Some swaps are observationally indistinguishable; do not force a mismatch verdict. |
+
+With 65 eligible test galaxies, the redshift panel yields up to **195 predictions on 65 galaxies**; evaluating one joint question under three input conditions yields up to **195 answers on those same galaxies**. These are evaluation records, not additional independent objects. With this small test set, report uncertainty and treat results as feasibility evidence rather than a general astronomy benchmark. Count spectrum-feature items and morphology labels separately. Bootstrap confidence intervals by galaxy, not by duplicated question, and report class, brightness, signal-to-noise and redshift slices.
+
+For the proposed redshift scorer, define `r = (z_pred − z_ref)/(1 + z_ref)`, bias as `median(r)`, and `NMAD = 1.4826 × median(|r − median(r)|)`. Report the fixed-threshold fraction `|r| > 0.05` separately from parse failures; report answer coverage on the complete eligible roster. These explicit scorer conventions are our proposal, not a claim that AstroSpecQA used them.
+
+For explanation scoring, retain concise claim-level evidence references in private audit records. The scoring facts should come from independently checked measurements/labels and expert review, rather than reuse of the same teacher's generated answer. A fluent explanation can be wrong even when a final label happens to be correct.
+
+### Controls and baselines
+
+- Compare direct instruction tuning against caption alignment followed by instruction tuning, with matched training data and evaluation sources.
+- Evaluate image-only, spectrum-only, joint and text-only conditions. Include question-only or shuffled-input controls to expose answer priors.
+- Train simple frozen-AION probes on the same training roster for numerical/categorical tasks. Match available inputs: an AION photometry+image+spectrum result is not an input-matched baseline for image+spectrum-only PRISM.
+- Keep observational captions separate from supervised property estimates. Catalog masses, SFRs or redshifts may be valid prediction labels; their presence in a reference table does not make them directly visible evidence in every modality.
+- In independent redshift evaluation, retain observed-frame spectra and exclude true-redshift rest-frame preprocessing and target-derived student features. Ensure a later conversation turn does not receive the target through an earlier answer.
+- Use only information retained after AION preprocessing. A teacher's wider image, additional band, or uncropped wavelength range cannot justify a claim about evidence the student does not receive.
+- Score final answers and evidence separately. Fix parsing rules, numeric conventions and invalid-answer handling before the test run. Never hide missing/malformed predictions by reporting accuracy only on successful parses.
+
+## Interleaving: input formatting versus scientific capability
+
+**Interleaving is a way to arrange the inputs, not a scientific task by itself.** A model may receive `text → image → text → spectrum → question` yet answer using only one modality. Conversely, a model receiving an image prefix followed by a spectrum prefix can perform a genuinely joint task. Demonstrate joint use through the scientific panels and observation interventions above, independently of serialization format.
+
+The initial proposal uses **all declared image/spectrum prefixes before the text conversation** for one galaxy. This is the baseline to implement and verify first. Incremental interleaved observations, repeated modality spans, and multiple-object conversations are optional future work; this note does not establish that the current PRISM loader, attention masks, or generation path support them.
+
+### Three distinct tests
+
+| Test | Example / contract | What it establishes |
+|---|---|---|
+| **Same object, different ordering** | Present the same galaxy image, spectrum, metadata and question before the answer, using image-first, spectrum-first, or text-separated layouts | Robustness to presentation order when the available evidence is unchanged. Hold the facts and permitted answer fixed; do not count permutations as new objects. |
+| **New evidence over multiple turns** | Turn 1 provides an image and asks about visible morphology. Turn 2 introduces that galaxy's spectrum and asks for spectral interpretation or a combined assessment | Whether the model incorporates newly available evidence without using it prematurely. The first answer must have an image-only reference, the later answer a reference valid for the expanded input. |
+| **Multiple objects and source binding** | Local alias A has image A and spectrum A; alias B has image B and spectrum B; the question explicitly asks about A, B or their comparison | Whether findings attach to the correct source. Reorder whole source-bound spans and include controlled distractors. Real survey IDs remain in the audit manifest; neutral A/B aliases may appear in the conversation. |
+
+These examples describe **proposed test structures**, not published example records or executed model results. Keep text references unambiguous: every “this galaxy,” “its spectrum,” or “the second source” must resolve to a source and an observation already available at that answer position.
+
+### Causal availability is part of the gold answer
+
+For each assistant turn, record the set of available observation IDs. Its question and reference answer must use that set, not modalities supplied later in the conversation. An image-only turn cannot claim line detections from a spectrum introduced afterward. Teacher generation and judging must follow the same turn-specific restriction.
+
+A prefix-only implementation that places **all** observations at the beginning would leak future spectra into earlier answers in an incremental-observation test. Do not run or claim this test until the model and collator implement the intended positions and causal visibility. For the initial prefix baseline, supply all observations up front and evaluate only that declared condition.
+
+Order permutations are scientifically comparable only when evidence availability and referents remain equivalent. Introducing a spectrum after an earlier answer changes the information available and therefore requires a different reference answer; it is not merely a harmless formatting permutation.
+
+### Engineering checks before scientific benchmarking
+
+Validate the serialization contract with small deterministic fixtures before model-quality comparisons:
+
+- Each modality placeholder maps to exactly one intended tensor/embedding span, with matching sample, source alias, modality type, ordering, token length and positions.
+- Repeated images/spectra and multiple objects retain correct bindings through batching, collation, caching and generation; no span or cached feature crosses sample boundaries.
+- Padding, truncation and missing modalities preserve valid span boundaries and attention masks; truncation must not silently retain a question whose required evidence was removed.
+- Assistant-only loss labels exclude user text, modality prefixes, padding and audit sidecars. For incremental inputs, verify that earlier predictions cannot attend to later modality spans, including during cached generation.
+- Dataset pixel-validity/bad-pixel masks, modality-presence masks, and language-model attention masks have separate named contracts and must not be conflated.
+
+These are **placeholder/alignment and causal-visibility checks**, not evidence of astronomy competence. After they pass, evaluate scientific correctness, source attribution, partial-answer behavior, order sensitivity and observation interventions on the held-out sources. Report both engineering pass/fail results and scientific scores; success in either category alone does not prove the other.
+
+### Evaluation after each training stage
+
+| Saved checkpoint | Required check | What a positive result establishes |
+|---|---|---|
+| Image caption alignment | Image claims against independently checked morphology/quality facts; compare image-conditioned and text-only output | Whether the image projector supplies useful visual evidence on unseen sources |
+| Spectrum and joint caption alignment | Spectral feature correctness, uncertainty, image retention, and per-claim modality attribution | Whether adding the spectrum branch supports its own facts without erasing image grounding |
+| Instruction SFT | All applicable task panels, answerability and multi-turn consistency; matched direct-SFT baseline if feasible | Whether the model follows instructions and keeps supported scientific content |
+| Later RL pilot | The identical locked panels plus reward edge-case audit and comparison with the saved SFT checkpoint | Whether optimization of the independent verifier improves held-out scientific results, rather than only training reward |
+
+Save all intermediate checkpoints and define the comparisons before seeing test results. Use validation for early stopping and prompt changes. Score predeclared checkpoints on the test set in one frozen evaluation; repeated test inspection followed by edits turns it into development data. If exploratory tests are unavoidable, disclose them and obtain a fresh independent test pool for final claims.
+
+### RL reward and evaluation are separate
+
+The proposed RL pool contains training sources only and tasks with independent numerical/categorical checkers. Fix units, tolerances, quality eligibility and parse rules on training/validation data. Do not reward teacher agreement, answer length, eloquence or a model's self-assessment. A reward for a correct scalar does not verify the accompanying scientific explanation. Test the checker using valid, invalid, boundary and malformed outputs before policy optimization.
+
+Measure held-out correctness, unsupported claims, appropriate abstention and retained image/spectrum capabilities separately from reward. Keep a non-RL SFT checkpoint as the baseline. A gain on a reward built from the same synthetic teacher that produced training answers is not independent evidence; use catalog/measurement targets and expert-reviewed claims as applicable. Given the small pool, proceed beyond a bounded RL pilot only if the verifier is sound and the independent results justify it.
+
+### Optional external checks
+
+1. **Galaxy10 / human Galaxy Zoo labels:** morphology test after identity-overlap checks and matching the image view. A re-split or re-crop is a new evaluation condition, not an exact reproduction of a published score.
+2. **AstroVLBench ASIB:** inspect and reproduce its original input protocol first, then separately identify any numerical-spectrum adaptation. Verify survey release, preprocessing, wavelength frame and target definitions; do not assume DESI releases are interchangeable.
+3. **AstroMMBench:** only through an interface that handles its paper figures. Report figure QA separately from raw-observation performance.
+4. **AION / AstroCLIP probes and retrieval:** use the same held-out source roster and declared inputs. Retrieval and property estimation do not by themselves establish instruction following.
+
+## Acceptance before implementing benchmark code
+
+Before scoring, assemble a reviewable evaluation manifest: exact source roster, overlap report, task eligibility counts, independent targets, questions/reference answers, per-view permissions, frozen scorers, and baseline configurations. Complete a small expert audit spanning ambiguous morphology, low signal-to-noise, nondetections, out-of-coverage features and aperture mismatch. Then freeze the benchmark and run the proposed comparisons.
+
+Document the resulting native image+spectrum QA panel as a **new AstroPrism pilot evaluation**, with its small sample size, construction and limitations. It should not be called AstroSpecQA reproduction unless the original benchmark data and evaluation protocol are obtained and preserved.
diff --git a/docs/applications/astroprism_data.md b/docs/applications/astroprism_data.md
new file mode 100644
index 0000000..fc99813
--- /dev/null
+++ b/docs/applications/astroprism_data.md
@@ -0,0 +1,117 @@
+# The 646 AION image–caption overlap candidates
+
+`overlap_candidates_646.csv.gz` is the complete, byte-identical candidate roster from the metadata audit completed on 2026-09-16. It contains identifiers, sky positions, angular separations, and exact source shard/row locators. **It does not contain caption text, image or spectrum arrays, features, or encoder weights.** Recover these on Aurora before accepting examples for training.
+
+- Compressed size: **39,219 bytes**; CSV data rows: **646**.
+- SHA-256: `650935faa58953c9929e7379429bfb332a7486dab54544fedef64ca385abb3bc`.
+- All 646 paired IDs and all 646 caption IDs are distinct. Every row has exactly one candidate in each direction within **1 arcsecond**; all three candidate/uniqueness flags are `True`.
+- Median separation: **0.0191149 arcsec**; maximum: **0.2309893 arcsec**. All distances were independently checked with a spherical haversine calculation.
+- The audit covered all **95,895** paired-feature objects and **164,911** Legacy DR10 caption records; it excluded the **110,702** HSC caption records. This is a positional candidate roster, not 646 already accepted training pairs.
+
+## License and attribution
+
+The roster is a derived table — identifiers, coordinates, separations and shard/row
+locators — computed from the three pinned sources below. It carries no caption text, pixels,
+spectra, features or weights.
+
+| Source | License |
+|---|---|
+| [`EiffL/desi_x_ls_aion`](https://huggingface.co/datasets/EiffL/desi_x_ls_aion) | CC BY 4.0 |
+| [`astronolan/galaxy-descriptions`](https://huggingface.co/datasets/astronolan/galaxy-descriptions) | MIT |
+| [AION tutorial HDF5](https://users.flatironinstitute.org/~polymathic/data/provabgs_desi_ls.hdf5) (Polymathic AI / Flatiron Institute) | see upstream |
+
+Both Hugging Face datasets permit redistribution of derived work with attribution, which this
+section and the pinned revisions below provide. The underlying observations come from DESI and
+the DESI Legacy Imaging Surveys; cite them directly in any downstream publication.
+
+## Pinned sources and identifier meanings
+
+| Source | Repository / location | Exact revision |
+|---|---|---|
+| Paired AION features | [`EiffL/desi_x_ls_aion`](https://huggingface.co/datasets/EiffL/desi_x_ls_aion/tree/57f7c7575260811ea076163bc743787d75b9aa49) | `57f7c7575260811ea076163bc743787d75b9aa49` |
+| AION-Search image descriptions | [`astronolan/galaxy-descriptions`](https://huggingface.co/datasets/astronolan/galaxy-descriptions/tree/6890dc0c8fc793867c07a6ce4fce11e51c167d6e) | `6890dc0c8fc793867c07a6ce4fce11e51c167d6e` |
+| Raw paired observations | [Official AION tutorial HDF5](https://users.flatironinstitute.org/~polymathic/data/provabgs_desi_ls.hdf5) | HTTP ETag `"6a62487d-1c2e3b0ad4"`; **121,034,705,620 bytes** |
+
+For captions use **config `default`, split `default`**, and require `survey == "legacy_dr10"`. Do not substitute the `tiny` split. The recorded caption locators refer directly to the pinned Parquet files, not an index in a filtered Hugging Face dataset.
+
+| Roster columns | Meaning |
+|---|---|
+| `paired_targetid` | String representation of the feature table's `object_id_provabgs`; use its PROVABGS namespace when recovering raw rows. |
+| `caption_object_id` | Legacy imaging identifier from the caption table's `object_id`; it is a different namespace from the paired ID. |
+| `paired_ra`, `paired_dec`, `caption_ra`, `caption_dec` | Source coordinates in degrees. |
+| `separation_arcsec` | Spherical angular separation in arcseconds. |
+| `paired_candidate_count`, `caption_candidate_count` | Number of neighbors within the matching radius in each direction. |
+| `one_to_one_row_candidate`, `unique_source_ids`, `provisional_unambiguous_candidate` | Metadata matching flags; they do not certify physical identity, image quality, or caption correctness. |
+| `paired_shard`, `caption_shard` | Paths relative to the corresponding pinned dataset repository. |
+| `paired_row`, `caption_row` | **Zero-based row offset in the original, entire Parquet shard**, before any filtering. Row numbering continues across its row groups. |
+
+Read identifiers as strings or exact integers, never floating-point numbers; these astronomical IDs can exceed JavaScript's exact integer range. The row numbering was verified against the audit's `enumerate` over the complete Parquet table.
+
+## Recover the original captions on Aurora
+
+For either source, construct its file URL as:
+
+```text
+https://huggingface.co/datasets/{repository}/resolve/{revision}/{shard_path}
+```
+
+The 646 caption candidates occupy just these two files, totaling **501,366,207 bytes** if downloaded in full:
+
+| Caption shard | SHA-256 of full Parquet file, from the pinned repository's LFS metadata |
+|---|---|
+| `data/default-00008-of-00097.parquet` | `b23b3742eb905213299d320d808841a5d208f140482c351d85436465c94cfd46` |
+| `data/default-00009-of-00097.parquet` | `fea00a8b73d16fda8883009a1c401f0f7af99a57e508f25cfec7e47a05a2b652` |
+
+Minimal extraction recipe (future remote work; no downloads were executed to prepare this package):
+
+```python
+# Requires huggingface_hub and pyarrow in the remote environment.
+import csv, gzip
+from collections import defaultdict
+from huggingface_hub import hf_hub_download
+import pyarrow.parquet as pq
+
+with gzip.open("docs/applications/astroprism_data/overlap_candidates_646.csv.gz", "rt") as f:
+ roster = list(csv.DictReader(f)) # IDs remain strings.
+groups = defaultdict(list)
+for row in roster:
+ groups[row["caption_shard"]].append(row)
+for shard, candidates in groups.items():
+ path = hf_hub_download(
+ repo_id="astronolan/galaxy-descriptions", repo_type="dataset",
+ revision="6890dc0c8fc793867c07a6ce4fce11e51c167d6e", filename=shard,
+ )
+ # Verify the downloaded file's SHA-256 against the table above.
+ table = pq.read_table(path, columns=[
+ "object_id", "survey", "ra", "dec", "caption", "summary",
+ ])
+ for row in candidates:
+ item = table.slice(int(row["caption_row"]), 1).to_pylist()[0]
+ assert str(item["object_id"]) == row["caption_object_id"]
+ assert item["survey"] == "legacy_dr10"
+ # Also check coordinates against the roster before accepting the row.
+ # Save full caption, summary, original locator, and provenance together.
+```
+
+Use the full `caption` as the candidate target. Keep `summary` for provenance and inspection; it is not a substitute for validating the full description. These descriptions are synthetic weak supervision. Preserve the original text separately from any crop-validated revision. Grouping by shard avoids downloading the same file per object. Selective Parquet column/range reads can reduce transferred bytes further.
+
+The paired-feature locators span **18 shards** totaling **4,705,846,446 bytes** if downloaded in full. They allow retrieval of the released `embedding_image` and `embedding_spectrum` and verification of `object_id_provabgs`, `ra_provabgs`, and `dec_provabgs`. They do **not** locate raw pixels or flux arrays; use the raw tutorial source below for encoder inputs. Downloading these feature shards is optional when recomputing features with the official encoders.
+
+## Recover matching raw images and spectra
+
+The official tutorial HDF5 contains **115,404** unique PROVABGS objects, including every ID in the paired-feature release. Its URL is not a Git revision: check the recorded ETag and length before reading, and stop to revalidate if either changes. Avoid downloading the entire 121 GB file solely to extract these candidates; use a seekable HTTP-range reader with HDF5, require HTTP 206 responses, and restrict reads to IDs and selected rows.
+
+1. Read `provabgs_desi_edr_sv3/object_id` as exact IDs and construct `str(ID) -> raw row index`. Require uniqueness. Map each roster `paired_targetid` through this lookup; do not use the Parquet row offset as an HDF5 row offset.
+2. Cross-check `desi/object_id`, `desi/ra`, `desi/dec`, the PROVABGS coordinates, and `legacysurvey/object_id`, `legacysurvey/ra`, `legacysurvey/dec`. Retain all original IDs and the raw source index.
+3. Extract each selected row's `legacysurvey/image_array`, `image_band`, `image_ivar`, `image_mask`, `image_scale`, and `object_mask` under the `legacysurvey/` group, plus `desi/spectrum_flux`, `desi/spectrum_ivar`, `desi/spectrum_lambda`, and `desi/spectrum_mask`. Preserve units and masks. The image shape is `[4,160,160]`; the spectrum length is `7781`.
+4. Read per-object band labels and use the official AION preprocessing; do not infer channel order from the array shape. Retain calibrated arrays, not only RGB previews. Apply and record the AION image crop and spectral preprocessing before caption acceptance.
+
+The full raw file has **95 rows where DESI and PROVABGS IDs differ**, so blindly joining through `desi/object_id` is unsafe. A local cross-check of this particular roster found **646 unique raw PROVABGS matches, with DESI and PROVABGS IDs agreeing for all 646**. Recheck this on extraction; any discrepancy must be resolved from source coordinates/identities or excluded, never silently relabeled. Identifier agreement still does not establish image–spectrum physical correspondence.
+
+## Acceptance and training handoff
+
+Before any candidate becomes supervision, verify the physical galaxy match, imaging survey/release, spectroscopic aperture, and that every accepted caption claim is supported by the **exact student image crop**. Remove or revise claims about structures or companions outside the crop. Record raw-source checksums, preprocessing, rejected claims, quality flags, and accepted/rejected status per object. Retain the original roster unchanged and write a separate accepted manifest.
+
+Freeze object-level train/validation/test assignments before deriving image captions, spectral captions, conversations, preferences, or reward examples. Every view, question, and generated derivative of one galaxy stays in the same split. The accepted count can be smaller than 646; 646 is the total candidate pool, not the training count. The training plan defines the subsequent image alignment, spectral captioning, joint model, instruction SFT, evaluation, and optional RL sequence.
+
+This directory is self-contained provenance and source-location information. The local analysis scripts, reports, environments, caches, and 32-object diagnostic arrays are not required to interpret this roster and are intentionally outside the portable package.
diff --git a/docs/applications/astroprism_data/overlap_candidates_646.csv.gz b/docs/applications/astroprism_data/overlap_candidates_646.csv.gz
new file mode 100644
index 0000000..b890099
Binary files /dev/null and b/docs/applications/astroprism_data/overlap_candidates_646.csv.gz differ
diff --git a/docs/applications/astroprism_synthetic_data_prompts.md b/docs/applications/astroprism_synthetic_data_prompts.md
new file mode 100644
index 0000000..d20f1ba
--- /dev/null
+++ b/docs/applications/astroprism_synthetic_data_prompts.md
@@ -0,0 +1,271 @@
+# AstroPrism synthetic-data prompts for the 646-candidate pilot
+
+These are proposed reusable prompts and data contracts, **adapted for this project, not copied paper prompts**. They accompany the [training plan](astroprism_training_plan.md) and [evaluation plan](astroprism_benchmark_landscape.md). No synthetic records have been generated by this document.
+
+Use the [646 positional candidates](astroprism_data/overlap_candidates_646.csv.gz), resolve physical identities and raw observations using the [data contract](astroprism_data.md), then split by physical source. A provisional split is **516 training / 65 validation / 65 test** if all 646 are eligible and grouping permits it; reductions and grouped-split deviations must be recorded. Existing image captions are candidate supervision, not verified ground truth. The earlier 10,000-training-galaxy design remains a later expansion.
+
+The sequence is image-caption alignment while measuring spectral facts and generating spectral captions, followed by spectrum/joint caption alignment, instruction SFT, and a later verifier-grounded RL experiment. AION-1 Base stays frozen; its modality-specific image and spectrum codecs feed the shared encoder, followed by separate projectors into Qwen3-4B. Caption alignment initially trains projectors; SFT adds Qwen LoRA. These prompts do not launch any stage.
+
+AstroSpecLM uses intermediate measured/catalog facts for conversations (Appendices A.2–A.3) and captions (A.5). Its main model uses SpecMAE plus SFT; caption alignment is an ablation training projector and LoRA. Our frozen-AION/projector-first schedule is an adaptation. [AstroSpecLM methodology](https://arxiv.org/html/2609.07102)
+
+## Workflow and generation boundary
+
+1. Retrieve each matched image, spectrum, and existing caption; confirm that they refer to the intended source and that the caption's crop/bands match the image actually presented to the student.
+2. Audit the existing image captions with Prompts 1 and 5. Keep approved original captions; edit or regenerate only unsupported/incompatible ones, preserving the original and change provenance outside student text. Image alignment can begin on the training split while the spectral work proceeds.
+3. Extract spectral facts with Contract 2 and generate spectrum-only captions with Prompt 3. Produce paired captions from independently supported image and spectral facts only after both views are available.
+4. Train the spectrum projector and continue image/spectrum/joint alignment using their corresponding valid targets. Different captions and modality views remain descendants of the same split source.
+5. Generate image-only, spectrum-only, and joint instructions with Prompt 4, verify with Prompt 5, then perform SFT. For this small pilot, start with one caption per available view and two accepted conversations per modality per training source, with approximately 20% containing short evidence/uncertainty follow-ups where supported. This gives a nominal 3,096 training conversations at 516 sources, subject to acceptance. Do not multiply paraphrases to imply a larger independent sample count.
+6. Only after SFT evaluation, build the small verifiable task pool specified in Prompt 6 for a controlled RL experiment. Keep held-out source IDs and all their descendants outside the training/reward pool.
+
+Freeze the split before new captions or conversations are generated. Training teachers must not inspect held-out examples for prompt tuning. Evaluation references require separately checked measurements and expert review; a generated held-out caption is not automatically an independent gold answer. Record generator version, decoding settings, prompt hash, source revision, exact input view and all filtering decisions.
+
+## Shared contract: observations, references, and permissions
+
+Create source-disjoint splits **before** generation; keep repeated observations of the same physical source together. Apply an additional sky/field separation where overlapping cutouts or observing conditions could leak between splits. Teachers may receive privileged reference information; the student receives only the declared observational inputs and ordinary conversation text.
+
+Each example has three separate components:
+
+1. `student_input`: image and/or observed-frame spectrum, public input metadata, and conversation history.
+2. `teacher_reference`: typed facts, independently obtained supervision labels, and scientific constraints.
+3. `audit_sidecar`: source IDs, fact IDs, measurement provenance, target permissions, claim-to-fact links, generator/verifier versions, and split membership. **Do not concatenate this sidecar into student prompts or targets.**
+
+Minimal fact record:
+
+```json
+{
+ "fact_id": "spec.continuum.001",
+ "name": "descriptive_or_measured_property",
+ "value": null,
+ "unit": null,
+ "uncertainty": null,
+ "provenance": "image_pixels | spectral_measurement | catalog_label | supplied_metadata",
+ "source_reference": "versioned_record_or_pipeline_output",
+ "evidence_kind": "direct | derived | inferred | external_label",
+ "required_student_inputs": ["spectrum"],
+ "quality": "valid | low_confidence | unavailable",
+ "status": "measured | upper_limit | not_detected | out_of_coverage | masked | not_assessed",
+ "evidence_locator": {},
+ "allowed_uses": ["caption", "grounded_explanation"],
+ "restrictions": []
+}
+```
+
+`evidence_locator` contains image regions/bands or observed wavelength intervals and relevant quality flags. Preserve original measurement units, conventions, uncertainty definitions, and extraction versions. A measured detection significance is distinct from a teacher's subjective confidence. Do not invent calibrated probabilities for VLM descriptions.
+
+**Permission rule:** a fact is usable in a grounded description only when its required inputs are available and its scientific restrictions are satisfied. A catalog label may also be an explicitly authorized supervised prediction target, but that does not make it directly observed evidence. For example, image-only redshift estimation can use spectroscopic redshift as a training label; its answer must be framed as an estimate, and its explanation cannot claim to observe spectral lines. Keep these prediction tasks separate from observational captions.
+
+**Processed-input rule:** admissibility refers to the exact processed student view, including crop, resize, selected bands, wavelength coverage, normalization, and supplied calibration metadata. A fact visible in a teacher’s larger image may fall outside the student’s calibrated four-band AION 96-by-96 crop. The teacher rendering is only a view of those bands, not a replacement for the calibrated student input. Add requirements such as `metadata.pixel_scale` or `metadata.flux_scale` when needed. Validate the fact after input transformations.
+
+**Redshift rule:** when evaluating redshift inference, keep spectra in the observed frame. Do not feed the true redshift, true-redshift rest-frame transformations, answer-bearing filenames, or target-derived input features to the student. An independently inferred candidate redshift is a model/measurement output, not a hidden input. Line identifications used as supporting evidence must have defensible support; associations selected only by consulting the hidden answer must not be described as independent confirmation.
+
+## Prompt 1 — Image fact extraction with a VLM
+
+Teacher input: the actual image; declared band-to-channel mapping; pixel scale/WCS when available; image masks and quality flags. Omit answer-bearing catalog properties during this perception pass.
+
+```text
+You are extracting a conservative image evidence record for astronomical
+training data. Inspect only the supplied image and declared image metadata.
+
+Record visible source extent, shape, concentration, resolved structure,
+relative appearance across supplied bands, nearby sources, and artifacts.
+Distinguish observation from interpretation. “Unresolved source” does not
+establish “star”; a color composite does not establish calibrated color.
+
+Use coarse descriptive categories when appropriate. Do not invent numerical
+photometry, angular sizes, redshift, stellar mass, spectral lines, or physical
+classifications. Numerical measurements must come from a supplied documented
+measurement; angular scale claims require valid scale metadata. Mark ambiguous
+features as low confidence. Do not equate “not visibly resolved” with absent.
+
+For each fact, return its image region/bands, evidence kind, quality, and
+limitations. Set required_student_inputs=["image"]. Include only facts you can
+support. Return strict JSON:
+{"facts": [...], "image_quality": [...], "limitations": [...]}.
+```
+
+Review uncertain morphology facts before using them as categorical targets. A second VLM agreeing is useful screening, not scientific ground truth.
+
+## Contract 2 — Spectrum facts from measurement code
+
+**This is an extraction contract for deterministic, versioned analysis code, not a request for an LLM to invent measurements.** Existing trusted pipeline products can supply measurements when provenance and quality flags are retained.
+
+Input: observed wavelength array, calibrated flux and units, inverse variance/uncertainty, valid-pixel mask, instrumental resolution, arm boundaries, sky/telluric flags, and declared calibration metadata. Do not silently impute missing wavelengths or treat zero weight as zero flux.
+
+Output where measurable:
+
+- Valid wavelength coverage, masked intervals, and documented continuum/quality summaries.
+- Emission/absorption candidates with observed centroids, widths, integrated fluxes or equivalent widths, uncertainties, detection statistics, fit diagnostics, and evidence intervals.
+- Defensible line identifications, competing identifications when ambiguous, and any independently inferred redshift with uncertainty and method.
+- Line ratios and other derived quantities only when their components, units, uncertainty propagation, and required corrections are valid.
+- Separate `not_detected`, `upper_limit`, `out_of_coverage`, `masked`, and `not_assessed`. Nondetection requires a valid search region and a documented detection threshold. An upper limit requires an explicit numerical method and confidence convention. Neither establishes physical absence.
+- Separate external catalog labels from observation-derived measurements. Preserve aperture/fiber geometry; a spectrum may sample only part of an extended image source.
+
+All thresholds, fit windows, quality rules, and uncertainty conventions belong in versioned configuration. The LLM receives measured facts and may verbalize them; it must not backfill missing numerical fields. Null means unavailable, not zero.
+
+For this galaxy pilot, assess continuum shape, a valid 4000-Å break measure, and covered galaxy emission/absorption features. Suggested candidates include [O II], [O III], Hα/Hβ, [N II], [S II], Ca H/K, Mg b and Na D. This list is a measurement checklist, never evidence that every source contains these lines. Record blending and resolution limits rather than force individual detections. Stellar temperature and molecular-band tasks are outside the initial galaxy-only scope.
+
+If the measurement pipeline uses a catalog redshift to select rest-frame windows, tag the result `conditioned_on_catalog_redshift=true` and retain the corresponding observed-frame interval. Such facts can supervise a description when their support is checked in the observed spectrum, but they are not an independent redshift recovery test. Keep catalog-redshift rest-frame arrays and selected-line identities out of the student inputs for redshift inference. Preserve the AION checkpoint's own preprocessing contract rather than substituting AstroSpecLM preprocessing.
+
+For each candidate produce a machine-readable fact record even when no useful caption is possible. An unavailable measurement should reduce the task pool, not trigger an invented answer. Caption the training, validation and test sources in isolated split workflows; only training captions enter alignment. Expert-check evaluation facts independently of the training caption teacher.
+
+## Prompt 3 — Stage 1 generative caption construction
+
+Teacher input: `active_modalities`, admissible facts, modality-specific limitations, and optional image observation. Produce image-only, spectrum-only, and optionally paired captions. The training input contains the corresponding observation tokens and a short caption request; the target is the caption. Use next-token loss on the assistant caption.
+
+```text
+Write a concise observational caption for the declared student inputs.
+Use only facts whose allowed_uses include caption and whose required inputs
+are present. External catalog labels are not observational caption facts
+unless the student explicitly receives that metadata and you identify it as
+supplied information. Do not mention hidden labels, source IDs, fact IDs,
+the reference table, or the data-generation process in the caption.
+
+Describe 2–4 salient supported facts, preserving uncertainty and limitations.
+Use numerical values only when measured and useful. Never replace missing
+coverage with a nondetection. For paired inputs, attribute image findings to
+the image and spectral findings to the spectrum; avoid unsupported physical
+connections and account for aperture differences. Avoid generic filler.
+
+Return strict JSON with two distinct fields:
+{"student_caption": "...",
+ "audit_sidecar": {"claims": [{"text": "...", "fact_ids": [...]}],
+ "active_modalities": [...], "limitations": [...]}}.
+If no defensible caption can be written, return student_caption=null and a
+rejection reason in audit_sidecar.
+```
+
+Sample student requests: “Describe the visible morphology in this image.” / “Describe the main features supported by this observed-frame spectrum.” / “Summarize what the image and spectrum each show.” Do not use the same caption unchanged for all three input conditions.
+
+## Prompt 4 — Stage 2 instruction conversation construction
+
+Teacher input: admissible image/spectral facts; separately marked supervised labels; available student inputs; task family; single- or multi-turn requirement; answerability policy. Teachers may use reference facts to write grounded answers, but the student must learn to recover the relevant information from its own inputs.
+
+```text
+Generate one scientifically grounded instruction-following training example.
+Inputs specify active modalities, permitted tasks, reference facts, separately
+authorized prediction targets, and single-turn or multi-turn format.
+
+Choose a question that can be answered at the declared level of certainty.
+Use an external label as an answer only for an explicitly authorized
+supervised prediction task. Do not present that label as direct observational
+evidence or construct an explanation that the available input cannot support.
+Include a concise evidence summary when useful; do not invent reasoning,
+measurements, causal mechanisms, or certainty. Do not output hidden reasoning
+traces. Use ordinary observational language rather than fact IDs.
+
+Image-only: use visible morphology and valid image measurements. Spectral
+lines and exact catalog properties are not image evidence.
+Spectrum-only: use measured spectral features and justified inferences;
+do not describe morphology absent an image.
+Joint: ask at least one question that uses information from both modalities,
+or asks which modality supports each finding. Combining two facts does not
+prove causation. Respect fiber/aperture mismatch and possible ambiguity.
+
+For unanswerable/partially answerable examples, explain the specific missing
+measurement or limitation and answer the supported part. Distinguish missing
+coverage, low quality, upper limits, and nondetections. Do not invent an
+insufficiency just to create a refusal.
+
+For multi-turn examples, maintain the declared available modalities and all
+prior stated facts. Vary follow-ups meaningfully: clarification, evidence,
+uncertainty, or comparison. Record when an earlier turn reveals a target;
+a later answer conditioned on that reveal is not independent prediction.
+Do not put the desired answer or target-derived metadata into the question
+for a task intended to evaluate independent estimation.
+
+Return strict JSON:
+{"student_messages": [{"role": "user", "content": "..."},
+ {"role": "assistant", "content": "..."}],
+ "audit_sidecar": {"task_family": "...", "active_modalities": [...],
+ "answerability": "full | partial | insufficient",
+ "claims": [{"turn": 1, "text": "...", "fact_ids": [...],
+ "use": "observation | inference | supervised_label"}],
+ "conditioning_on_prior_targets": [], "limitations": [...]}}.
+```
+
+Grounded question patterns, instantiated only when the record supports them:
+
+- **Image:** “What visible structure does the source have, and which features are uncertain?”
+- **Spectrum:** “Which measured features support your interpretation, and what does the data quality prevent you from concluding?”
+- **Joint:** “How does the image morphology complement the spectral characterization? State the evidence from each and any aperture limitation.”
+- **Partial answer:** “Can these inputs establish this requested property? Explain the information that is present and the measurement still needed.”
+- **Follow-up:** “Which observation supports that statement?” / “How strong is the evidence?” / “Would the conclusion remain supported if only the image were available?”
+
+Use both single-turn and multi-turn records. Set mixture proportions after measuring coverage and difficulty, rather than presenting arbitrary ratios as established scientific practice. Mask user/instruction tokens in the supervised loss; audit sidecars never enter the student sequence. With modality dropout, select or regenerate a valid target for the remaining inputs; blindly reusing a joint answer can manufacture hallucination training.
+
+## Prompt 5 — Independent verification and filtering
+
+Use a separate verifier pass with access to the original observations where it can assess them, measurement outputs, provenance, permissions, student inputs, and the candidate record. Numerical consistency, schemas, split integrity, forbidden metadata, and coverage/detection rules should also be checked in code. A model verifier does not replace independent expert evaluation.
+
+```text
+Audit this generated training record against the exact student-visible inputs
+and the supplied evidence. Treat the candidate answer and its claimed fact
+links as untrusted. Use only evidence retained in the student’s processed inputs and visible
+calibration metadata. Check every scientific assertion, numerical value/unit,
+uncertainty, modality attribution, and limitation. Flag external labels
+presented as observations, unsupported causal stories, unjustified precision,
+false nondetections, and target leakage through inputs or earlier turns.
+
+Check that the question is useful and answerable as labeled. For joint tasks,
+check whether both modalities actually contribute. Relevance to both does
+not mean the observations prove the same property. Flag aperture mismatch
+where it changes the interpretation. Do not infer that mismatched inputs
+must be detectably inconsistent; some mismatches are observationally ambiguous.
+
+Return strict JSON:
+{"decision": "accept | revise | reject",
+ "checks": {"schema": true, "supported_claims": true,
+ "modality_permissions": true, "numeric_consistency": true,
+ "uncertainty": true, "answerability": true,
+ "no_target_leakage": true, "conversation_consistency": true},
+ "issues": [{"code": "...", "severity": "error | warning",
+ "turn": 1, "claim": "...", "reason": "...",
+ "fact_ids": [...]}],
+ "requires_expert_review": false,
+ "revision_instructions": []}.
+Do not silently rewrite a faulty record and mark it accepted. Revisions must
+be regenerated and verified again. Reject examples whose required evidence
+cannot be established; do not fabricate evidence to rescue them.
+```
+
+## Prompt 6 — Later verifiable task construction for RL
+
+This is a proposed data contract for the stage **after SFT**, not a claim that RL is implemented or justified by the present sample size. Use only training-source observations. A reward is eligible only when a versioned independent checker can assess the requested answer from trusted labels or measurements. Suitable initial tasks include a measured numerical property with a fixed tolerance, a covered line-state classification, or attribution of a known supported feature to the correct modality. Avoid broad physical interpretations for which the checker cannot establish correctness.
+
+```text
+Construct one short task for a verifier-grounded reinforcement-learning pool.
+You receive training-source observations, vetted facts, and an explicit list
+of available deterministic verifier functions. Select only a task supported
+by these inputs and one of those functions. Do not invent a checker, a target,
+a tolerance, or a human-agreement score. The reward specification is private
+supervision and must never be included in the student's question or context.
+
+Give the question, student-visible modalities/metadata, structured answer
+schema, and the exact authorized verifier identifier and target reference.
+Use the supplied units, tolerances and quality requirements unchanged.
+If the answer cannot be independently verified, return eligible=false.
+If a task is intended to assess insufficiency, its coverage/quality rule must
+be explicit and independently checkable; do not label an arbitrary unknown
+property as unanswerable just to create a rewardable refusal.
+
+Return strict JSON:
+{"eligible": true,
+ "student_question": "...",
+ "student_inputs": ["image", "spectrum"],
+ "answer_schema": {"answer": "...", "unit": "..."},
+ "private_reward_spec": {"verifier_id": "...", "verifier_version": "...",
+ "target_reference": "...", "tolerance_reference": "...",
+ "quality_rule_reference": "..."},
+ "audit_sidecar": {"source_split": "train", "fact_ids": [...],
+ "task_family": "...", "limitations": [...]}}.
+```
+
+Implement and unit-check the reward independently of the generator. Scientific correctness should dominate; schema compliance is a small auxiliary signal, not a substitute for correctness. Log parse failures and check reward edge cases, unit conversions, boundary values and empty answers. Do not reward answer length, persuasive prose, agreement with the same teacher, unsupported explanations, or student self-assessment. When an answer has a correct scalar and an unverifiable explanation, do not treat the scalar's reward as validation of that explanation.
+
+Start from the saved SFT checkpoint. For the first RL pilot, freeze AION and the aligned projectors/type embeddings and update only Qwen LoRA, matching the training plan. Select the algorithm and regularization against the SFT reference only after validating the reward and Aurora runtime; no PPO/GRPO implementation is assumed here. Compare SFT and RL checkpoints on the same locked independent test set and separate reward score from scientific quality. With approximately 516 training galaxies, monitor memorization and retain SFT as the default if the RL pilot does not improve held-out results.
+
+## Evaluation boundary
+
+Evaluate on held-out physical sources with original data and independently checked targets, including an expert-audited explanation set. Report quantitative task accuracy and calibration separately from prose quality. Include image-only, spectrum-only, joint, and language-only controls; remove/swap relevant modalities and mask scientifically relevant evidence. A teacher/verifier agreement rate is a data-quality screening measure, not proof of astrophysical correctness or successful cross-modal reasoning.
+
+For incremental conversations, an answer may use only observations introduced by that turn. The first implementation places both observations before the conversation; it cannot claim a test in which the spectrum becomes available only later without implementing and checking the intended causal positions. See [evaluation and interleaving contracts](astroprism_benchmark_landscape.md).
diff --git a/docs/applications/astroprism_training_plan.md b/docs/applications/astroprism_training_plan.md
new file mode 100644
index 0000000..43fa4ab
--- /dev/null
+++ b/docs/applications/astroprism_training_plan.md
@@ -0,0 +1,183 @@
+# AstroPrism: image, spectrum and language pilot with 646 candidates
+
+**Last updated**: 2026-09-16. The plan below is written against PRISM's existing Zone-A training path; it proposes work, it does not describe shipped code.
+
+## 1. Goal and sequence
+
+Prototype a model that accepts an image, a spectrum or their matched pair and produces grounded descriptions and answers. Begin with the **646 candidate matches** between released paired AION galaxies and AION-Search image captions, using **AION-1 Base** and **Qwen3-4B**.
+
+```text
+Recover candidates -> verify identity, crop and captions -> freeze source splits
+ |
+ +------------------+-------------------+
+ | |
+ Stage 1A: image-caption alignment Measure spectral facts and generate
+ with AION image tokens verified captions for the same objects
+ | |
+ +------------------+-------------------+
+ |
+ Stage 1B: spectrum-projector warmup, then mixed image / spectrum / joint captions
+ |
+ Stage 2: verified instruction conversations -> projectors + Qwen LoRA SFT
+ |
+ Stage 3: small verifier-grounded RL experiment -> compare against SFT
+```
+
+Stages 1A/1B are **caption alignment**, Stage 2 is **instruction SFT**, and Stage 3 is a later **RL extension**. These are our proposed stages. [AION-Search](https://arxiv.org/abs/2512.11982) supplies candidate image descriptions but trains contrastive retrieval. [AstroSpecLM](https://arxiv.org/abs/2609.07102) motivates evidence-first spectral supervision; its main model uses SpecMAE and SFT, and its caption-alignment ablation trains projector plus LoRA. Our AION-based, frozen-Qwen caption stage is a deliberate adaptation.
+
+**Status:** candidate metadata and this plan exist. Caption text has not been extracted for this roster. The accepted corpus/splits, AION/Qwen integration, generated instructions and training runs remain future work. Local data audits and offline tests did not run AION inference.
+
+## 2. Recover and accept the candidate observations
+
+Use the [646-row roster](astroprism_data/overlap_candidates_646.csv.gz) and [pinned source/recovery instructions](astroprism_data.md). It joins [`EiffL/desi_x_ls_aion`](https://huggingface.co/datasets/EiffL/desi_x_ls_aion) paired features to `legacy_dr10` records in [`astronolan/galaxy-descriptions`](https://huggingface.co/datasets/astronolan/galaxy-descriptions).
+
+All 646 paired IDs and caption IDs are unique, with one-to-one positional candidates within one arcsecond; the largest separation is 0.231 arcseconds. All 646 also map uniquely into the official tutorial raw-file index with matching DESI/PROVABGS IDs. This resolves the larger raw catalog's identifier-discrepancy concern for this roster, but still requires physical identity and input-quality verification.
+
+Recover raw observations from the [official tutorial HDF5 file](https://users.flatironinstitute.org/~polymathic/data/provabgs_desi_ls.hdf5), following the [AION tutorial](https://github.com/PolymathicAI/AION/blob/main/notebooks/Tutorial.ipynb). Select by verified IDs; the first 646 rows are not this roster. Public full-token features may accelerate training after checking extraction provenance and numerical agreement with representative raw-input recomputation.
+
+Where additional raw fields or reference measurements are needed, use release-compatible sources:
+
+| Resource | Role |
+|---|---|
+| [MMU Legacy DR10 South](https://huggingface.co/datasets/hugging-science/mmu_legacysurvey_dr10_south_21) | Calibrated multiband images, uncertainties, masks and imaging metadata |
+| [MMU DESI EDR SV3](https://huggingface.co/datasets/UniverseTBD/mmu_desi_edr_sv3) | Observed spectra, wavelength, inverse variance, masks and reduction metadata |
+| [MMU DESI PROVABGS](https://huggingface.co/datasets/UniverseTBD/mmu_desi_provabgs) | Catalog identity and separately marked model-derived physical-property references |
+| [FastSpecFit EDR](https://data.desi.lbl.gov/doc/releases/edr/vac/fastspecfit/) | Compatible continuum/line measurements, with fit quality and provenance |
+| [Galaxy Zoo](https://data.galaxyzoo.org/) | Independent morphology references where matched human classifications exist |
+
+Do not mix EDR and DR1 fits without reconciling releases. Some fitted properties use photometry jointly with spectra; preserve their actual evidence requirements. Catalog properties are reference targets, not automatically student inputs or directly visible facts.
+
+Acceptance procedure:
+
+1. Read each pinned shard/row and verify its ID and coordinates. Keep large identifiers as strings. Verify raw image/spectrum association, blends, nearby companions and fiber aperture.
+2. Recover calibrated four-band images and observed-frame spectra with wavelength, inverse variance and mask. Check finite values, usable coverage and metadata.
+3. Verify each original caption against a rendering of the **exact AION image crop**. Revise unsupported/out-of-crop claims, retain original text in provenance, and label `released_reused`, `released_revised` or `rejected`.
+4. Group repeat observations, aliases and overlapping sources before assigning splits. Record every exclusion; do not silently add unrelated objects to maintain 646.
+
+Use the full caption as a weak reference, not an unchecked summary or expert label. This overlap is a selected galaxy sample; it does not establish broad performance on other surveys, stars or quasars.
+
+## 3. Split once before fitting or instruction generation
+
+If all candidates survive, target **516 train / 65 validation / 65 test galaxies**, with seed 42. Treat this as an approximately 80/10/10 allocation. Keep physical-source and overlapping-cutout groups intact even when exact totals change. Report the accepted counts after rejection/grouping.
+
+- A source keeps its split through image alignment, spectrum alignment, SFT and RL. All views, paraphrases, conversations and rollouts inherit it.
+- Captions can be generated for every eligible object, but held-out captions/facts stay outside training, reward fitting and training-teacher prompt tuning. Validation sources may support inference-prompt/scorer selection; test sources may not support any tuning. Training on all 646 images and only later holding out some spectra would leak those objects.
+- Use training-only sources for interactive debugging, training-teacher prompt revision, tiny overfit checks and demonstrations. Use validation for checkpoint, inference-prompt and scorer selection. Independently curate and seal test references before selection.
+- Stratify approximately over brightness, redshift, spectral S/N and reliable coarse morphology. Leakage control takes precedence over exact balance.
+- Describe this as evaluation on objects held out from adaptation. Disjointness from AION/Qwen foundation-model pretraining is not established.
+
+The accepted manifest must bind `source_group_id`, split, survey/release/IDs, raw and feature locators/checksums, crop signature, encoder/codec revisions, caption provenance, permitted fact references and quality flags. The committed candidate roster is the input to this process, not an accepted split manifest.
+
+## 4. AION-1 encoders and Qwen interface
+
+AION has modality-specific codecs followed by **one shared pretrained transformer encoder**. Use separate logical image/spectrum passes with shared frozen weights, allowing either modality independently. Follow the [official code](https://github.com/PolymathicAI/AION) and [Base checkpoint](https://huggingface.co/polymathic-ai/aion-base).
+
+| Component | Image branch | Spectrum branch |
+|---|---|---|
+| Raw input | Calibrated `g,r,i,z`, typically `4 × 160 × 160` | Observed wavelength, flux, inverse variance, mask |
+| Official codec | `LegacySurveyImage`, `96 × 96` center crop | `DESISpectrum`, official normalization/padding |
+| Full Base feature tensor | `B × 576 × 768` | `B × 273 × 768` |
+| Trainable projector | `Linear(768,2560) → GELU → Linear(2560,2560)` | Separate projector with identical dimensions |
+| Qwen observation embeddings | `B × 576 × 2560` | `B × 273 × 2560` |
+
+The released feature card specifies the token shapes; Qwen3-4B's [configuration](https://huggingface.co/Qwen/Qwen3-4B/blob/main/config.json) specifies width 2560. Verify shapes on the pinned real checkpoint. Start with full tokens and no resampler, trainable modality type embeddings, and an explicit projector-normalization setting of `none`; monitor embedding norms.
+
+- Pin AION package, model, codecs and Qwen/tokenizer revisions. Local AION metadata recorded checkpoint revision `40541618104bab0fa85c8af68daeb867a720bb8c`; verify existing Aurora weights against that or document the chosen alternative. Codec pins must be explicit, not assumed from the main model loader.
+- Request 576/273 tokens explicitly from `AION.encode`; its default budget can truncate output. The spectrum sequence includes a normalization token plus spectral codes, not 273 raw wavelength bins. Preserve token positions/masks and deterministic ordering.
+- Keep AION codecs and encoder in evaluation mode. Use official preprocessing instead of generic time-series normalization. Check that narrow-line claims survive the lossy codec before making them caption targets.
+- Cache full frozen AION tokens only after provenance and raw/cache agreement checks. Trainable projector outputs cannot be cached across updates. The nominal float16 cache for all 646 paired objects is about **0.84 GB decimal**, excluding metadata/raw inputs.
+
+The first joint baseline supplies `[image embeddings][spectrum embeddings][question/history][assistant answer]`. Its **849 observation positions** leave 1,199 positions for text/delimiters in a 2,048-position budget. Reject overflow instead of silently truncating observations or targets.
+
+Label observation positions, system/user text and padding with `-100`; label only assistant answer tokens for next-token prediction. The model predicts text, not the input observation-code IDs. Use the pinned Qwen non-thinking chat template and verify every multi-turn assistant span. This all-upfront prefix baseline is distinct from a later conversation in which new observations arrive between turns; that needs explicit causal/source-binding checks.
+
+## 5. Stage 1A: image-caption alignment
+
+Train the image projector and image type embedding using accepted training image/caption pairs. Freeze AION and Qwen. Optimize assistant-caption cross entropy and validate grounding and image sensitivity, not only wording similarity. Save the selected image-only checkpoint as a control.
+
+Start with a representative training-only encoder check and a tiny overfit check to establish real gradients, masking and save/reload. These are future implementation checks. At 516 accepted training objects, the nominal corpus is **516 image-caption records**; paraphrases do not add independent galaxies.
+
+Meanwhile, measure spectral facts and generate captions in a separate process using the same source splits. This parallel work requires separately provisioned resources; do not run teacher or training inference on an Aurora login node.
+
+## 6. Parallel spectral caption preparation
+
+Use the [synthetic-data prompts](astroprism_synthetic_data_prompts.md). First distill a compact fact record from measured spectra and compatible catalogs: wavelength coverage, quality, continuum/breaks, reliable line identifications, strengths/EWs and uncertainties. Every fact carries provenance, units, uncertainty, detection/coverage status and permitted student modalities.
+
+Catalog-redshift rest-frame measurements may help the teacher, but the student receives an observed-frame spectrum. Never leak true redshift through student preprocessing or prompt text. Distinguish nondetection, masked/missing coverage and an unassessed feature. Keep model-derived physical properties separately marked.
+
+Generate one concise spectral caption per accepted source, with a separate claim-to-fact audit. Verify numerical targets and recoverability from the processed AION input. Generate joint captions from verified image and spectral facts, attribute evidence to each observation, and respect whole-galaxy versus fiber-aperture limitations. Coexistence of two features does not establish a causal explanation.
+
+Generation may cover the accepted roster, but outputs stay separated by split. Independently review evaluation references; the teacher agreeing with its own prose is insufficient. Record model/prompt versions, generation settings and accepted/revised/rejected counts.
+
+## 7. Stage 1B: spectrum and joint caption alignment
+
+Initialize from the selected Stage 1A checkpoint and add a fresh spectrum projector.
+
+1. **Spectrum warmup:** freeze the image projector/type embedding, AION and Qwen; train only the spectrum projector/type embedding on training spectral captions.
+2. **Mixed alignment:** train both projectors/type embeddings using image-only, spectrum-only and paired captions while AION/Qwen remain frozen. Replay image-only examples to preserve their grounding.
+3. Select using per-modality validation evidence accuracy. Compare image-only performance with the saved Stage 1A checkpoint.
+
+With 516 accepted training sources, mixed alignment has a nominal **1,548 caption records** (516 for each modality condition), subject to caption acceptance. Report actual records and objects separately. Never fabricate a missing observation or treat unavailable facts as zero measurements.
+
+## 8. Stage 2: instruction generation and SFT
+
+Build conversations from validated fact records, not raw unfiltered catalogs or caption paraphrases alone:
+
+- **Image-only:** morphology, visible structures, quality and limitations supported by the image.
+- **Spectrum-only:** line/continuum evidence, supported redshift/numerical tasks and uncertainty/coverage questions.
+- **Joint:** complementary evidence, modality attribution, aperture limitations and tasks that require both observations.
+
+Start with two accepted conversations per training source per modality condition: nominally **3,096 conversations** for 516 sources. Most are single-turn; approximately 20% can contain a short evidence/uncertainty follow-up. Reject unsupported tasks rather than filling a quota. Count conversations and assistant turns separately.
+
+Initialize from mixed alignment. Freeze AION and Qwen base weights; train projectors/type embeddings plus Qwen LoRA. Verify trainable names and optimizer membership. Use explicit assistant-only labels for every turn and homogeneous modality batches initially. All source splits remain fixed.
+
+### Proposed starting settings
+
+These settings are initial proposals, not measured optima. This small corpus primarily tests integration and grounding.
+
+| Setting | Caption alignment | Instruction SFT |
+|---|---|---|
+| Projector LR | `1e-4` initially | `5e-5` |
+| Qwen adaptation | Frozen | LoRA rank 16, alpha 32, dropout 0.05, q/k/v/o projections |
+| LoRA LR | — | `2e-5` |
+| Optimizer | AdamW, weight decay 0.01, clip 1.0 | Same |
+| Effective batch | 16 records via accumulation | 16 conversations via accumulation |
+| Initial duration | Evaluate each epoch; initially 1–3 per subphase | Initially 1–2 epochs |
+| Precision/schedule | bf16 where supported; ~5% warmup then cosine | Same |
+| Full sequence budget | 2,048 including observation tokens | Same; reject overflow |
+
+Nominal updates per epoch are about 33 for image alignment, 33 for spectrum warmup, 97 for mixed captions, and 194 for SFT. Actual counts depend on accepted records and sampling. Check validation at least each epoch and reduce training when held-out grounding regresses. Establish real XPU memory/throughput before estimating duration.
+
+## 9. Stage 3: verifier-grounded RL
+
+Treat RL as a separately evaluated extension initialized from SFT. With roughly 516 training galaxies, memorization and unreliable rewards are substantial concerns. Build and calibrate deterministic verifiers before generating optimization rollouts.
+
+Start with objectively scored tasks: independently supported morphology labels, line identity/coverage decisions, modality attribution, or numeric answers with declared units/tolerances. Use training-source references only and require the relevant student inputs. Explanatory prose may accompany an answer, but verbosity or teacher preference alone must not determine reward.
+
+A possible first algorithm is group-relative policy optimization: sample several completions per training prompt, keep a frozen SFT reference for KL control, and update Qwen LoRA while AION and aligned projectors stay frozen. Choose exact rollout counts, LR and KL settings after Aurora memory/throughput and reward-calibration checks; this document supplies no runnable RL implementation.
+
+Log separate reward components for answer correctness, unit/schema validity, verified evidence/coverage consistency and checkable unsupported claims. Appropriate abstention needs a reference answerability label; blanket refusal must not earn reward. Reject malformed outputs deterministically and test reward exploitation against independently checked cases.
+
+Compare **SFT versus SFT+RL** on the same locked, source-disjoint tasks and decoding budget. Track independent scientific accuracy, uncertainty and image-only retention alongside reward. Higher training reward alone does not select the model. Leave broad unverified free-form reasoning outside the first RL objective.
+
+## 10. Evaluation and implementation requirements
+
+The [benchmark protocol](astroprism_benchmark_landscape.md) maps AstroSpecQA, image resources, paired evaluations and later interleaving checks. Review references for all eligible held-out sources in this small pilot. A 200-source expert panel belongs to future expansion, not a 65-object test split.
+
+Measure supported-claim precision/recall, unsupported claims, uncertainty, task-specific classification/numeric accuracy and modality attribution. Bootstrap by physical object, not paraphrase; report small-sample uncertainty and selection bias. Retain a fixed text-only instruction control.
+
+Compare image-only retention after mixed alignment; image versus spectrum versus joint on identical tasks; matched versus removed/swapped inputs; aligned SFT versus direct SFT where resources permit; and SFT versus RL. A random swapped pair tests sensitivity, not necessarily a recognizable physical contradiction. Removing an irrelevant modality should not alter an answer.
+
+For later interleaved arrival (image → question → spectrum → joint question), earlier answers must not attend to later spectral tokens. Verify this masking/source-binding contract before claiming interleaving support.
+
+Required PRISM work includes AION codecs/shared encoder adapters, a provenance-checked raw/cached feature path, paired source loading, modality masks and assistant-only labels, stage-specific optimizers, Qwen LoRA/save-reload, astronomy evaluation, and separate RL rollouts/verifiers. None is delivered by this planning commit.
+
+Inspect [`src/model.py`](../../src/model.py), [`src/data/multimodal.py`](../../src/data/multimodal.py), [`src/data/collate.py`](../../src/data/collate.py) and [`trainer_zone_a.py`](../../src/training/trainer_zone_a.py). Existing multimodal fusion is not proof of AION support. At this baseline the trainer passes text IDs as labels and toggles backbone gradients broadly; verify explicit multi-turn labels and LoRA-specific freeze rules before use.
+
+On Aurora, first inventory existing weights and the XPU Python environment. Preserve its accelerator-enabled PyTorch. Model loading, feature extraction, teacher inference and training belong on allocated compute resources following repository Aurora/PBS guidance. The publication/initial continuation task downloads neither data nor weights and launches no experiments; execution starts separately on the remote.
+
+## 11. Return to 10,000 training galaxies later
+
+After this pipeline works, expand to **10,000 distinct training galaxies plus separate validation/test**, proposed 1,000 each. The released overlap alone cannot supply this: even allocating all 646 candidates to training leaves at least 9,354 training images needing new captions; holdouts and rejections increase that number.
+
+Generate additional image captions with the same evidence controls, recover spectra, and preserve evaluation integrity when defining larger splits. More conversations from the same 646 objects do not substitute for more independent galaxies. Scale unique qualified sources and independently reviewed evaluation references before reassessing alignment, SFT and RL gains.
diff --git a/docs/index.md b/docs/index.md
index 4bb281a..536be25 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -98,6 +98,19 @@ loadable as Claude agent skills or read as docs. See the
| [getting-started.md](getting-started.md) | Clone, pick a platform, build an environment, run a first job |
| [training/cli.md](training/cli.md) | Unified CLI (`prism`) reference: installation, commands, options, examples |
+### Tutorials — `tutorials/`
+
+End-to-end walkthroughs, grouped by application area. Pages live in
+`tutorials//`; notebooks live next to their scripts in the top-level
+[`applications//`](../applications/) directory.
+
+| Document | Description |
+|----------|-------------|
+| [tutorials/README.md](tutorials/README.md) | Every tutorial by application area, and where a new one goes |
+| [tutorials/astroprism/astroprism_tutorial.md](tutorials/astroprism/astroprism_tutorial.md) | AstroPRISM: align a frozen AION-1 encoder and Qwen3-4B on galaxy images and DESI spectra, from the three alignment stages through inference, on Aurora |
+| [../applications/bioreason/bioreason_demo.ipynb](../applications/bioreason/bioreason_demo.ipynb) | BioReason notebook: train a DNA + language model through projector alignment, SFT, GRPO, and KEGG evaluation on Aurora |
+| [tutorials/material_science/materials_prism_tutorial.md](tutorials/material_science/materials_prism_tutorial.md) | Materials: jointly train on crystal structures and text descriptions to predict a material property |
+
### Platforms — `platforms/`
Environment build and job submission, one page per machine, plus the
@@ -177,6 +190,11 @@ Recorded experiment outcomes. These are dated reports, not usage guides.
| [applications/kegg_curation_pipeline.md](applications/kegg_curation_pipeline.md) | KEGG dataset curation pipeline — design doc |
| [applications/kegg_curation_implementation_status.md](applications/kegg_curation_implementation_status.md) | KEGG curation — what is built and live-tested, stage by stage |
| [applications/materials_multimodal.md](applications/materials_multimodal.md) | Materials property prediction from text + crystal graphs (`src/materials`, `CrystalGraphTokenEncoder`) |
+| [applications/astroprism.md](applications/astroprism.md) | AstroPrism: the 646-candidate image + spectrum + language pilot — start here |
+| [applications/astroprism_training_plan.md](applications/astroprism_training_plan.md) | AstroPrism: AION-1/Qwen architecture, training phases, splits and settings (proposed) |
+| [applications/astroprism_synthetic_data_prompts.md](applications/astroprism_synthetic_data_prompts.md) | AstroPrism: prompt templates and data contracts for caption and instruction synthesis |
+| [applications/astroprism_benchmark_landscape.md](applications/astroprism_benchmark_landscape.md) | AstroPrism: astronomy benchmark survey and the proposed evaluation panels |
+| [applications/astroprism_data.md](applications/astroprism_data.md) | AstroPrism: the 646-candidate roster — pins, checksum, field definitions, recovery |
### API reference — `api/`
diff --git a/docs/tutorials/README.md b/docs/tutorials/README.md
new file mode 100644
index 0000000..90be58b
--- /dev/null
+++ b/docs/tutorials/README.md
@@ -0,0 +1,30 @@
+# Tutorials
+
+Hands-on walkthroughs of PRISM on a scientific application, grouped by
+application area. Each one walks a single configuration from environment and
+data through training.
+
+Tutorial pages live here, in `docs/tutorials//`. Notebooks live in
+[`applications//`](../../applications/), next to the scripts they call.
+
+For reference material on the same applications, such as pipeline designs and
+evaluation invocations, see [`docs/applications/`](../applications/).
+
+| Area | Tutorial | Modalities | Platform |
+|---|---|---|---|
+| Astronomy | [Align AION and Qwen3-4B on galaxy images and spectra](astroprism/astroprism_tutorial.md): image-caption and spectrum-caption alignment, a joint stage, and inference from the published checkpoint | Galaxy image + spectrum + text | Aurora |
+| Biology | [BioReason: train a DNA + language model](../../applications/bioreason/bioreason_demo.ipynb) (notebook): projector alignment, SFT, GRPO, and KEGG evaluation, submitted as chained PBS jobs | DNA + text | Aurora |
+| Materials science | [Jointly train on material descriptions and crystal structures](material_science/materials_prism_tutorial.md): predict a material property from a CIF crystal graph and a text description | Crystal graph + text | CUDA, XPU, or CPU |
+
+## Adding a tutorial
+
+- Put Markdown pages in `docs/tutorials//`, where `` is the
+ application's directory under `applications/`, or the name that directory
+ would take if the application has no scripts there yet.
+- Put a notebook in `applications//`, next to the scripts it runs, and
+ link to it from this page.
+- Add a row to the table above and to the Tutorials section of
+ [`docs/index.md`](../index.md).
+- Run `make docs` before pushing. It checks relative links in Markdown files
+ the same way CI does. Links inside notebooks are not checked, so test those
+ by hand after moving a notebook.
diff --git a/docs/tutorials/astroprism/astroprism_tutorial.md b/docs/tutorials/astroprism/astroprism_tutorial.md
new file mode 100644
index 0000000..ccab641
--- /dev/null
+++ b/docs/tutorials/astroprism/astroprism_tutorial.md
@@ -0,0 +1,364 @@
+# PRISM tutorial: align AION and Qwen3-4B on galaxy images and spectra
+
+Train AstroPRISM on one Aurora node: a frozen AION-1 Base astronomy encoder
+and a frozen Qwen3-4B language model, joined by two small projectors that
+learn to describe calibrated galaxy images and DESI spectra in words.
+
+```
+ image float32[4,160,160] ──► AION image codec ──► 576 tokens ×768 ──► image projector ───┐ (frozen AION)
+ spectrum flux/ivar/mask/λ ──► AION spectrum ──► 273 tokens ×768 ──► spectrum projector ┤
+ float32[7781] codec ▼ spliced into the prompt
+ "Describe the galaxy in this calibrated four-band image." ────────────► Qwen3-4B (frozen) ──► caption
+```
+
+Only the two projectors train — 17,059,840 parameters out of 4.45 B.
+
+Every command below was run on an Aurora compute node while this page was
+written: all three stages as two-step smoke runs, and section 7's inference
+against the released weights, which reproduces the published spectrum caption
+exactly.
+
+The three stages are a sequence, not three independent runs. 1A and 1B each
+produce half the adapters; a composer merges them; the joint stage *resumes*
+from that merge rather than starting cold:
+
+```
+1A ──► image projector ─┐
+ ├──► compose_astroprism_joint_adapters.py ──► joint
+1B ──► spectrum projector ─┘ (joint_adapter_init)
+```
+
+| Stage | Training config | Trains | Starts from |
+|---|---|---|---|
+| 1A. Image alignment | `astroprism_stage1a` | image projector | random init |
+| 1B. Spectrum alignment | `astroprism_stage1b_spectrum` | spectrum projector | random init |
+| Joint | `astroprism_joint_aion` | both, alternating | 1A + 1B, composed |
+
+A reference checkpoint from the joint stage is published at
+[`Sand33p/AstroPRISM`](https://huggingface.co/Sand33p/AstroPRISM), trained with
+exactly the config in this repo. Section 7 runs inference against it.
+
+## Before you start
+
+An Aurora allocation, and access to the two gated Hugging Face repos:
+[`Sand33p/AstroPRISM`](https://huggingface.co/Sand33p/AstroPRISM) (weights) and
+[`Sand33p/AstroPRISM`](https://huggingface.co/datasets/Sand33p/AstroPRISM)
+(data). Request access on each page, then put a token in your environment:
+
+```bash
+export HF_TOKEN=hf_... # from https://huggingface.co/settings/tokens
+```
+
+Use a **classic read** token. A fine-grained token scoped to specific repos
+returns 401 on anything outside its list, which looks identical to not having
+access at all.
+
+## 1. Set up the environment
+
+Build on a login node — it has direct internet access, while compute nodes
+reach the network only through the ALCF proxy. Clone onto `/flare`, not
+`$HOME`, which is too small for the caches this needs:
+
+```bash
+cd /flare//$USER
+git clone https://github.com/AI-ModCon/PRISM.git
+cd PRISM
+
+module use /soft/modulefiles && module load frameworks/2026.1.0
+python -m venv --system-site-packages /flare//$USER/prism-envs/astroprism
+source /flare//$USER/prism-envs/astroprism/bin/activate
+pip install --no-deps -e .
+pip install "hydra-core>=1.3" pyarrow huggingface_hub wandb
+pip install --no-deps polymathic-aion==0.0.2 jaxtyping astropy
+```
+
+`--system-site-packages` plus `--no-deps` is the Aurora rule: the frameworks
+module supplies the XPU-patched PyTorch, and a plain
+`pip install -r requirements/base.txt` pulls a CUDA build from PyPI that
+shadows it. See [platforms/aurora_operations.md](../../platforms/aurora_operations.md).
+
+`polymathic-aion` must go in with `--no-deps` for that reason, but then its own
+requirements do not come along. The module already supplies all but two, so
+`jaxtyping` and `astropy` are named explicitly — without `jaxtyping`,
+`aion.codecs` raises `ModuleNotFoundError` the moment the encoder is built.
+Check before training:
+
+```bash
+python -c "from aion.codecs import ImageCodec, SpectrumCodec; print('aion OK')"
+```
+
+Built this way from a clean clone the venv is about 1.4 MB: everything heavy
+comes from the frameworks module, and `torch` stays the XPU build
+(`2.13.0a0+gitcf30153`). If `site-packages/nvidia/` appears, a CUDA wheel
+stack has shadowed it and the environment is wrong.
+
+Point the caches at `/flare`; home quotas are too small for 4 B-parameter
+checkpoints:
+
+```bash
+export HF_HOME=/flare//$USER/cache/model_weights
+export PRISM_CKPT_ROOT=/flare//$USER/checkpoints
+export ASTROPRISM_ROOT=/flare//$USER/astroprism
+```
+
+**`ASTROPRISM_ROOT` is required.** The model configs resolve the backbone from
+it rather than from a Hub id, so that compute nodes load from a staged copy
+instead of reaching the network:
+
+```yaml
+backbone_id: "${oc.env:ASTROPRISM_ROOT,_local/astroprism}/hf/models/qwen3-4b"
+```
+
+Stage Qwen3-4B there once, from a login node:
+
+```bash
+mkdir -p "$ASTROPRISM_ROOT/hf/models" "$ASTROPRISM_ROOT/hf/cache"
+hf download Qwen/Qwen3-4B \
+ --revision 1cfa9a7208912126459214e8b04321603b3df60c \
+ --local-dir "$ASTROPRISM_ROOT/hf/models/qwen3-4b"
+```
+
+Leave it unset and the configs fall back to a relative `_local/astroprism`,
+which will not exist on a compute node.
+
+Stage the AION encoder too, 1.8 GB at the revision the released checkpoint
+pins. The configs set `aion_local_files_only: true`, so a compute node will
+not fetch it for you:
+
+```bash
+hf download polymathic-ai/aion-base \
+ --revision 40541618104bab0fa85c8af68daeb867a720bb8c
+```
+
+That goes to `$HF_HOME`, which is where the configs look (`aion_cache_dir:
+null`). Without it the run dies in `AION.from_pretrained` with
+`LocalEntryNotFoundError`.
+
+## 2. Check the inputs
+
+```bash
+python -c "
+import torch, src.modalities as m
+print('XPU tiles :', torch.xpu.device_count(), '(expect 12 on a compute node)')
+print('modalities :', [x.value for x in m.Modality])
+from src.encoders import AIONImageEncoder, AIONSpectrumEncoder, AIONMultimodalEncoder
+print('AION encoders importable')
+"
+```
+
+`spectrum` must appear in the modality list. **0 XPU tiles means you are on a
+login node** — fine for this check and for downloads, but training needs a
+compute node.
+
+## 3. Get the data
+
+The dataset holds 646 galaxies with captions for both modalities, split
+**516 train / 65 validation / 65 test**, grouped by physical source so no
+galaxy appears in two splits:
+
+```bash
+python -c "
+from datasets import load_dataset
+images = load_dataset('Sand33p/AstroPRISM', 'image_caption')
+print({k: len(v) for k, v in images.items()})
+"
+```
+
+Model inputs are the raw arrays, not the preview columns:
+`assets/raw/images/image_array.npy` is `float32[N,4,160,160]` in DES-G/R/I/Z
+band order, and the spectra are `float32[7781]` flux, ivar, wavelength plus a
+bool mask. `candidate_index` indexes axis 0 of every array directly.
+
+Two limits the dataset card states plainly, and they bound what any result
+here means: the image captions are **synthetic weak supervision** that can
+contain unsupported interpretation, and the spectrum captions are generated
+from measured facts but remain **unreviewed**. The
+`spectrum_caption_inspection` config is held-out and marked
+`never_train` — do not train on it.
+
+## 4. Stage 1A: image-caption alignment
+Materialize it into the layout the training configs read. The Hugging Face
+package and the configs use different manifest envelopes — the configs want
+`artifact_type: astroprism_raw_images` with an `outputs` block and a
+per-candidate list, while the HF tree ships an upload index and a bare
+`candidates.jsonl` — so one step bridges them:
+
+```bash
+python tools/materialize_astroprism_hf_dataset.py --output-dir _local/astroprism
+```
+
+That writes four small manifests under `_local/astroprism/` pointing at the
+arrays where `snapshot_download` already put them, so nothing large is copied
+twice. `_local/astroprism` is what `src/conf/data/astroprism_*.yaml` reads by
+default. Pass `--snapshot-dir` to reuse a snapshot you already have, and
+regenerate the tree if you ever move or clear the Hugging Face cache, since the
+manifests reference absolute paths.
+
+
+```bash
+python tools/launch_aurora.py \
+ --file experiments/astroprism_stage1a.yaml \
+ --design ASTROPRISM-STAGE1A \
+ --id ASTROPRISM-STAGE1A \
+ --nodes 1 --batch
+```
+
+This trains the image projector only; AION and Qwen stay frozen. `--batch`
+submits with `qsub` and prints a job id. Drop `--batch` to run interactively
+on a node you already hold.
+
+Watch for the merged sequence length in the log: 576 image tokens plus the
+prompt must fit `max_merged_seq_length` (2048 in the shipped config), or
+examples are silently dropped.
+
+A two-step smoke run on one node is the fastest way to confirm the whole path
+before spending queue time:
+
+```bash
+python src/train.py model=prism_astroprism_qwen3_4b \
+ training=astroprism_stage1a +data=astroprism_stage1a \
+ training.max_steps=2 training.warmup_steps=1 \
+ training.device=xpu training.batch_size=1 wandb.mode=offline
+```
+
+It should reach a loss and write a checkpoint:
+
+```text
+Step 1: Loss 2.9253
+Training Zone A: 100%|██████████| 2/2 [00:10<00:00, 5.47s/it, loss=2.93]
+Checkpoint Saved.
+```
+
+The checkpoint is ~9 GB: the frozen Qwen and AION tensors are saved alongside
+the projector that actually trained.
+
+## 5. Stage 1B: spectrum-caption alignment
+
+```bash
+python tools/launch_aurora.py \
+ --file experiments/astroprism_stage1b_spectrum.yaml \
+ --design ASTROPRISM-STAGE1B-SPECTRUM \
+ --id ASTROPRISM-STAGE1B-SPECTRUM \
+ --nodes 1 --batch
+```
+
+Same shape, 273 spectrum tokens, training the spectrum projector. A two-step
+smoke run here prints `loss=4.70` and saves a checkpoint holding
+`projectors.spectrum` only. Note the two
+projectors are configured differently and deliberately so: the image projector
+uses `norm_mode: none`, the spectrum projector `text_scaled_layernorm` with
+`target_norm: 1.04`. Spectra and images do not arrive at the backbone with
+comparable scales.
+
+## 6. Joint alignment
+
+First merge the two adapters. The joint config's `resume_weights_only` points
+at `_local/astroprism/joint_adapter_init`, and the run fails with
+`FileNotFoundError` if it is not there:
+
+```bash
+python tools/compose_astroprism_joint_adapters.py \
+ --image-checkpoint /checkpoints/step_final \
+ --spectrum-checkpoint /checkpoints/step_final \
+ --output-dir _local/astroprism/joint_adapter_init
+```
+
+Each stage-1 checkpoint holds only its own projector — 1A `projectors.image`,
+1B `projectors.spectrum` — and the composer writes the small checkpoint that
+carries both. Set `ASTROPRISM_JOINT_ADAPTER_INIT` to put it elsewhere.
+
+Then launch:
+
+```bash
+python tools/launch_aurora.py \
+ --file experiments/astroprism_joint_aion.yaml \
+ --design ASTROPRISM-JOINT-AION \
+ --id ASTROPRISM-JOINT-AION \
+ --nodes 1 --batch
+```
+
+One AION core serves both routes, and a rank-synchronous batch sampler
+alternates modalities so each projector gets the same number of updates it
+would have had alone. Each example carries **one** non-text modality;
+simultaneous image-plus-spectrum prompting was never trained.
+
+The published checkpoint is this stage at 200 steps, batch 2 per rank on 12
+XPU ranks, BF16, cosine LR 1e-3 with 10 warmup steps, seed 42. Its held-out
+teacher-forced losses:
+
+| Stream | Test loss |
+|---|---:|
+| Combined | 1.4708 |
+| Image-caption | 1.5210 |
+| Spectrum-caption | 1.3006 |
+
+These measure target-token prediction, not scientific correctness.
+
+A two-step smoke run of this stage is the one place both modalities are
+exercised together, and it reports them separately:
+
+```text
+Training Zone A: 100%|██████████| 2/2 [00:31<00:00, 15.96s/it, loss=6.18]
+Test loss metrics: {'test/loss': 5.728, 'test/image_loss': 3.057,
+ 'test/spectrum_loss': 14.780, ...}
+```
+
+Its checkpoint carries both projectors. Those numbers are two steps against
+the released run's 200 — they show the alternating sampler works, nothing
+more.
+
+## 7. Generate from the published checkpoint
+
+The checkpoint is a PRISM `UnifiedTransformer`, **not** loadable with
+`AutoModelForCausalLM.from_pretrained`. The runner used for the published
+examples ships with the weights:
+
+```bash
+hf download Sand33p/AstroPRISM \
+ inference/generate_astroprism_joint_examples.py model.safetensors \
+ --local-dir $HF_HOME/astroprism
+```
+
+It builds the pinned Qwen3-4B (`1cfa9a72…`) and AION Base (`40541618…`),
+then loads `model.safetensors` through PRISM's checkpoint loader. Greedy
+decoding on a held-out spectrum produces:
+
+> At the catalog redshift z ≈ 0.113, the spectrum shows a red, rising
+> continuum and a red, rising [O III] emission line.
+
+**Read that output next to its reference before trusting the model.** The
+published examples disagree with their references in both cases: the image
+example calls a galaxy lenticular where the reference says elliptical, and
+this spectrum asserts `[O III]` where the reference emphasizes a possible
+4000 Å break. A 200-step projector alignment over 516 weakly-supervised
+captions produces fluent astronomy-shaped text, and fluency is not evidence.
+
+## 8. Troubleshooting
+
+**`GatedRepoError` / 401.** Request access on both HF repo pages, and check
+your token is classic-read rather than fine-grained:
+`curl -s -H "Authorization: Bearer $HF_TOKEN" https://huggingface.co/api/whoami-v2`.
+
+**Downloads hang or raise `LocalEntryNotFoundError`.** On a **compute node**
+export `HTTPS_PROXY=http://proxy.alcf.anl.gov:3128` (and `HTTP_PROXY`). On a
+**login node** do the opposite — it reaches Hugging Face directly, and setting
+the proxy there makes requests time out. Check with `env | grep -i proxy`.
+
+**`ValueError: The pinned AION image codec emits exactly 576 tokens`.** The
+config validates the AION contract at construction. Changing
+`aion_num_image_tokens`, `d_img`, the band order or the 160×160 cutout size
+fails loudly rather than silently feeding the encoder something it was not
+trained on.
+
+**`Disk quota exceeded`.** Something is writing to `$HOME`. `HF_HOME` and
+`PRISM_CKPT_ROOT` both belong on `/flare`.
+
+**Identifiers ending in `000`.** Astronomical ids exceed JavaScript's exact
+integer range; read them as strings end to end.
+
+## Where to read further
+
+- [applications/astroprism_training_plan.md](../../applications/astroprism_training_plan.md) — the full architecture and stage rationale
+- [applications/astroprism_data.md](../../applications/astroprism_data.md) — roster provenance, pinned revisions, recovery from upstream sources
+- [applications/astroprism_benchmark_landscape.md](../../applications/astroprism_benchmark_landscape.md) — what a real astronomy evaluation would require
+- [applications/astroprism_synthetic_data_prompts.md](../../applications/astroprism_synthetic_data_prompts.md) — how the captions were generated
diff --git a/docs/materials_prism_tutorial.md b/docs/tutorials/material_science/materials_prism_tutorial.md
similarity index 98%
rename from docs/materials_prism_tutorial.md
rename to docs/tutorials/material_science/materials_prism_tutorial.md
index 9d4461e..e41a5eb 100644
--- a/docs/materials_prism_tutorial.md
+++ b/docs/tutorials/material_science/materials_prism_tutorial.md
@@ -89,7 +89,7 @@ python applications/material_science/stage_materials_text_model.py \
--cache-dir /path/to/shared/huggingface/hub
```
-Then point Hugging Face at that cache when launching training, according to your cluster's setup. The tokenizer and model name passed to `--text-model` must match the staged model. On Polaris, see [the materials guide](applications/materials_multimodal.md) for the PBS launcher and cache environment settings.
+Then point Hugging Face at that cache when launching training, according to your cluster's setup. The tokenizer and model name passed to `--text-model` must match the staged model. On Polaris, see [the materials guide](../../applications/materials_multimodal.md) for the PBS launcher and cache environment settings.
## 5. What PRISM does with the two inputs
diff --git a/experiments/astroprism_joint_aion.yaml b/experiments/astroprism_joint_aion.yaml
new file mode 100644
index 0000000..1cb10f9
--- /dev/null
+++ b/experiments/astroprism_joint_aion.yaml
@@ -0,0 +1,34 @@
+# Joint raw-image + raw-spectrum alignment with one frozen AION core and Qwen3-4B.
+
+experiments:
+ - id: ASTROPRISM-JOINT-AION
+ name: "AstroPRISM joint AION image/spectrum caption alignment"
+ resources:
+ ngpus: 12
+ topology: "12xpu_aurora"
+ walltime: "01:00:00"
+ common_overrides:
+ model: prism_astroprism_joint_qwen3_4b
+ training: astroprism_joint_aion
+ +data: astroprism_joint_aion
+
+ training.task: astroprism_joint_aion
+ training.device: xpu
+ training.batch_size: 2
+ training.gradient_accumulation_steps: 1
+ training.max_steps: 200
+ training.save_every_n_steps: 50
+ training.eval_enabled: true
+ training.eval_every_n_steps: 50
+ training.validation_batches: 0
+ training.freeze_llm: true
+ training.freeze_vit: true
+
+ model.freeze_backbone: true
+ model.freeze_encoders: true
+ data.split: train
+
+ wandb.project: astroprism-joint-aion
+ # This host has no W&B credential. The complete run is retained locally
+ # and can be uploaded later with `wandb sync`.
+ wandb.mode: offline
diff --git a/experiments/astroprism_stage1a.yaml b/experiments/astroprism_stage1a.yaml
new file mode 100644
index 0000000..3c32e5c
--- /dev/null
+++ b/experiments/astroprism_stage1a.yaml
@@ -0,0 +1,42 @@
+# AstroPRISM Stage 1A: raw calibrated images -> frozen AION -> adapter -> Qwen3-4B.
+#
+# The submission command supplies absolute shared-filesystem paths for the
+# checkout, Qwen weights, AION cache, and reviewed data manifests. Keeping those
+# operator-owned paths out of this design makes the experiment reusable while
+# preserving an offline/fail-closed runtime.
+
+experiments:
+ - id: ASTROPRISM-STAGE1A
+ name: "AstroPRISM Stage 1A: frozen AION/Qwen image-caption alignment"
+ resources:
+ ngpus: 12
+ topology: "12xpu_aurora"
+ walltime: "01:00:00"
+ common_overrides:
+ model: prism_astroprism_qwen3_4b
+ training: astroprism_stage1a
+ +data: astroprism_stage1a
+
+ # Keep this explicit so launch_aurora_web.py selects its supported
+ # Accelerate path before Hydra composes the training group.
+ training.task: astroprism_stage1a
+ training.device: xpu
+ training.batch_size: 2
+ training.gradient_accumulation_steps: 1
+ training.max_steps: 100
+ training.save_every_n_steps: 50
+ training.freeze_llm: true
+ training.freeze_vit: true
+
+ # AION is selected by image_encoder_type in the model preset; prevent the
+ # launcher from staging the inactive SigLIP compatibility default.
+ model.image_encoder_id: null
+ model.freeze_backbone: true
+ model.freeze_encoders: true
+
+ data.split: train
+ data.allow_unsafe_candidates: false
+ data.verify_file_hashes: true
+
+ wandb.project: astroprism-stage1a
+ wandb.mode: offline
diff --git a/experiments/astroprism_stage1b_spectrum.yaml b/experiments/astroprism_stage1b_spectrum.yaml
new file mode 100644
index 0000000..e4fdebc
--- /dev/null
+++ b/experiments/astroprism_stage1b_spectrum.yaml
@@ -0,0 +1,43 @@
+# AstroPRISM Stage 1B spectrum-only reviewed-caption alignment.
+#
+# This run aligns raw DESI spectra to promoted, reviewed spectrum captions.
+# Only the immutable train split is enabled for gradients; validation and test
+# remain evaluation-only in the artifact and evaluation stays disabled here.
+# The launch command must supply shared paths through ASTROPRISM_ROOT.
+
+experiments:
+ - id: ASTROPRISM-STAGE1B-SPECTRUM
+ name: "AstroPRISM Stage 1B: frozen AION/Qwen reviewed spectrum captions"
+ resources:
+ ngpus: 12
+ topology: "12xpu_aurora"
+ walltime: "01:00:00"
+ common_overrides:
+ model: prism_astroprism_spectrum_qwen3_4b
+ training: astroprism_stage1b_spectrum
+ +data: astroprism_stage1b_spectrum
+
+ training.task: astroprism_stage1b_spectrum
+ training.device: xpu
+ training.batch_size: 2
+ training.gradient_accumulation_steps: 1
+ training.max_steps: 100
+ training.save_every_n_steps: 50
+ training.freeze_llm: true
+ training.freeze_vit: true
+
+ # No image tower is active; AION spectrum assets are resolved from the
+ # pinned offline cache configured by the model preset.
+ model.image_encoder_id: null
+ model.freeze_backbone: true
+ model.freeze_encoders: true
+ model.projector_norm_mode: text_scaled_layernorm
+ model.projector_target_norm: 1.04
+ model.projector_modality_embed_pos: before_norm
+
+ data.split: train
+ data.accepted_caption_statuses: [accepted]
+ data.verify_file_hashes: true
+
+ wandb.project: astroprism-stage1b-spectrum
+ wandb.mode: offline
diff --git a/src/conf/data/astroprism_joint_aion.yaml b/src/conf/data/astroprism_joint_aion.yaml
new file mode 100644
index 0000000..70ee306
--- /dev/null
+++ b/src/conf/data/astroprism_joint_aion.yaml
@@ -0,0 +1,22 @@
+# Reviewed image-caption and spectrum-caption pairs over the shared 646 roster.
+split: "train"
+max_merged_seq_length: 2048
+
+image:
+ raw_manifest_path: "_local/astroprism/raw_images/manifest.json"
+ candidates_manifest_path: "_local/astroprism/stage1a_accepted/manifest.json"
+ candidates_path: null
+ caption_field: "auto"
+ allow_unsafe_candidates: false
+ accepted_statuses: ["accepted", "approved"]
+ accepted_caption_statuses: ["released_reused", "released_revised"]
+ verify_file_hashes: true
+
+spectrum:
+ raw_manifest_path: "_local/astroprism/raw_spectra/manifest.json"
+ captions_manifest_path: "_local/astroprism/stage1b_spectrum_caption_accepted/manifest.json"
+ captions_path: null
+ accepted_caption_statuses: ["accepted"]
+ verify_file_hashes: true
+
+strict_modality_check: true
diff --git a/src/conf/data/astroprism_stage1a.yaml b/src/conf/data/astroprism_stage1a.yaml
new file mode 100644
index 0000000..0dbed97
--- /dev/null
+++ b/src/conf/data/astroprism_stage1a.yaml
@@ -0,0 +1,18 @@
+# Calibrated Legacy Survey images extracted for the 646-object overlap roster.
+raw_manifest_path: "_local/astroprism/raw_images/manifest.json"
+candidates_manifest_path: "_local/astroprism/stage1a_accepted/manifest.json"
+candidates_path: null
+
+# Production-safe defaults. The accepted artifact preserves the original
+# candidate roster while recording crop-overlap groups and frozen splits.
+split: "train"
+allow_unsafe_candidates: false
+accepted_statuses: ["accepted", "approved"]
+accepted_caption_statuses: ["released_reused", "released_revised"]
+verify_file_hashes: true
+
+caption_field: "auto"
+max_merged_seq_length: 2048
+
+# Require the model and dataset to agree on the emitted non-text modality.
+strict_modality_check: true
diff --git a/src/conf/data/astroprism_stage1b_spectrum.yaml b/src/conf/data/astroprism_stage1b_spectrum.yaml
new file mode 100644
index 0000000..48523f6
--- /dev/null
+++ b/src/conf/data/astroprism_stage1b_spectrum.yaml
@@ -0,0 +1,14 @@
+# Raw observed-frame DESI spectra with promoted, reviewed spectrum captions.
+raw_manifest_path: "_local/astroprism/raw_spectra/manifest.json"
+captions_manifest_path: "_local/astroprism/stage1b_spectrum_caption_accepted/manifest.json"
+captions_path: null
+
+# Preserve the immutable source-level split; only train enters gradient updates.
+split: "train"
+accepted_caption_statuses: ["accepted"]
+verify_file_hashes: true
+
+max_merged_seq_length: 2048
+
+# Require the model and dataset to agree on the emitted non-text modality.
+strict_modality_check: true
diff --git a/src/conf/model/prism_astroprism_joint_qwen3_4b.yaml b/src/conf/model/prism_astroprism_joint_qwen3_4b.yaml
new file mode 100644
index 0000000..28af09c
--- /dev/null
+++ b/src/conf/model/prism_astroprism_joint_qwen3_4b.yaml
@@ -0,0 +1,55 @@
+# Joint AstroPRISM alignment: raw images and spectra -> shared AION -> Qwen3-4B.
+
+backbone_id: "${oc.env:ASTROPRISM_ROOT,_local/astroprism}/hf/models/qwen3-4b"
+tokenizer_id: "${oc.env:ASTROPRISM_ROOT,_local/astroprism}/hf/models/qwen3-4b"
+freeze_backbone: true
+freeze_encoders: true
+
+d_text: 2560
+d_img: 768
+d_spectrum: 768
+
+# A single registered AION transformer serves both frozen raw-modality codecs.
+image_encoder_type: "aion"
+spectrum_encoder_type: "aion"
+image_encoder_id: null
+image_processor_id: null
+image_processor_strict: false
+image_size: 160
+aion_model_name_or_path: "polymathic-ai/aion-base"
+aion_model_revision: "40541618104bab0fa85c8af68daeb867a720bb8c"
+aion_codec_revision: "40541618104bab0fa85c8af68daeb867a720bb8c"
+aion_spectrum_codec_revision: "40541618104bab0fa85c8af68daeb867a720bb8c"
+# huggingface_hub's cache_dir is the hub directory itself, unlike HF_HOME
+# which has hub/ appended. null uses the ambient cache, which is what the
+# released checkpoint's training_config.yaml records.
+aion_cache_dir: null
+aion_local_files_only: true
+aion_image_bands: ["DES-G", "DES-R", "DES-I", "DES-Z"]
+aion_num_image_tokens: 576
+aion_raw_image_size: [160, 160]
+aion_num_spectrum_tokens: 273
+aion_raw_spectrum_length: 7781
+
+attn_implementation: "sdpa"
+
+modalities:
+ - "text"
+ - "image"
+ - "spectrum"
+
+# Preserve the architectures used by the two independently trained adapters.
+projector_norm_mode: "none"
+projector_target_norm: 0.25
+projector_modality_embed_pos: "after_norm"
+projector_modality_embed_scale: 0.02
+image_projector_norm_mode: "none"
+image_projector_target_norm: 0.25
+image_projector_modality_embed_pos: "after_norm"
+image_projector_modality_embed_scale: 0.02
+spectrum_projector_norm_mode: "text_scaled_layernorm"
+spectrum_projector_target_norm: 1.04
+spectrum_projector_modality_embed_pos: "before_norm"
+spectrum_projector_modality_embed_scale: 0.02
+projector_hidden_mult: 1
+projector_num_layers: 2
diff --git a/src/conf/model/prism_astroprism_qwen3_4b.yaml b/src/conf/model/prism_astroprism_qwen3_4b.yaml
new file mode 100644
index 0000000..965fe31
--- /dev/null
+++ b/src/conf/model/prism_astroprism_qwen3_4b.yaml
@@ -0,0 +1,40 @@
+# AstroPRISM Stage 1A: raw Legacy Survey images -> AION -> adapter -> Qwen3-4B.
+
+backbone_id: "${oc.env:ASTROPRISM_ROOT,_local/astroprism}/hf/models/qwen3-4b"
+tokenizer_id: "${oc.env:ASTROPRISM_ROOT,_local/astroprism}/hf/models/qwen3-4b"
+freeze_backbone: true
+freeze_encoders: true
+
+d_text: 2560
+d_img: 768
+
+# AION consumes calibrated four-band 160x160 cutouts directly. The released
+# feature tensors are parity references only, not model inputs.
+image_encoder_type: "aion"
+image_encoder_id: "google/siglip2-base-patch16-224" # Inactive compatibility default.
+image_processor_id: null
+image_processor_strict: false
+image_size: 160
+aion_model_name_or_path: "polymathic-ai/aion-base"
+aion_model_revision: "40541618104bab0fa85c8af68daeb867a720bb8c"
+aion_codec_revision: "40541618104bab0fa85c8af68daeb867a720bb8c"
+# huggingface_hub's cache_dir is the hub directory itself, unlike HF_HOME
+# which has hub/ appended. null uses the ambient cache, which is what the
+# released checkpoint's training_config.yaml records.
+aion_cache_dir: null
+aion_local_files_only: true
+aion_image_bands: ["DES-G", "DES-R", "DES-I", "DES-Z"]
+aion_num_image_tokens: 576
+aion_raw_image_size: [160, 160]
+
+attn_implementation: "sdpa"
+
+modalities:
+ - "text"
+ - "image"
+
+# Existing two-layer 768 -> 2560 adapter.
+projector_norm_mode: "none"
+projector_modality_embed_pos: "after_norm"
+projector_hidden_mult: 1
+projector_num_layers: 2
diff --git a/src/conf/model/prism_astroprism_spectrum_qwen3_4b.yaml b/src/conf/model/prism_astroprism_spectrum_qwen3_4b.yaml
new file mode 100644
index 0000000..1e28467
--- /dev/null
+++ b/src/conf/model/prism_astroprism_spectrum_qwen3_4b.yaml
@@ -0,0 +1,39 @@
+# AstroPRISM Stage 1B proxy: raw DESI spectra -> AION -> adapter -> Qwen3-4B.
+
+backbone_id: "${oc.env:ASTROPRISM_ROOT,_local/astroprism}/hf/models/qwen3-4b"
+tokenizer_id: "${oc.env:ASTROPRISM_ROOT,_local/astroprism}/hf/models/qwen3-4b"
+freeze_backbone: true
+freeze_encoders: true
+
+d_text: 2560
+d_spectrum: 768
+
+# AION consumes the raw observed-frame DESI arrays. Released spectrum
+# embeddings remain parity references and are not model inputs.
+spectrum_encoder_type: "aion"
+aion_model_name_or_path: "polymathic-ai/aion-base"
+aion_model_revision: "40541618104bab0fa85c8af68daeb867a720bb8c"
+aion_spectrum_codec_revision: "40541618104bab0fa85c8af68daeb867a720bb8c"
+# huggingface_hub's cache_dir is the hub directory itself, unlike HF_HOME
+# which has hub/ appended. null uses the ambient cache, which is what the
+# released checkpoint's training_config.yaml records.
+aion_cache_dir: null
+aion_local_files_only: true
+aion_num_spectrum_tokens: 273
+aion_raw_spectrum_length: 7781
+
+# No image tower is active in this spectrum-only model.
+image_encoder_id: null
+attn_implementation: "sdpa"
+
+modalities:
+ - "text"
+ - "spectrum"
+
+# Initialize the adapter output at the measured Qwen text-token norm.
+projector_norm_mode: "text_scaled_layernorm"
+projector_target_norm: 1.04
+projector_modality_embed_pos: "before_norm"
+projector_modality_embed_scale: 0.02
+projector_hidden_mult: 1
+projector_num_layers: 2
diff --git a/src/conf/training/astroprism_joint_aion.yaml b/src/conf/training/astroprism_joint_aion.yaml
new file mode 100644
index 0000000..018c85b
--- /dev/null
+++ b/src/conf/training/astroprism_joint_aion.yaml
@@ -0,0 +1,31 @@
+# Joint adapter continuation from the separately aligned image/spectrum adapters.
+defaults:
+ - zone_a
+ - _self_
+
+task: "astroprism_joint_aion"
+batch_size: 2
+
+# Rank-synchronous alternating modalities give each adapter 100 optimizer
+# updates and the same 2,400 sample presentations as its prior 100-step run.
+max_steps: 200
+finite_epoch_steps: 0
+
+freeze_llm: true
+freeze_vit: true
+learning_rate: 1e-3
+lr_connector: 1e-3
+lr_vit: 1e-5
+weight_decay: 0.0
+warmup_steps: 10
+
+data_num_workers: 0
+viz_every_n_steps: 0
+save_every_n_steps: 50
+
+# Validation is periodic; test is evaluated once after the final update.
+eval_enabled: true
+eval_every_n_steps: 50
+validation_batches: 0
+
+resume_weights_only: "${oc.env:ASTROPRISM_JOINT_ADAPTER_INIT,_local/astroprism/joint_adapter_init}"
diff --git a/src/conf/training/astroprism_stage1a.yaml b/src/conf/training/astroprism_stage1a.yaml
new file mode 100644
index 0000000..76dcaed
--- /dev/null
+++ b/src/conf/training/astroprism_stage1a.yaml
@@ -0,0 +1,25 @@
+# AstroPRISM Stage 1A: align calibrated raw AION image tokens to language.
+defaults:
+ - zone_a
+ - _self_
+
+task: "astroprism_stage1a"
+batch_size: 2
+max_steps: 100
+finite_epoch_steps: 0
+
+# First align the adapter/projector. The AION wrapper keeps its hard codec
+# frozen/fp32 even when this is later overridden to false for encoder tuning.
+freeze_llm: true
+freeze_vit: true
+learning_rate: 1e-3
+lr_connector: 1e-3
+lr_vit: 1e-5
+weight_decay: 0.0
+warmup_steps: 10
+
+# The 646-row mmap-backed dataset is small; avoid multiplying open mappings by
+# default. Operators can raise this after validating a launch environment.
+data_num_workers: 0
+eval_enabled: false
+validation_batches: 0
diff --git a/src/conf/training/astroprism_stage1b_spectrum.yaml b/src/conf/training/astroprism_stage1b_spectrum.yaml
new file mode 100644
index 0000000..6085c92
--- /dev/null
+++ b/src/conf/training/astroprism_stage1b_spectrum.yaml
@@ -0,0 +1,25 @@
+# AstroPRISM Stage 1B reviewed spectrum-caption alignment.
+defaults:
+ - zone_a
+ - _self_
+
+task: "astroprism_stage1b_spectrum"
+batch_size: 2
+max_steps: 100
+finite_epoch_steps: 0
+
+# Train only the spectrum adapter/type embedding against reviewed captions.
+# The AION codec/shared encoder and Qwen backbone remain frozen.
+freeze_llm: true
+freeze_vit: true
+learning_rate: 1e-3
+lr_connector: 1e-3
+lr_vit: 1e-5
+weight_decay: 0.0
+warmup_steps: 10
+
+# The reviewed corpus is small and mmap-backed. Evaluation remains disabled
+# for this alignment run even though held-out captions are evaluation-ready.
+data_num_workers: 0
+eval_enabled: false
+validation_batches: 0
diff --git a/src/config.py b/src/config.py
index 0e3bb4b..df4e77e 100644
--- a/src/config.py
+++ b/src/config.py
@@ -201,6 +201,7 @@ class ModelConfig:
d_table: int = 768
d_ts: int = 512
d_img: int = 1152 # SigLIP-large native dim (use 768 for SigLIP2-base)
+ d_spectrum: int = 768
d_geo: int = 512
d_graph: int = 768
d_dna: int = 1024 # Nucleotide Transformer hidden dim
@@ -214,6 +215,9 @@ class ModelConfig:
max_dna_length: int = 1024
# Image encoder / processor
+ # ``siglip`` preserves the existing RGB image path. ``aion`` selects the
+ # calibrated four-band Legacy Survey path and its pinned raw-image codec.
+ image_encoder_type: str = "siglip"
image_encoder_id: str = "google/siglip2-base-patch16-224"
image_processor_id: str | None = None
image_processor_strict: bool = False
@@ -221,6 +225,31 @@ class ModelConfig:
image_mean: tuple[float, float, float] = (0.5, 0.5, 0.5)
image_std: tuple[float, float, float] = (0.5, 0.5, 0.5)
+ # Pinned AION raw-image path. These fields are intentionally separate
+ # from the SigLIP identifiers above: AION loads both a hard-quantized
+ # image codec and the differentiable shared transformer encoder.
+ aion_model_name_or_path: str = "polymathic-ai/aion-base"
+ aion_model_revision: str = "40541618104bab0fa85c8af68daeb867a720bb8c"
+ aion_codec_revision: str = "40541618104bab0fa85c8af68daeb867a720bb8c"
+ aion_cache_dir: str | None = None
+ aion_local_files_only: bool = True
+ aion_image_bands: tuple[str, str, str, str] = (
+ "DES-G",
+ "DES-R",
+ "DES-I",
+ "DES-Z",
+ )
+ aion_num_image_tokens: int = 576
+ aion_raw_image_size: tuple[int, int] = (160, 160)
+
+ # Pinned AION raw-spectrum path. The base transformer/cache settings above
+ # are shared with the image wrapper, while the modality codec has its own
+ # explicit revision and input/token contract.
+ spectrum_encoder_type: str = "aion"
+ aion_spectrum_codec_revision: str = "40541618104bab0fa85c8af68daeb867a720bb8c"
+ aion_num_spectrum_tokens: int = 273
+ aion_raw_spectrum_length: int = 7781
+
# Transformer / MoE
num_layers: int = 20
num_heads: int = 10
@@ -327,10 +356,11 @@ class ModelConfig:
# === Projector Configuration (Ablation-Ready) ===
# Normalization mode for ModalityProjector
- # Options: "none", "layernorm", "rmsnorm", "l2_token", "l2_sequence", "scale_only",
- # "match_text_stats", "match_text_elemstats"
+ # Options: "none", "layernorm", "text_scaled_layernorm", "rmsnorm", "l2_token",
+ # "l2_sequence", "scale_only", "match_text_stats", "match_text_elemstats"
# - none: No normalization (simple MLP output)
# - layernorm: Standard LayerNorm on output (original behavior)
+ # - text_scaled_layernorm: LayerNorm initialized to the target text-token norm
# - rmsnorm: RMSNorm (like LLaMA/OLMo, no centering)
# - l2_token: Per-token L2 normalization to target_norm
# - l2_sequence: Sequence-level L2 normalization (preserves token differences)
@@ -339,7 +369,7 @@ class ModelConfig:
# - match_text_elemstats: Match element-wise distribution to text (like LayerNorm but scaled)
projector_norm_mode: str = "layernorm" # Default: original behavior
- # Target norm for l2_token, l2_sequence, and scale_only modes
+ # Target norm for text_scaled_layernorm, l2_token, l2_sequence, and scale_only modes
projector_target_norm: float = 0.25
# Where to add the modality embedding
@@ -351,6 +381,20 @@ class ModelConfig:
# For text-matched norms, use: target_norm / sqrt(d_model) ≈ 0.25/45.25 ≈ 0.0055
projector_modality_embed_scale: float = 0.02 # Default: original behavior (norm ~0.9)
+ # Optional adapter-specific settings. None preserves the global setting
+ # above, so every existing single-modality preset remains bit-compatible.
+ # Joint AstroPRISM can retain its established image adapter
+ # (none/after_norm) and spectrum adapter
+ # (text_scaled_layernorm/1.04/before_norm) in the same model.
+ image_projector_norm_mode: str | None = None
+ image_projector_target_norm: float | None = None
+ image_projector_modality_embed_pos: str | None = None
+ image_projector_modality_embed_scale: float | None = None
+ spectrum_projector_norm_mode: str | None = None
+ spectrum_projector_target_norm: float | None = None
+ spectrum_projector_modality_embed_pos: str | None = None
+ spectrum_projector_modality_embed_scale: float | None = None
+
# === Text Statistics Matching (for match_text_stats and match_text_elemstats modes) ===
# These values are based on observed OLMo-1B text embedding statistics
projector_text_norm_mean: float = 0.25 # Target mean of token norms
@@ -389,6 +433,64 @@ def __post_init__(self) -> None:
if self.is_vla and "action" not in self.output_decoders:
self.output_decoders = ["text", "action"]
+ self.image_encoder_type = str(self.image_encoder_type).strip().lower()
+ valid_image_encoder_types = {"siglip", "aion"}
+ if self.image_encoder_type not in valid_image_encoder_types:
+ raise ValueError(
+ "image_encoder_type must be one of "
+ f"{sorted(valid_image_encoder_types)}, got {self.image_encoder_type!r}"
+ )
+
+ # OmegaConf ListConfig values are accepted at the Hydra boundary but
+ # normalized here so downstream construction receives stable tuples.
+ self.aion_image_bands = tuple(str(band) for band in self.aion_image_bands)
+ self.aion_raw_image_size = tuple(int(size) for size in self.aion_raw_image_size)
+
+ if self.image_encoder_type == "aion":
+ expected_bands = ("DES-G", "DES-R", "DES-I", "DES-Z")
+ if self.d_img != 768:
+ raise ValueError(
+ "The pinned AION shared encoder emits d_img=768 features; "
+ f"got d_img={self.d_img}."
+ )
+ if self.aion_num_image_tokens != 576:
+ raise ValueError(
+ "The pinned AION image codec emits exactly 576 tokens; "
+ f"got aion_num_image_tokens={self.aion_num_image_tokens}."
+ )
+ if self.aion_image_bands != expected_bands:
+ raise ValueError(
+ "AION requires calibrated Legacy Survey channels in "
+ f"{expected_bands!r} order; got {self.aion_image_bands!r}."
+ )
+ if self.aion_raw_image_size != (160, 160):
+ raise ValueError(
+ "The verified AION raw-image path expects 160x160 cutouts; "
+ f"got {self.aion_raw_image_size!r}."
+ )
+
+ self.spectrum_encoder_type = str(self.spectrum_encoder_type).strip().lower()
+ if self.spectrum_encoder_type != "aion":
+ raise ValueError(
+ "spectrum_encoder_type must be 'aion', got "
+ f"{self.spectrum_encoder_type!r}"
+ )
+ if self.d_spectrum != 768:
+ raise ValueError(
+ "The pinned AION shared encoder emits d_spectrum=768 features; "
+ f"got d_spectrum={self.d_spectrum}."
+ )
+ if self.aion_num_spectrum_tokens != 273:
+ raise ValueError(
+ "The pinned AION spectrum codec emits exactly 273 tokens; "
+ f"got aion_num_spectrum_tokens={self.aion_num_spectrum_tokens}."
+ )
+ if self.aion_raw_spectrum_length != 7781:
+ raise ValueError(
+ "The verified AION raw-spectrum path expects 7781 samples; "
+ f"got aion_raw_spectrum_length={self.aion_raw_spectrum_length}."
+ )
+
@classmethod
def from_preset(cls, preset_name: str):
"""Build a ``ModelConfig`` from a named entry in ``PRISM_CONFIGS``.
diff --git a/src/data/astroprism_joint_batch_sampler.py b/src/data/astroprism_joint_batch_sampler.py
new file mode 100644
index 0000000..5f3d916
--- /dev/null
+++ b/src/data/astroprism_joint_batch_sampler.py
@@ -0,0 +1,225 @@
+"""Rank-synchronous homogeneous batching for joint AstroPRISM alignment."""
+
+from __future__ import annotations
+
+import math
+import os
+import random
+from collections.abc import Iterator, Sequence
+
+from torch.utils.data import Sampler
+
+from .astroprism_joint_dataset import (
+ ASTROPRISM_JOINT_MODALITIES,
+ ASTROPRISM_JOINT_SPLITS,
+ AstroPRISMJointDataset,
+)
+
+_WORLD_SIZE_ENVIRONMENT_VARIABLES = (
+ "WORLD_SIZE",
+ "ACCELERATE_NUM_PROCESSES",
+ "OMPI_COMM_WORLD_SIZE",
+ "PMI_SIZE",
+)
+
+
+class AstroPRISMJointBatchSampler(Sampler[list[int]]):
+ """Yield rank-synchronous image-only or spectrum-only batches.
+
+ Training produces equal numbers of full batches for both modalities. The
+ shorter modality is deterministically reshuffled and cycled as needed.
+ Batches are emitted in blocks of ``num_replicas`` with one modality per
+ block. With Accelerate whole-batch sharding (``split_batches=False``), all
+ ranks therefore execute the same modality on each optimizer step.
+
+ ``set_epoch`` selects a deterministic epoch-specific ordering. Validation
+ and test never shuffle or repeat examples. Their final batches may be
+ partial, so distributed evaluation should use ``even_batches=False`` to
+ prevent Accelerate from filling a partial batch with another modality.
+ """
+
+ def __init__(
+ self,
+ dataset: AstroPRISMJointDataset,
+ *,
+ batch_size: int,
+ seed: int = 0,
+ drop_last: bool | None = None,
+ num_replicas: int | None = None,
+ ) -> None:
+ if not isinstance(batch_size, int) or isinstance(batch_size, bool) or batch_size <= 0:
+ raise ValueError(f"batch_size must be a positive integer, got {batch_size!r}")
+ if dataset.split not in ASTROPRISM_JOINT_SPLITS:
+ raise ValueError(f"unsupported joint dataset split {dataset.split!r}")
+
+ self.dataset = dataset
+ self.batch_size = batch_size
+ self.seed = int(seed)
+ self.split = dataset.split
+ self.is_training = self.split == "train"
+ self.drop_last = False if drop_last is None else bool(drop_last)
+ self.num_replicas = self._resolve_num_replicas(num_replicas)
+ self.shuffle = self.is_training
+ self.balanced = self.is_training
+ self.epoch = 0
+
+ indices = dataset.modality_indices
+ self._indices = {
+ modality: tuple(indices[modality])
+ for modality in ASTROPRISM_JOINT_MODALITIES
+ }
+ for modality, modality_indices in self._indices.items():
+ if not modality_indices:
+ raise ValueError(f"joint batch sampler requires non-empty {modality} indices")
+
+ if self.is_training and self._training_blocks_per_modality() == 0:
+ longest = max(len(indices) for indices in self._indices.values())
+ raise ValueError(
+ "drop_last=True would produce no complete distributed training block: "
+ f"longest modality has {longest} examples, batch_size={self.batch_size}, "
+ f"num_replicas={self.num_replicas}"
+ )
+
+ @staticmethod
+ def _resolve_num_replicas(explicit: int | None) -> int:
+ if explicit is not None:
+ if not isinstance(explicit, int) or isinstance(explicit, bool) or explicit <= 0:
+ raise ValueError(
+ f"num_replicas must be a positive integer, got {explicit!r}"
+ )
+ return explicit
+
+ for variable in _WORLD_SIZE_ENVIRONMENT_VARIABLES:
+ value = os.environ.get(variable)
+ if value is None or not value.strip():
+ continue
+ try:
+ resolved = int(value)
+ except ValueError as exc:
+ raise ValueError(
+ f"distributed environment variable {variable} must be an integer, "
+ f"got {value!r}"
+ ) from exc
+ if resolved <= 0:
+ raise ValueError(
+ f"distributed environment variable {variable} must be positive, "
+ f"got {value!r}"
+ )
+ return resolved
+ return 1
+
+ @property
+ def sampler(self) -> AstroPRISMJointBatchSampler:
+ """Expose the epoch hook through Accelerate's custom-sampler chain.
+
+ ``DataLoaderShard.set_epoch`` traverses
+ ``BatchSamplerShard.batch_sampler.sampler``. Returning this batch
+ sampler lets that hook reach :meth:`set_epoch` without changing the
+ batches consumed by PyTorch or Accelerate.
+ """
+ return self
+
+ def set_epoch(self, epoch: int) -> None:
+ if not isinstance(epoch, int) or isinstance(epoch, bool) or epoch < 0:
+ raise ValueError(f"epoch must be a non-negative integer, got {epoch!r}")
+ self.epoch = epoch
+
+ def _training_blocks_per_modality(self) -> int:
+ longest = max(len(indices) for indices in self._indices.values())
+ samples_per_world_block = self.batch_size * self.num_replicas
+ if self.drop_last:
+ return longest // samples_per_world_block
+ return math.ceil(longest / samples_per_world_block)
+
+ def _training_batches_per_modality(self) -> int:
+ return self._training_blocks_per_modality() * self.num_replicas
+
+ def _epoch_seed(self, stream: int) -> int:
+ return (
+ self.seed
+ + (self.epoch + 1) * 1_000_003
+ + (stream + 1) * 97_409
+ ) & ((1 << 63) - 1)
+
+ def _draw_training_indices(
+ self,
+ indices: Sequence[int],
+ *,
+ count: int,
+ stream: int,
+ ) -> list[int]:
+ rng = random.Random(self._epoch_seed(stream))
+ drawn: list[int] = []
+ while len(drawn) < count:
+ cycle = list(indices)
+ rng.shuffle(cycle)
+ drawn.extend(cycle[: count - len(drawn)])
+ return drawn
+
+ def _training_batches(self) -> list[list[int]]:
+ block_count = self._training_blocks_per_modality()
+ batch_count = self._training_batches_per_modality()
+ sample_count = batch_count * self.batch_size
+ by_modality: dict[str, list[list[int]]] = {}
+ for stream, modality in enumerate(ASTROPRISM_JOINT_MODALITIES):
+ drawn = self._draw_training_indices(
+ self._indices[modality],
+ count=sample_count,
+ stream=stream,
+ )
+ by_modality[modality] = [
+ drawn[start : start + self.batch_size]
+ for start in range(0, sample_count, self.batch_size)
+ ]
+
+ order_rng = random.Random(self._epoch_seed(len(ASTROPRISM_JOINT_MODALITIES)))
+ batches: list[list[int]] = []
+ for block_index in range(block_count):
+ modality_pair = list(ASTROPRISM_JOINT_MODALITIES)
+ if order_rng.getrandbits(1):
+ modality_pair.reverse()
+ start = block_index * self.num_replicas
+ stop = start + self.num_replicas
+ for modality in modality_pair:
+ batches.extend(by_modality[modality][start:stop])
+ return batches
+
+ def _evaluation_batches(self) -> list[list[int]]:
+ by_modality: dict[str, list[list[int]]] = {}
+ for modality in ASTROPRISM_JOINT_MODALITIES:
+ indices = self._indices[modality]
+ batches = [
+ list(indices[start : start + self.batch_size])
+ for start in range(0, len(indices), self.batch_size)
+ ]
+ if self.drop_last and batches and len(batches[-1]) < self.batch_size:
+ batches.pop()
+ by_modality[modality] = batches
+
+ batches: list[list[int]] = []
+ batch_count = max(len(value) for value in by_modality.values())
+ for batch_index in range(batch_count):
+ for modality in ASTROPRISM_JOINT_MODALITIES:
+ modality_batches = by_modality[modality]
+ if batch_index < len(modality_batches):
+ batches.append(modality_batches[batch_index])
+ return batches
+
+ def __iter__(self) -> Iterator[list[int]]:
+ batches = self._training_batches() if self.is_training else self._evaluation_batches()
+ yield from batches
+
+ def __len__(self) -> int:
+ if self.is_training:
+ return 2 * self._training_batches_per_modality()
+
+ total = 0
+ for indices in self._indices.values():
+ if self.drop_last:
+ total += len(indices) // self.batch_size
+ else:
+ total += math.ceil(len(indices) / self.batch_size)
+ return total
+
+
+__all__ = ["AstroPRISMJointBatchSampler"]
diff --git a/src/data/astroprism_joint_collate.py b/src/data/astroprism_joint_collate.py
new file mode 100644
index 0000000..1b9e1e3
--- /dev/null
+++ b/src/data/astroprism_joint_collate.py
@@ -0,0 +1,96 @@
+"""Modality-routing collator for joint AstroPRISM alignment."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+from .astroprism_joint_dataset import ASTROPRISM_JOINT_MODALITIES
+from .astroprism_stage1a_collate import AstroPRISMStage1ACollator
+from .astroprism_stage1b_spectrum_collate import AstroPRISMStage1BSpectrumCollator
+
+
+class AstroPRISMJointCollator:
+ """Delegate a homogeneous batch to its existing modality collator.
+
+ The explicit ``batch_modality`` tag is required on every input sample. A
+ mixed batch fails before either source collator or the model sees it.
+ """
+
+ active_modalities = {"image", "spectrum"}
+
+ def __init__(
+ self,
+ image_collator: Callable[[list[dict[str, Any]]], dict[str, Any]],
+ spectrum_collator: Callable[[list[dict[str, Any]]], dict[str, Any]],
+ ) -> None:
+ self._collators = {
+ "image": image_collator,
+ "spectrum": spectrum_collator,
+ }
+
+ @classmethod
+ def from_tokenizer(
+ cls,
+ tokenizer: Any,
+ *,
+ max_merged_seq_length: int,
+ ) -> AstroPRISMJointCollator:
+ return cls(
+ AstroPRISMStage1ACollator(
+ tokenizer,
+ max_merged_seq_length=max_merged_seq_length,
+ ),
+ AstroPRISMStage1BSpectrumCollator(
+ tokenizer,
+ max_merged_seq_length=max_merged_seq_length,
+ ),
+ )
+
+ def __call__(self, batch: list[dict[str, Any]]) -> dict[str, Any]:
+ if not batch:
+ return {}
+
+ modalities: list[str] = []
+ payload: list[dict[str, Any]] = []
+ for index, sample in enumerate(batch):
+ if not isinstance(sample, dict):
+ raise TypeError(
+ f"joint sample {index} must be a dict, got {type(sample).__name__}"
+ )
+ modality = sample.get("batch_modality")
+ if modality not in ASTROPRISM_JOINT_MODALITIES:
+ raise ValueError(
+ f"joint sample {index} has invalid batch_modality={modality!r}; "
+ f"expected one of {list(ASTROPRISM_JOINT_MODALITIES)}"
+ )
+ if modality == "image" and "spectrum" in sample:
+ raise ValueError(f"joint image sample {index} unexpectedly contains spectrum")
+ if modality == "spectrum" and "image" in sample:
+ raise ValueError(f"joint spectrum sample {index} unexpectedly contains image")
+ modalities.append(modality)
+ payload.append({key: value for key, value in sample.items() if key != "batch_modality"})
+
+ unique_modalities = set(modalities)
+ if len(unique_modalities) != 1:
+ raise ValueError(
+ "AstroPRISM joint batches must be modality-homogeneous; got "
+ f"{modalities!r}"
+ )
+ modality = modalities[0]
+ collated = self._collators[modality](payload)
+ if not isinstance(collated, dict):
+ raise TypeError(
+ f"{modality} collator returned {type(collated).__name__}; expected a dict"
+ )
+ if "batch_modality" in collated and collated["batch_modality"] != modality:
+ raise ValueError(
+ f"{modality} collator returned conflicting "
+ f"batch_modality={collated['batch_modality']!r}"
+ )
+ output = dict(collated)
+ output["batch_modality"] = modality
+ return output
+
+
+__all__ = ["AstroPRISMJointCollator"]
diff --git a/src/data/astroprism_joint_dataset.py b/src/data/astroprism_joint_dataset.py
new file mode 100644
index 0000000..46eed76
--- /dev/null
+++ b/src/data/astroprism_joint_dataset.py
@@ -0,0 +1,185 @@
+"""Map-style union of the reviewed AstroPRISM image and spectrum datasets.
+
+The two source datasets retain their own validation and preprocessing
+contracts. This wrapper only provides a stable joint index space and tags
+each sample with the modality that its batch must contain. Homogeneous
+batches are enforced by :mod:`astroprism_joint_batch_sampler` and rechecked by
+:mod:`astroprism_joint_collate`.
+"""
+
+from __future__ import annotations
+
+import operator
+from dataclasses import dataclass, replace
+from typing import Any
+
+from torch.utils.data import Dataset
+
+from .astroprism_stage1a_dataset import (
+ AstroPRISMStage1ADataset,
+ AstroPRISMStage1ADatasetConfig,
+)
+from .astroprism_stage1b_spectrum_dataset import (
+ AstroPRISMStage1BSpectrumDataset,
+ AstroPRISMStage1BSpectrumDatasetConfig,
+)
+
+ASTROPRISM_JOINT_MODALITIES = ("image", "spectrum")
+ASTROPRISM_JOINT_SPLITS = frozenset({"train", "validation", "test"})
+
+
+class AstroPRISMJointDatasetError(ValueError):
+ """Raised when the two source datasets cannot form one joint view."""
+
+
+@dataclass(frozen=True)
+class AstroPRISMJointDatasetConfig:
+ """Child dataset configs plus the single split exposed by the joint view.
+
+ ``split`` is authoritative. :meth:`AstroPRISMJointDataset.from_config`
+ copies it into both child configs, preventing a caller from accidentally
+ pairing train examples from one modality with validation examples from the
+ other.
+ """
+
+ image: AstroPRISMStage1ADatasetConfig
+ spectrum: AstroPRISMStage1BSpectrumDatasetConfig
+ split: str = "train"
+
+
+class AstroPRISMJointDataset(Dataset):
+ """Joint index space over image-caption and spectrum-caption examples.
+
+ Image examples occupy ``[0, len(image_dataset))`` and spectrum examples
+ follow them. A shallow copy of each source sample is returned with a
+ ``batch_modality`` tag; source dataset samples are never mutated.
+ """
+
+ active_modalities = {"image", "spectrum"}
+
+ def __init__(
+ self,
+ image_dataset: Dataset,
+ spectrum_dataset: Dataset,
+ *,
+ split: str,
+ ) -> None:
+ split = str(split).strip()
+ if split not in ASTROPRISM_JOINT_SPLITS:
+ raise AstroPRISMJointDatasetError(
+ f"joint split must be one of {sorted(ASTROPRISM_JOINT_SPLITS)}, got {split!r}"
+ )
+
+ self.split = split
+ self.image_dataset = image_dataset
+ self.spectrum_dataset = spectrum_dataset
+ self._validate_child("image", image_dataset)
+ self._validate_child("spectrum", spectrum_dataset)
+
+ image_count = len(image_dataset)
+ spectrum_count = len(spectrum_dataset)
+ if image_count <= 0 or spectrum_count <= 0:
+ raise AstroPRISMJointDatasetError(
+ "joint dataset requires at least one image and one spectrum example; "
+ f"got image={image_count}, spectrum={spectrum_count}"
+ )
+
+ self._modality_indices = {
+ "image": range(0, image_count),
+ "spectrum": range(image_count, image_count + spectrum_count),
+ }
+
+ @classmethod
+ def from_config(
+ cls,
+ config: AstroPRISMJointDatasetConfig,
+ tokenizer: Any,
+ ) -> AstroPRISMJointDataset:
+ """Build both validated child datasets for ``config.split``."""
+
+ split = str(config.split).strip()
+ image_config = replace(config.image, split=split)
+ spectrum_config = replace(config.spectrum, split=split)
+ return cls(
+ AstroPRISMStage1ADataset(image_config, tokenizer),
+ AstroPRISMStage1BSpectrumDataset(spectrum_config, tokenizer),
+ split=split,
+ )
+
+ def _validate_child(self, modality: str, dataset: Dataset) -> None:
+ active_modalities = getattr(dataset, "active_modalities", None)
+ if set(active_modalities or ()) != {modality}:
+ raise AstroPRISMJointDatasetError(
+ f"{modality} child must expose active_modalities={{{modality!r}}}, "
+ f"got {active_modalities!r}"
+ )
+
+ child_config = getattr(dataset, "config", None)
+ child_split = getattr(child_config, "split", None)
+ if child_split != self.split:
+ raise AstroPRISMJointDatasetError(
+ f"{modality} child split must be {self.split!r}, got {child_split!r}"
+ )
+
+ @property
+ def modality_indices(self) -> dict[str, range]:
+ """Return immutable ranges in the joint index space by modality."""
+
+ return dict(self._modality_indices)
+
+ @property
+ def modality_lengths(self) -> dict[str, int]:
+ return {
+ modality: len(indices)
+ for modality, indices in self._modality_indices.items()
+ }
+
+ def modality_for_index(self, item: int) -> str:
+ index = self._normalize_index(item)
+ return "image" if index < len(self.image_dataset) else "spectrum"
+
+ def _normalize_index(self, item: int) -> int:
+ try:
+ index = operator.index(item)
+ except TypeError as exc:
+ raise TypeError(f"joint dataset index must be an integer, got {type(item).__name__}") from exc
+ if index < 0:
+ index += len(self)
+ if index < 0 or index >= len(self):
+ raise IndexError(f"joint dataset index {item!r} is out of range")
+ return index
+
+ def __len__(self) -> int:
+ return len(self.image_dataset) + len(self.spectrum_dataset)
+
+ def __getitem__(self, item: int) -> dict[str, Any]:
+ index = self._normalize_index(item)
+ image_count = len(self.image_dataset)
+ if index < image_count:
+ modality = "image"
+ sample = self.image_dataset[index]
+ else:
+ modality = "spectrum"
+ sample = self.spectrum_dataset[index - image_count]
+
+ if not isinstance(sample, dict):
+ raise AstroPRISMJointDatasetError(
+ f"{modality} child returned {type(sample).__name__}; expected a dict"
+ )
+ existing_tag = sample.get("batch_modality")
+ if existing_tag is not None and existing_tag != modality:
+ raise AstroPRISMJointDatasetError(
+ f"{modality} child returned conflicting batch_modality={existing_tag!r}"
+ )
+ tagged = dict(sample)
+ tagged["batch_modality"] = modality
+ return tagged
+
+
+__all__ = [
+ "ASTROPRISM_JOINT_MODALITIES",
+ "ASTROPRISM_JOINT_SPLITS",
+ "AstroPRISMJointDataset",
+ "AstroPRISMJointDatasetConfig",
+ "AstroPRISMJointDatasetError",
+]
diff --git a/src/data/astroprism_raw_images.py b/src/data/astroprism_raw_images.py
new file mode 100644
index 0000000..b8b816c
--- /dev/null
+++ b/src/data/astroprism_raw_images.py
@@ -0,0 +1,897 @@
+"""Select the pinned AstroPRISM raw Legacy Survey image rows.
+
+The official AION tutorial HDF5 is about 121 GB, but its datasets are
+contiguous and the server supports HTTP byte ranges. This module maps the
+committed candidate roster through the exact PROVABGS identifier column and
+copies only the selected image rows into mmap-friendly NPY files. Spectrum
+datasets are intentionally never read or written.
+"""
+
+from __future__ import annotations
+
+import csv
+import gzip
+import hashlib
+import json
+import math
+import os
+import re
+import shutil
+import tempfile
+import time
+import urllib.error
+import urllib.request
+import uuid
+from collections.abc import Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, TextIO
+from urllib.parse import urlsplit
+
+import fsspec
+import h5py
+import numpy as np
+
+RAW_SOURCE_URL = "https://users.flatironinstitute.org/~polymathic/data/provabgs_desi_ls.hdf5"
+RAW_SOURCE_ETAG = '"6a62487d-1c2e3b0ad4"'
+RAW_SOURCE_LENGTH = 121_034_705_620
+ROSTER_SHA256 = "650935faa58953c9929e7379429bfb332a7486dab54544fedef64ca385abb3bc"
+EXPECTED_CANDIDATES = 646
+
+_PROVABGS_GROUP = "provabgs_desi_edr_sv3"
+_EXPECTED_SOURCE_BANDS = ("des-g", "des-r", "des-i", "des-z")
+_AION_BAND_MAPPING = {
+ "des-g": "DES-G",
+ "des-r": "DES-R",
+ "des-i": "DES-I",
+ "des-z": "DES-Z",
+}
+_REQUIRED_ROSTER_COLUMNS = {
+ "paired_targetid",
+ "caption_object_id",
+ "paired_ra",
+ "paired_dec",
+ "caption_ra",
+ "caption_dec",
+}
+
+
+@dataclass(frozen=True)
+class _DatasetSpec:
+ source_path: str
+ filename: str
+ dtype: np.dtype[Any]
+ tail_shape: tuple[int, ...]
+ finite: bool = False
+
+
+_IMAGE_DATASETS = {
+ "image_array": _DatasetSpec(
+ "legacysurvey/image_array", "image_array.npy", np.dtype("float32"), (4, 160, 160), True
+ ),
+ "image_band": _DatasetSpec("legacysurvey/image_band", "image_band.npy", np.dtype("S5"), (4,)),
+ "image_ivar": _DatasetSpec(
+ "legacysurvey/image_ivar", "image_ivar.npy", np.dtype("float32"), (4, 160, 160), True
+ ),
+ "image_mask": _DatasetSpec(
+ "legacysurvey/image_mask", "image_mask.npy", np.dtype("bool"), (160, 160)
+ ),
+ "image_scale": _DatasetSpec(
+ "legacysurvey/image_scale", "image_scale.npy", np.dtype("float32"), (4,), True
+ ),
+ "object_mask": _DatasetSpec(
+ "legacysurvey/object_mask", "object_mask.npy", np.dtype("uint8"), (160, 160)
+ ),
+}
+
+_METADATA_DATASETS = {
+ "provabgs_id": (f"{_PROVABGS_GROUP}/object_id", np.dtype("int64")),
+ "provabgs_ra": (f"{_PROVABGS_GROUP}/ra", np.dtype("float64")),
+ "provabgs_dec": (f"{_PROVABGS_GROUP}/dec", np.dtype("float64")),
+ "desi_id": ("desi/object_id", np.dtype("int64")),
+ "desi_ra": ("desi/ra", np.dtype("float64")),
+ "desi_dec": ("desi/dec", np.dtype("float64")),
+ "legacy_id": ("legacysurvey/object_id", np.dtype("S16")),
+ "legacy_ra": ("legacysurvey/ra", np.dtype("float64")),
+ "legacy_dec": ("legacysurvey/dec", np.dtype("float64")),
+}
+
+
+class RawImageValidationError(ValueError):
+ """Raised when a source violates the pinned raw-image contract."""
+
+
+@dataclass(frozen=True)
+class RawImageConfig:
+ """Inputs and validation pins for one selective raw-image extraction."""
+
+ roster_path: Path
+ output_dir: Path
+ source: str | Path = RAW_SOURCE_URL
+ overwrite: bool = False
+ expected_count: int = EXPECTED_CANDIDATES
+ expected_roster_sha256: str | None = ROSTER_SHA256
+ expected_source_length: int | None = RAW_SOURCE_LENGTH
+ expected_source_etag: str | None = RAW_SOURCE_ETAG
+ http_block_size: int = 256 * 1024
+ http_timeout_seconds: float = 30.0
+ http_probe_attempts: int = 3
+ http_probe_backoff_seconds: float = 1.0
+ roster_coordinate_tolerance_arcsec: float = 0.1
+ desi_coordinate_tolerance_arcsec: float = 1e-4
+ legacy_coordinate_tolerance_arcsec: float = 1.0
+ caption_legacy_coordinate_tolerance_arcsec: float = 1e-4
+
+
+@dataclass(frozen=True)
+class _SourceInfo:
+ kind: str
+ location: str
+ length_bytes: int
+ etag: str | None
+ range_verified: bool
+ status_code: int | None
+ content_range: str | None
+ local_mtime_ns: int | None = None
+
+
+@dataclass(frozen=True)
+class _Run:
+ start: int
+ stop: int
+ candidate_indices: tuple[int, ...]
+
+
+def sha256_file(path: Path, chunk_size: int = 8 * 1024 * 1024) -> str:
+ """Return the SHA-256 digest of *path* without loading it into memory."""
+
+ digest = hashlib.sha256()
+ with path.open("rb") as handle:
+ while chunk := handle.read(chunk_size):
+ digest.update(chunk)
+ return digest.hexdigest()
+
+
+def _is_http_source(source: str | Path) -> bool:
+ return urlsplit(str(source)).scheme.lower() in {"http", "https"}
+
+
+def _open_roster(path: Path) -> TextIO:
+ if path.suffix == ".gz":
+ return gzip.open(path, "rt", encoding="utf-8", newline="")
+ return path.open("r", encoding="utf-8", newline="")
+
+
+def _parse_roster_float(value: str, *, field: str, index: int) -> float:
+ try:
+ parsed = float(value)
+ except (TypeError, ValueError) as exc:
+ raise RawImageValidationError(
+ f"roster row {index}: {field} must be a finite float, got {value!r}"
+ ) from exc
+ if not math.isfinite(parsed):
+ raise RawImageValidationError(f"roster row {index}: {field} is not finite")
+ return parsed
+
+
+def _parse_roster_id(value: str, *, index: int) -> int:
+ if not re.fullmatch(r"[0-9]+", value or ""):
+ raise RawImageValidationError(
+ f"roster row {index}: paired_targetid must be an unsigned decimal integer"
+ )
+ parsed = int(value)
+ if parsed > np.iinfo(np.int64).max:
+ raise RawImageValidationError(
+ f"roster row {index}: paired_targetid is outside signed int64"
+ )
+ return parsed
+
+
+def _read_roster(config: RawImageConfig) -> tuple[list[dict[str, str]], list[int], str]:
+ path = Path(config.roster_path)
+ if not path.is_file():
+ raise RawImageValidationError(f"roster does not exist: {path}")
+ digest = sha256_file(path)
+ if config.expected_roster_sha256 is not None and digest != config.expected_roster_sha256:
+ raise RawImageValidationError(
+ f"roster SHA-256 mismatch: expected {config.expected_roster_sha256}, got {digest}"
+ )
+
+ with _open_roster(path) as handle:
+ reader = csv.DictReader(handle)
+ fieldnames = reader.fieldnames or []
+ if len(fieldnames) != len(set(fieldnames)):
+ raise RawImageValidationError("roster contains duplicate column names")
+ missing = sorted(_REQUIRED_ROSTER_COLUMNS - set(fieldnames))
+ if missing:
+ raise RawImageValidationError(f"roster is missing required columns: {missing}")
+ rows = list(reader)
+
+ if len(rows) != config.expected_count:
+ raise RawImageValidationError(
+ f"roster must contain exactly {config.expected_count} rows, got {len(rows)}"
+ )
+
+ parsed_ids: list[int] = []
+ seen: set[int] = set()
+ caption_ids_seen: set[str] = set()
+ for index, row in enumerate(rows):
+ parsed_id = _parse_roster_id(row["paired_targetid"], index=index)
+ if parsed_id in seen:
+ raise RawImageValidationError(f"duplicate paired_targetid: {parsed_id}")
+ seen.add(parsed_id)
+ parsed_ids.append(parsed_id)
+
+ ra = _parse_roster_float(row["paired_ra"], field="paired_ra", index=index)
+ dec = _parse_roster_float(row["paired_dec"], field="paired_dec", index=index)
+ caption_id = row["caption_object_id"]
+ caption_ra = _parse_roster_float(row["caption_ra"], field="caption_ra", index=index)
+ caption_dec = _parse_roster_float(row["caption_dec"], field="caption_dec", index=index)
+ if not caption_id:
+ raise RawImageValidationError(
+ f"roster row {index}: caption_object_id must not be empty"
+ )
+ if caption_id in caption_ids_seen:
+ raise RawImageValidationError(f"duplicate caption_object_id: {caption_id!r}")
+ caption_ids_seen.add(caption_id)
+ if not 0.0 <= ra < 360.0 or not 0.0 <= caption_ra < 360.0:
+ raise RawImageValidationError(
+ f"roster row {index}: paired_ra/caption_ra is outside [0, 360)"
+ )
+ if not -90.0 <= dec <= 90.0 or not -90.0 <= caption_dec <= 90.0:
+ raise RawImageValidationError(
+ f"roster row {index}: paired_dec/caption_dec is outside [-90, 90]"
+ )
+ return rows, parsed_ids, digest
+
+
+def _probe_http_source(
+ url: str,
+ *,
+ expected_length: int,
+ expected_etag: str,
+ timeout_seconds: float,
+ attempts: int = 3,
+ backoff_seconds: float = 1.0,
+) -> _SourceInfo:
+ """Verify a one-byte HTTP range response and return pinned source metadata."""
+
+ if attempts < 1:
+ raise ValueError("attempts must be at least 1")
+ if not math.isfinite(backoff_seconds) or backoff_seconds < 0:
+ raise ValueError("backoff_seconds must be nonnegative")
+
+ for attempt in range(attempts):
+ request = urllib.request.Request(
+ url,
+ headers={"Range": "bytes=0-0", "If-Match": expected_etag},
+ method="GET",
+ )
+ try:
+ with urllib.request.urlopen(request, timeout=timeout_seconds) as response:
+ status = response.getcode()
+ headers = response.headers
+ body = response.read(2)
+ final_url = response.geturl()
+ except urllib.error.HTTPError as exc:
+ raise RawImageValidationError(f"HTTP range probe returned status {exc.code}") from exc
+ except OSError as exc:
+ if attempt + 1 == attempts:
+ raise RawImageValidationError(
+ f"HTTP range probe failed for {url} after {attempts} attempts: {exc}"
+ ) from exc
+ time.sleep(backoff_seconds * (2**attempt))
+ continue
+ break
+
+ if status != 206:
+ raise RawImageValidationError(f"HTTP range probe requires status 206, got {status}")
+ if len(body) != 1:
+ raise RawImageValidationError(
+ f"HTTP range probe expected exactly one response byte, got {len(body)}"
+ )
+ content_range = headers.get("Content-Range")
+ match = re.fullmatch(r"bytes 0-0/([0-9]+)", content_range or "")
+ if match is None:
+ raise RawImageValidationError(
+ f"HTTP range probe returned invalid Content-Range: {content_range!r}"
+ )
+ observed_length = int(match.group(1))
+ if observed_length != expected_length:
+ raise RawImageValidationError(
+ f"source length mismatch: expected {expected_length}, got {observed_length}"
+ )
+ content_length = headers.get("Content-Length")
+ if content_length is not None and content_length != "1":
+ raise RawImageValidationError(f"one-byte range returned Content-Length {content_length!r}")
+ observed_etag = headers.get("ETag")
+ if observed_etag != expected_etag:
+ raise RawImageValidationError(
+ f"source ETag mismatch: expected {expected_etag!r}, got {observed_etag!r}"
+ )
+ return _SourceInfo(
+ kind="http",
+ location=final_url,
+ length_bytes=observed_length,
+ etag=observed_etag,
+ range_verified=True,
+ status_code=status,
+ content_range=content_range,
+ )
+
+
+def _inspect_source(config: RawImageConfig) -> _SourceInfo:
+ if _is_http_source(config.source):
+ if config.expected_source_length is None or config.expected_source_etag is None:
+ raise ValueError("HTTP sources require expected_source_length and expected_source_etag")
+ return _probe_http_source(
+ str(config.source),
+ expected_length=config.expected_source_length,
+ expected_etag=config.expected_source_etag,
+ timeout_seconds=config.http_timeout_seconds,
+ attempts=config.http_probe_attempts,
+ backoff_seconds=config.http_probe_backoff_seconds,
+ )
+
+ path = Path(config.source).resolve()
+ if not path.is_file():
+ raise RawImageValidationError(f"local HDF5 source does not exist: {path}")
+ stat = path.stat()
+ if config.expected_source_length is not None and stat.st_size != config.expected_source_length:
+ raise RawImageValidationError(
+ f"source length mismatch: expected {config.expected_source_length}, got {stat.st_size}"
+ )
+ if config.expected_source_etag is not None:
+ raise ValueError("expected_source_etag must be None for a local HDF5 source")
+ return _SourceInfo(
+ kind="local",
+ location=str(path),
+ length_bytes=stat.st_size,
+ etag=None,
+ range_verified=False,
+ status_code=None,
+ content_range=None,
+ local_mtime_ns=stat.st_mtime_ns,
+ )
+
+
+@contextmanager
+def _open_hdf5(config: RawImageConfig) -> Iterator[h5py.File]:
+ if _is_http_source(config.source):
+ headers = {"If-Match": config.expected_source_etag}
+ with fsspec.open(
+ str(config.source),
+ mode="rb",
+ block_size=config.http_block_size,
+ cache_type="none",
+ headers=headers,
+ ) as remote:
+ with h5py.File(remote, "r") as handle:
+ yield handle
+ return
+
+ with h5py.File(Path(config.source), "r") as handle:
+ yield handle
+
+
+def _require_dataset(
+ handle: h5py.File,
+ path: str,
+ *,
+ shape: tuple[int, ...],
+ dtype: np.dtype[Any],
+ require_contiguous: bool,
+) -> h5py.Dataset:
+ if path not in handle or not isinstance(handle[path], h5py.Dataset):
+ raise RawImageValidationError(f"source is missing dataset {path!r}")
+ dataset = handle[path]
+ if dataset.shape != shape:
+ raise RawImageValidationError(f"{path} has shape {dataset.shape}, expected {shape}")
+ if dataset.dtype != dtype:
+ raise RawImageValidationError(f"{path} has dtype {dataset.dtype}, expected {dtype}")
+ if require_contiguous and (dataset.chunks is not None or dataset.compression is not None):
+ raise RawImageValidationError(
+ f"{path} must be contiguous and uncompressed for bounded HTTP range extraction"
+ )
+ return dataset
+
+
+def _validate_schema(handle: h5py.File, *, require_contiguous: bool) -> int:
+ id_path = _METADATA_DATASETS["provabgs_id"][0]
+ if id_path not in handle or not isinstance(handle[id_path], h5py.Dataset):
+ raise RawImageValidationError(f"source is missing dataset {id_path!r}")
+ id_dataset = handle[id_path]
+ if id_dataset.ndim != 1:
+ raise RawImageValidationError(f"{id_path} must be one-dimensional")
+ source_count = id_dataset.shape[0]
+
+ for path, dtype in _METADATA_DATASETS.values():
+ _require_dataset(
+ handle,
+ path,
+ shape=(source_count,),
+ dtype=dtype,
+ require_contiguous=require_contiguous,
+ )
+ for spec in _IMAGE_DATASETS.values():
+ _require_dataset(
+ handle,
+ spec.source_path,
+ shape=(source_count, *spec.tail_shape),
+ dtype=spec.dtype,
+ require_contiguous=require_contiguous,
+ )
+ return source_count
+
+
+def _build_raw_row_map(handle: h5py.File, roster_ids: list[int]) -> tuple[list[int], np.ndarray]:
+ source_ids = np.asarray(handle[_METADATA_DATASETS["provabgs_id"][0]][:])
+ unique_ids, counts = np.unique(source_ids, return_counts=True)
+ duplicate_ids = unique_ids[counts != 1]
+ if duplicate_ids.size:
+ raise RawImageValidationError(
+ f"PROVABGS object_id is not unique; first duplicate is {int(duplicate_ids[0])}"
+ )
+ lookup = {int(value): index for index, value in enumerate(source_ids)}
+ missing = [value for value in roster_ids if value not in lookup]
+ if missing:
+ raise RawImageValidationError(
+ f"{len(missing)} roster PROVABGS IDs are absent; first missing ID is {missing[0]}"
+ )
+ raw_rows = [lookup[value] for value in roster_ids]
+ if len(set(raw_rows)) != len(raw_rows):
+ raise RawImageValidationError("multiple roster IDs map to the same raw HDF5 row")
+ return raw_rows, source_ids
+
+
+def _contiguous_runs(raw_rows: list[int]) -> list[_Run]:
+ ordered = sorted((raw_row, candidate_index) for candidate_index, raw_row in enumerate(raw_rows))
+ if not ordered:
+ return []
+ runs: list[_Run] = []
+ run_start = ordered[0][0]
+ previous = ordered[0][0]
+ indices = [ordered[0][1]]
+ for raw_row, candidate_index in ordered[1:]:
+ if raw_row == previous + 1:
+ previous = raw_row
+ indices.append(candidate_index)
+ continue
+ runs.append(_Run(run_start, previous + 1, tuple(indices)))
+ run_start = previous = raw_row
+ indices = [candidate_index]
+ runs.append(_Run(run_start, previous + 1, tuple(indices)))
+ return runs
+
+
+def _coordinate(ra: Any, dec: Any, *, label: str, index: int) -> tuple[float, float]:
+ ra_float, dec_float = float(ra), float(dec)
+ if not math.isfinite(ra_float) or not 0.0 <= ra_float < 360.0:
+ raise RawImageValidationError(f"candidate {index}: invalid {label} RA {ra_float}")
+ if not math.isfinite(dec_float) or not -90.0 <= dec_float <= 90.0:
+ raise RawImageValidationError(f"candidate {index}: invalid {label} Dec {dec_float}")
+ return ra_float, dec_float
+
+
+def _angular_separation_arcsec(ra1: float, dec1: float, ra2: float, dec2: float) -> float:
+ ra1_rad, dec1_rad, ra2_rad, dec2_rad = map(math.radians, (ra1, dec1, ra2, dec2))
+ sin_dec = math.sin((dec2_rad - dec1_rad) / 2.0)
+ sin_ra = math.sin((ra2_rad - ra1_rad) / 2.0)
+ haversine = sin_dec**2 + math.cos(dec1_rad) * math.cos(dec2_rad) * sin_ra**2
+ haversine = min(1.0, max(0.0, haversine))
+ return math.degrees(2.0 * math.asin(math.sqrt(haversine))) * 3600.0
+
+
+def _decode_legacy_id(value: np.bytes_, *, index: int) -> str:
+ try:
+ decoded = bytes(value).rstrip(b"\x00").decode("ascii")
+ except UnicodeDecodeError as exc:
+ raise RawImageValidationError(
+ f"candidate {index}: Legacy Survey object_id is not ASCII"
+ ) from exc
+ if not decoded:
+ raise RawImageValidationError(f"candidate {index}: Legacy Survey object_id is empty")
+ return decoded
+
+
+def _crosscheck_metadata(
+ config: RawImageConfig,
+ handle: h5py.File,
+ rows: list[dict[str, str]],
+ roster_ids: list[int],
+ raw_rows: list[int],
+ source_ids: np.ndarray,
+) -> list[dict[str, Any]]:
+ values: dict[str, np.ndarray] = {}
+ selected_indices = np.asarray(raw_rows, dtype=np.intp)
+ for key, (path, _) in _METADATA_DATASETS.items():
+ if key == "provabgs_id":
+ values[key] = source_ids[selected_indices]
+ else:
+ full_column = np.asarray(handle[path][:])
+ values[key] = full_column[selected_indices]
+
+ records: list[dict[str, Any]] = []
+ legacy_ids_seen: set[str] = set()
+ for index, (row, roster_id, raw_row) in enumerate(zip(rows, roster_ids, raw_rows, strict=True)):
+ provabgs_id = int(values["provabgs_id"][index])
+ desi_id = int(values["desi_id"][index])
+ legacy_id = _decode_legacy_id(values["legacy_id"][index], index=index)
+ if provabgs_id != roster_id:
+ raise RawImageValidationError(
+ f"candidate {index}: PROVABGS ID mismatch; expected {roster_id}, got {provabgs_id}"
+ )
+ if desi_id != provabgs_id:
+ raise RawImageValidationError(
+ f"candidate {index}: DESI ID {desi_id} does not match PROVABGS ID {provabgs_id}"
+ )
+ if legacy_id in legacy_ids_seen:
+ raise RawImageValidationError(
+ f"candidate {index}: duplicate Legacy Survey object_id {legacy_id!r}"
+ )
+ legacy_ids_seen.add(legacy_id)
+ caption_id = row["caption_object_id"]
+ if legacy_id != caption_id:
+ raise RawImageValidationError(
+ f"candidate {index}: Legacy Survey ID {legacy_id!r} does not match "
+ f"caption_object_id {caption_id!r}"
+ )
+
+ provabgs_coord = _coordinate(
+ values["provabgs_ra"][index],
+ values["provabgs_dec"][index],
+ label="PROVABGS",
+ index=index,
+ )
+ desi_coord = _coordinate(
+ values["desi_ra"][index],
+ values["desi_dec"][index],
+ label="DESI",
+ index=index,
+ )
+ legacy_coord = _coordinate(
+ values["legacy_ra"][index],
+ values["legacy_dec"][index],
+ label="Legacy Survey",
+ index=index,
+ )
+ roster_coord = (
+ _parse_roster_float(row["paired_ra"], field="paired_ra", index=index),
+ _parse_roster_float(row["paired_dec"], field="paired_dec", index=index),
+ )
+ caption_coord = (
+ _parse_roster_float(row["caption_ra"], field="caption_ra", index=index),
+ _parse_roster_float(row["caption_dec"], field="caption_dec", index=index),
+ )
+ roster_separation = _angular_separation_arcsec(*roster_coord, *provabgs_coord)
+ desi_separation = _angular_separation_arcsec(*provabgs_coord, *desi_coord)
+ legacy_separation = _angular_separation_arcsec(*provabgs_coord, *legacy_coord)
+ caption_legacy_separation = _angular_separation_arcsec(*caption_coord, *legacy_coord)
+ if roster_separation > config.roster_coordinate_tolerance_arcsec:
+ raise RawImageValidationError(
+ f"candidate {index}: roster/PROVABGS coordinate separation "
+ f"{roster_separation:.6g} arcsec exceeds "
+ f"{config.roster_coordinate_tolerance_arcsec}"
+ )
+ if desi_separation > config.desi_coordinate_tolerance_arcsec:
+ raise RawImageValidationError(
+ f"candidate {index}: PROVABGS/DESI coordinate separation "
+ f"{desi_separation:.6g} arcsec exceeds "
+ f"{config.desi_coordinate_tolerance_arcsec}"
+ )
+ if legacy_separation > config.legacy_coordinate_tolerance_arcsec:
+ raise RawImageValidationError(
+ f"candidate {index}: PROVABGS/Legacy Survey coordinate separation "
+ f"{legacy_separation:.6g} arcsec exceeds "
+ f"{config.legacy_coordinate_tolerance_arcsec}"
+ )
+ if caption_legacy_separation > config.caption_legacy_coordinate_tolerance_arcsec:
+ raise RawImageValidationError(
+ f"candidate {index}: caption/Legacy Survey coordinate separation "
+ f"{caption_legacy_separation:.6g} arcsec exceeds "
+ f"{config.caption_legacy_coordinate_tolerance_arcsec}"
+ )
+
+ records.append(
+ {
+ "candidate_index": index,
+ "paired_targetid": str(roster_id),
+ "caption_object_id": caption_id,
+ "raw_row": raw_row,
+ "source_ids": {
+ "provabgs": str(provabgs_id),
+ "desi": str(desi_id),
+ "legacysurvey": legacy_id,
+ },
+ "coordinates_deg": {
+ "roster": {"ra": roster_coord[0], "dec": roster_coord[1]},
+ "caption": {"ra": caption_coord[0], "dec": caption_coord[1]},
+ "provabgs": {"ra": provabgs_coord[0], "dec": provabgs_coord[1]},
+ "desi": {"ra": desi_coord[0], "dec": desi_coord[1]},
+ "legacysurvey": {"ra": legacy_coord[0], "dec": legacy_coord[1]},
+ },
+ "coordinate_crosscheck_arcsec": {
+ "roster_to_provabgs": roster_separation,
+ "provabgs_to_desi": desi_separation,
+ "provabgs_to_legacysurvey": legacy_separation,
+ "caption_to_legacysurvey": caption_legacy_separation,
+ },
+ "acceptance_status": "pending_review",
+ "split": None,
+ }
+ )
+ return records
+
+
+def _map_aion_bands(block: np.ndarray, candidate_indices: tuple[int, ...]) -> np.ndarray:
+ mapped = np.empty(block.shape, dtype="S5")
+ for local_index, candidate_index in enumerate(candidate_indices):
+ source_names: list[str] = []
+ for value in block[local_index]:
+ try:
+ source_name = bytes(value).rstrip(b"\x00").decode("ascii")
+ except UnicodeDecodeError as exc:
+ raise RawImageValidationError(
+ f"candidate {candidate_index}: image_band contains non-ASCII bytes"
+ ) from exc
+ if source_name not in _AION_BAND_MAPPING:
+ raise RawImageValidationError(
+ f"candidate {candidate_index}: unsupported image band {source_name!r}"
+ )
+ source_names.append(source_name)
+ if tuple(source_names) != _EXPECTED_SOURCE_BANDS:
+ raise RawImageValidationError(
+ f"candidate {candidate_index}: expected image bands in strict "
+ f"des-g/des-r/des-i/des-z order, "
+ f"got {source_names}"
+ )
+ mapped[local_index] = np.asarray(
+ [_AION_BAND_MAPPING[name].encode("ascii") for name in source_names], dtype="S5"
+ )
+ return mapped
+
+
+def _extract_datasets(
+ handle: h5py.File,
+ runs: list[_Run],
+ candidate_count: int,
+ staging: Path,
+) -> tuple[dict[str, dict[str, Any]], np.ndarray, np.ndarray]:
+ outputs: dict[str, dict[str, Any]] = {}
+ source_bands: np.ndarray | None = None
+ aion_bands: np.ndarray | None = None
+
+ for key, spec in _IMAGE_DATASETS.items():
+ output_path = staging / spec.filename
+ output = np.lib.format.open_memmap(
+ output_path,
+ mode="w+",
+ dtype=spec.dtype,
+ shape=(candidate_count, *spec.tail_shape),
+ )
+ mapped_output: np.memmap | None = None
+ mapped_path = staging / "aion_image_band.npy"
+ if key == "image_band":
+ mapped_output = np.lib.format.open_memmap(
+ mapped_path,
+ mode="w+",
+ dtype=np.dtype("S5"),
+ shape=(candidate_count, 4),
+ )
+ try:
+ dataset = handle[spec.source_path]
+ for run in runs:
+ block = np.asarray(dataset[run.start : run.stop])
+ if block.shape != (len(run.candidate_indices), *spec.tail_shape):
+ raise RawImageValidationError(
+ f"short or malformed read from {spec.source_path}: got {block.shape}"
+ )
+ if spec.finite and not np.isfinite(block).all():
+ raise RawImageValidationError(
+ f"{spec.source_path} contains non-finite values in raw rows "
+ f"{run.start}:{run.stop}"
+ )
+ destination = np.asarray(run.candidate_indices, dtype=np.intp)
+ output[destination] = block
+ if mapped_output is not None:
+ mapped_output[destination] = _map_aion_bands(block, run.candidate_indices)
+ output.flush()
+ if mapped_output is not None:
+ mapped_output.flush()
+ finally:
+ del output
+ if mapped_output is not None:
+ del mapped_output
+
+ outputs[key] = {
+ "path": spec.filename,
+ "source_dataset": spec.source_path,
+ "dtype": str(spec.dtype),
+ "shape": [candidate_count, *spec.tail_shape],
+ "size_bytes": output_path.stat().st_size,
+ "sha256": sha256_file(output_path),
+ "mmap_safe": True,
+ }
+ if key == "image_band":
+ outputs["aion_image_band"] = {
+ "path": mapped_path.name,
+ "derived_from": spec.filename,
+ "dtype": str(np.dtype("S5")),
+ "shape": [candidate_count, 4],
+ "size_bytes": mapped_path.stat().st_size,
+ "sha256": sha256_file(mapped_path),
+ "mmap_safe": True,
+ }
+ source_bands = np.load(output_path, mmap_mode="r")
+ aion_bands = np.load(mapped_path, mmap_mode="r")
+
+ assert source_bands is not None and aion_bands is not None
+ # Materialize these tiny arrays before the temporary mmap objects are closed.
+ return outputs, np.asarray(source_bands).copy(), np.asarray(aion_bands).copy()
+
+
+def _path_exists(path: Path) -> bool:
+ return path.exists() or path.is_symlink()
+
+
+def _commit_staging(staging: Path, output: Path, *, overwrite: bool) -> None:
+ if not _path_exists(output):
+ os.replace(staging, output)
+ return
+ if not overwrite:
+ raise FileExistsError(f"output already exists: {output}")
+
+ backup = output.with_name(f".{output.name}.backup-{uuid.uuid4().hex}")
+ os.replace(output, backup)
+ try:
+ os.replace(staging, output)
+ except BaseException:
+ os.replace(backup, output)
+ raise
+ if backup.is_dir() and not backup.is_symlink():
+ shutil.rmtree(backup)
+ else:
+ backup.unlink(missing_ok=True)
+
+
+def _verify_source_unchanged(config: RawImageConfig, original: _SourceInfo) -> None:
+ if original.kind == "http":
+ current = _inspect_source(config)
+ if current.length_bytes != original.length_bytes or current.etag != original.etag:
+ raise RawImageValidationError("remote source changed during extraction")
+ return
+ stat = Path(original.location).stat()
+ if stat.st_size != original.length_bytes or stat.st_mtime_ns != original.local_mtime_ns:
+ raise RawImageValidationError("local source changed during extraction")
+
+
+def _decode_band_row(row: np.ndarray) -> list[str]:
+ return [bytes(value).rstrip(b"\x00").decode("ascii") for value in row]
+
+
+def extract_raw_images(config: RawImageConfig) -> dict[str, Any]:
+ """Extract the roster's raw image rows into an atomic mmap-safe artifact."""
+
+ if config.http_block_size < 1:
+ raise ValueError("http_block_size must be positive")
+ if config.http_timeout_seconds <= 0:
+ raise ValueError("http_timeout_seconds must be positive")
+ if config.http_probe_attempts < 1:
+ raise ValueError("http_probe_attempts must be at least 1")
+ if (
+ not math.isfinite(config.http_probe_backoff_seconds)
+ or config.http_probe_backoff_seconds < 0
+ ):
+ raise ValueError("http_probe_backoff_seconds must be nonnegative")
+ for name in (
+ "roster_coordinate_tolerance_arcsec",
+ "desi_coordinate_tolerance_arcsec",
+ "legacy_coordinate_tolerance_arcsec",
+ "caption_legacy_coordinate_tolerance_arcsec",
+ ):
+ if getattr(config, name) < 0:
+ raise ValueError(f"{name} must be nonnegative")
+
+ output_dir = Path(config.output_dir)
+ if _path_exists(output_dir) and not config.overwrite:
+ raise FileExistsError(f"output already exists: {output_dir}")
+
+ rows, roster_ids, roster_sha256 = _read_roster(config)
+ source_info = _inspect_source(config)
+ staging: Path | None = None
+ try:
+ with _open_hdf5(config) as handle:
+ source_count = _validate_schema(handle, require_contiguous=source_info.kind == "http")
+ raw_rows, source_ids = _build_raw_row_map(handle, roster_ids)
+ runs = _contiguous_runs(raw_rows)
+ records = _crosscheck_metadata(
+ config,
+ handle,
+ rows,
+ roster_ids,
+ raw_rows,
+ source_ids,
+ )
+
+ output_dir.parent.mkdir(parents=True, exist_ok=True)
+ staging = Path(
+ tempfile.mkdtemp(prefix=f".{output_dir.name}.tmp-", dir=output_dir.parent)
+ )
+ outputs, source_bands, aion_bands = _extract_datasets(handle, runs, len(rows), staging)
+
+ _verify_source_unchanged(config, source_info)
+ for index, record in enumerate(records):
+ record["bands"] = {
+ "source": _decode_band_row(source_bands[index]),
+ "aion": _decode_band_row(aion_bands[index]),
+ }
+
+ manifest = {
+ "schema_version": 1,
+ "artifact_type": "astroprism_raw_images",
+ "created_at_utc": datetime.now(timezone.utc).isoformat(),
+ "roster": {
+ "path": str(Path(config.roster_path)),
+ "sha256": roster_sha256,
+ "row_count": len(rows),
+ },
+ "source": {
+ "kind": source_info.kind,
+ "location": source_info.location,
+ "url": source_info.location if source_info.kind == "http" else None,
+ "etag": source_info.etag,
+ "length_bytes": source_info.length_bytes,
+ "http_range_verified": source_info.range_verified,
+ "http_status": source_info.status_code,
+ "content_range": source_info.content_range,
+ "hdf5_row_count": source_count,
+ },
+ "selection": {
+ "candidate_count": len(rows),
+ "roster_order_preserved": True,
+ "raw_rows_in_roster_order": raw_rows,
+ "raw_row_min": min(raw_rows),
+ "raw_row_max": max(raw_rows),
+ "contiguous_run_count": len(runs),
+ "http_cache_type": "none" if source_info.kind == "http" else None,
+ "acceptance_status": "pending_review",
+ "split": None,
+ },
+ "identity_tolerances_arcsec": {
+ "roster_to_provabgs": config.roster_coordinate_tolerance_arcsec,
+ "provabgs_to_desi": config.desi_coordinate_tolerance_arcsec,
+ "provabgs_to_legacysurvey": config.legacy_coordinate_tolerance_arcsec,
+ "caption_to_legacysurvey": (config.caption_legacy_coordinate_tolerance_arcsec),
+ },
+ "bands": {
+ "source_labels_preserved": True,
+ "aion_mapping": _AION_BAND_MAPPING,
+ "aion_mapping_verified_for_every_candidate": True,
+ },
+ "outputs": outputs,
+ "candidates": records,
+ "spectra_included": False,
+ }
+ assert staging is not None
+ manifest_path = staging / "manifest.json"
+ manifest_path.write_text(
+ json.dumps(manifest, indent=2, sort_keys=True, allow_nan=False) + "\n",
+ encoding="utf-8",
+ )
+ _commit_staging(staging, output_dir, overwrite=config.overwrite)
+ staging = None
+ return {
+ "status": "prepared",
+ "candidate_count": len(rows),
+ "source_length_bytes": source_info.length_bytes,
+ "source_etag": source_info.etag,
+ "contiguous_run_count": len(runs),
+ "output_dir": str(output_dir),
+ "manifest": str(output_dir / "manifest.json"),
+ "spectra_included": False,
+ }
+ finally:
+ if staging is not None and staging.exists():
+ shutil.rmtree(staging)
diff --git a/src/data/astroprism_raw_spectra.py b/src/data/astroprism_raw_spectra.py
new file mode 100644
index 0000000..ea39e3f
--- /dev/null
+++ b/src/data/astroprism_raw_spectra.py
@@ -0,0 +1,349 @@
+"""Select the pinned AstroPRISM raw DESI spectrum rows.
+
+The 646-object AstroPRISM roster was originally used to extract calibrated
+Legacy Survey cutouts from the official AION tutorial HDF5. This module uses
+the same pinned source, object identifiers, coordinate checks, and roster order
+to extract the corresponding observed-frame DESI spectra without downloading
+the 121 GB source in full.
+
+The four output arrays are mmap-friendly and preserve the inputs expected by
+``aion.modalities.DESISpectrum``: flux, inverse variance, mask, and wavelength.
+In AION's contract, ``True`` mask values are excluded from normalization.
+"""
+
+from __future__ import annotations
+
+import json
+import math
+import shutil
+import tempfile
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+from .astroprism_raw_images import (
+ EXPECTED_CANDIDATES,
+ RAW_SOURCE_ETAG,
+ RAW_SOURCE_LENGTH,
+ RAW_SOURCE_URL,
+ ROSTER_SHA256,
+ RawImageConfig,
+ RawImageValidationError,
+ _build_raw_row_map,
+ _commit_staging,
+ _contiguous_runs,
+ _crosscheck_metadata,
+ _inspect_source,
+ _open_hdf5,
+ _path_exists,
+ _read_roster,
+ _require_dataset,
+ _validate_schema,
+ _verify_source_unchanged,
+ sha256_file,
+)
+from .astroprism_spectrum_contract import (
+ ASTROPRISM_SPECTRUM_LENGTH,
+ ASTROPRISM_SPECTRUM_TOKENS,
+)
+
+
+class RawSpectrumValidationError(RawImageValidationError):
+ """Raised when a source violates the pinned raw-spectrum contract."""
+
+
+@dataclass(frozen=True)
+class _SpectrumDatasetSpec:
+ source_path: str
+ filename: str
+ dtype: np.dtype[Any]
+
+
+_SPECTRUM_DATASETS = {
+ "spectrum_flux": _SpectrumDatasetSpec(
+ "desi/spectrum_flux", "spectrum_flux.npy", np.dtype("float32")
+ ),
+ "spectrum_ivar": _SpectrumDatasetSpec(
+ "desi/spectrum_ivar", "spectrum_ivar.npy", np.dtype("float32")
+ ),
+ "spectrum_mask": _SpectrumDatasetSpec(
+ "desi/spectrum_mask", "spectrum_mask.npy", np.dtype("bool")
+ ),
+ "spectrum_wavelength": _SpectrumDatasetSpec(
+ "desi/spectrum_lambda", "spectrum_wavelength.npy", np.dtype("float32")
+ ),
+}
+
+
+@dataclass(frozen=True)
+class RawSpectrumConfig:
+ """Inputs and validation pins for one selective spectrum extraction."""
+
+ roster_path: Path
+ output_dir: Path
+ source: str | Path = RAW_SOURCE_URL
+ overwrite: bool = False
+ expected_count: int = EXPECTED_CANDIDATES
+ expected_roster_sha256: str | None = ROSTER_SHA256
+ expected_source_length: int | None = RAW_SOURCE_LENGTH
+ expected_source_etag: str | None = RAW_SOURCE_ETAG
+ http_block_size: int = 256 * 1024
+ http_timeout_seconds: float = 30.0
+ http_probe_attempts: int = 3
+ http_probe_backoff_seconds: float = 1.0
+ roster_coordinate_tolerance_arcsec: float = 0.1
+ desi_coordinate_tolerance_arcsec: float = 1e-4
+ legacy_coordinate_tolerance_arcsec: float = 1.0
+ caption_legacy_coordinate_tolerance_arcsec: float = 1e-4
+
+
+def _common_config(config: RawSpectrumConfig) -> RawImageConfig:
+ """Adapt to the shared pinned-source and identity-validation helpers."""
+
+ return RawImageConfig(
+ roster_path=config.roster_path,
+ output_dir=config.output_dir,
+ source=config.source,
+ overwrite=config.overwrite,
+ expected_count=config.expected_count,
+ expected_roster_sha256=config.expected_roster_sha256,
+ expected_source_length=config.expected_source_length,
+ expected_source_etag=config.expected_source_etag,
+ http_block_size=config.http_block_size,
+ http_timeout_seconds=config.http_timeout_seconds,
+ http_probe_attempts=config.http_probe_attempts,
+ http_probe_backoff_seconds=config.http_probe_backoff_seconds,
+ roster_coordinate_tolerance_arcsec=config.roster_coordinate_tolerance_arcsec,
+ desi_coordinate_tolerance_arcsec=config.desi_coordinate_tolerance_arcsec,
+ legacy_coordinate_tolerance_arcsec=config.legacy_coordinate_tolerance_arcsec,
+ caption_legacy_coordinate_tolerance_arcsec=(
+ config.caption_legacy_coordinate_tolerance_arcsec
+ ),
+ )
+
+
+def _validate_spectrum_schema(
+ handle: Any,
+ *,
+ source_count: int,
+ require_contiguous: bool,
+) -> None:
+ for spec in _SPECTRUM_DATASETS.values():
+ _require_dataset(
+ handle,
+ spec.source_path,
+ shape=(source_count, ASTROPRISM_SPECTRUM_LENGTH),
+ dtype=spec.dtype,
+ require_contiguous=require_contiguous,
+ )
+
+
+def _validate_spectrum_block(
+ key: str,
+ block: np.ndarray,
+ *,
+ raw_start: int,
+) -> None:
+ if key in {"spectrum_flux", "spectrum_ivar", "spectrum_wavelength"}:
+ if not np.isfinite(block).all():
+ raise RawSpectrumValidationError(
+ f"{_SPECTRUM_DATASETS[key].source_path} contains non-finite values "
+ f"in raw rows beginning at {raw_start}"
+ )
+ if key == "spectrum_ivar" and np.any(block < 0):
+ raise RawSpectrumValidationError(
+ f"desi/spectrum_ivar contains negative values in raw rows beginning at {raw_start}"
+ )
+ if key == "spectrum_wavelength" and np.any(np.diff(block, axis=-1) <= 0):
+ raise RawSpectrumValidationError(
+ "desi/spectrum_lambda must be strictly increasing for every selected spectrum"
+ )
+ if key == "spectrum_mask" and np.any(np.all(block, axis=-1)):
+ raise RawSpectrumValidationError(
+ "desi/spectrum_mask marks every wavelength as invalid for a selected spectrum"
+ )
+
+
+def _extract_spectrum_datasets(
+ handle: Any,
+ runs: list[Any],
+ candidate_count: int,
+ staging: Path,
+) -> dict[str, dict[str, Any]]:
+ outputs: dict[str, dict[str, Any]] = {}
+ for key, spec in _SPECTRUM_DATASETS.items():
+ output_path = staging / spec.filename
+ output = np.lib.format.open_memmap(
+ output_path,
+ mode="w+",
+ dtype=spec.dtype,
+ shape=(candidate_count, ASTROPRISM_SPECTRUM_LENGTH),
+ )
+ try:
+ source_dataset = handle[spec.source_path]
+ for run in runs:
+ block = np.asarray(source_dataset[run.start : run.stop])
+ expected_shape = (len(run.candidate_indices), ASTROPRISM_SPECTRUM_LENGTH)
+ if block.shape != expected_shape:
+ raise RawSpectrumValidationError(
+ f"short or malformed read from {spec.source_path}: "
+ f"got {block.shape}, expected {expected_shape}"
+ )
+ _validate_spectrum_block(key, block, raw_start=run.start)
+ destination = np.asarray(run.candidate_indices, dtype=np.intp)
+ output[destination] = block
+ output.flush()
+ finally:
+ del output
+
+ outputs[key] = {
+ "path": spec.filename,
+ "source_dataset": spec.source_path,
+ "dtype": str(spec.dtype),
+ "shape": [candidate_count, ASTROPRISM_SPECTRUM_LENGTH],
+ "size_bytes": output_path.stat().st_size,
+ "sha256": sha256_file(output_path),
+ "mmap_safe": True,
+ }
+ return outputs
+
+
+def extract_raw_spectra(config: RawSpectrumConfig) -> dict[str, Any]:
+ """Extract the roster's raw spectrum rows into an atomic artifact."""
+
+ if config.http_block_size < 1:
+ raise ValueError("http_block_size must be positive")
+ if config.http_timeout_seconds <= 0:
+ raise ValueError("http_timeout_seconds must be positive")
+ if config.http_probe_attempts < 1:
+ raise ValueError("http_probe_attempts must be at least 1")
+ if (
+ not math.isfinite(config.http_probe_backoff_seconds)
+ or config.http_probe_backoff_seconds < 0
+ ):
+ raise ValueError("http_probe_backoff_seconds must be nonnegative")
+ for name in (
+ "roster_coordinate_tolerance_arcsec",
+ "desi_coordinate_tolerance_arcsec",
+ "legacy_coordinate_tolerance_arcsec",
+ "caption_legacy_coordinate_tolerance_arcsec",
+ ):
+ if getattr(config, name) < 0:
+ raise ValueError(f"{name} must be nonnegative")
+
+ output_dir = Path(config.output_dir)
+ if _path_exists(output_dir) and not config.overwrite:
+ raise FileExistsError(f"output already exists: {output_dir}")
+
+ common = _common_config(config)
+ rows, roster_ids, roster_sha256 = _read_roster(common)
+ source_info = _inspect_source(common)
+ staging: Path | None = None
+ try:
+ with _open_hdf5(common) as handle:
+ source_count = _validate_schema(
+ handle,
+ require_contiguous=source_info.kind == "http",
+ )
+ _validate_spectrum_schema(
+ handle,
+ source_count=source_count,
+ require_contiguous=source_info.kind == "http",
+ )
+ raw_rows, source_ids = _build_raw_row_map(handle, roster_ids)
+ runs = _contiguous_runs(raw_rows)
+ records = _crosscheck_metadata(
+ common,
+ handle,
+ rows,
+ roster_ids,
+ raw_rows,
+ source_ids,
+ )
+
+ output_dir.parent.mkdir(parents=True, exist_ok=True)
+ staging = Path(
+ tempfile.mkdtemp(prefix=f".{output_dir.name}.tmp-", dir=output_dir.parent)
+ )
+ outputs = _extract_spectrum_datasets(handle, runs, len(rows), staging)
+
+ _verify_source_unchanged(common, source_info)
+ manifest = {
+ "schema_version": 1,
+ "artifact_type": "astroprism_raw_spectra",
+ "created_at_utc": datetime.now(timezone.utc).isoformat(),
+ "roster": {
+ "path": str(Path(config.roster_path)),
+ "sha256": roster_sha256,
+ "row_count": len(rows),
+ },
+ "source": {
+ "kind": source_info.kind,
+ "location": source_info.location,
+ "url": source_info.location if source_info.kind == "http" else None,
+ "etag": source_info.etag,
+ "length_bytes": source_info.length_bytes,
+ "http_range_verified": source_info.range_verified,
+ "http_status": source_info.status_code,
+ "content_range": source_info.content_range,
+ "hdf5_row_count": source_count,
+ },
+ "selection": {
+ "candidate_count": len(rows),
+ "roster_order_preserved": True,
+ "raw_rows_in_roster_order": raw_rows,
+ "raw_row_min": min(raw_rows),
+ "raw_row_max": max(raw_rows),
+ "contiguous_run_count": len(runs),
+ "http_cache_type": "none" if source_info.kind == "http" else None,
+ },
+ "identity_tolerances_arcsec": {
+ "roster_to_provabgs": config.roster_coordinate_tolerance_arcsec,
+ "provabgs_to_desi": config.desi_coordinate_tolerance_arcsec,
+ "provabgs_to_legacysurvey": config.legacy_coordinate_tolerance_arcsec,
+ "caption_to_legacysurvey": (
+ config.caption_legacy_coordinate_tolerance_arcsec
+ ),
+ },
+ "spectrum": {
+ "length": ASTROPRISM_SPECTRUM_LENGTH,
+ "observed_frame": True,
+ "mask_true_means_excluded": True,
+ "aion_modality": "DESISpectrum",
+ "aion_token_count": ASTROPRISM_SPECTRUM_TOKENS,
+ },
+ "outputs": outputs,
+ "candidates": records,
+ "images_included": False,
+ }
+ assert staging is not None
+ (staging / "manifest.json").write_text(
+ json.dumps(manifest, indent=2, sort_keys=True, allow_nan=False) + "\n",
+ encoding="utf-8",
+ )
+ _commit_staging(staging, output_dir, overwrite=config.overwrite)
+ staging = None
+ return {
+ "status": "prepared",
+ "candidate_count": len(rows),
+ "source_length_bytes": source_info.length_bytes,
+ "source_etag": source_info.etag,
+ "contiguous_run_count": len(runs),
+ "output_dir": str(output_dir),
+ "manifest": str(output_dir / "manifest.json"),
+ "spectrum_length": ASTROPRISM_SPECTRUM_LENGTH,
+ }
+ finally:
+ if staging is not None and staging.exists():
+ shutil.rmtree(staging)
+
+
+__all__ = [
+ "ASTROPRISM_SPECTRUM_LENGTH",
+ "RawSpectrumConfig",
+ "RawSpectrumValidationError",
+ "extract_raw_spectra",
+]
diff --git a/src/data/astroprism_spectrum_contract.py b/src/data/astroprism_spectrum_contract.py
new file mode 100644
index 0000000..8ebfd06
--- /dev/null
+++ b/src/data/astroprism_spectrum_contract.py
@@ -0,0 +1,13 @@
+"""Lightweight constants shared by AstroPRISM spectrum data paths.
+
+Keep this module free of extraction-only dependencies such as h5py so the
+training runtime can import the mmap dataset from the packed environment.
+"""
+
+ASTROPRISM_SPECTRUM_LENGTH = 7781
+ASTROPRISM_SPECTRUM_TOKENS = 273
+
+__all__ = [
+ "ASTROPRISM_SPECTRUM_LENGTH",
+ "ASTROPRISM_SPECTRUM_TOKENS",
+]
diff --git a/src/data/astroprism_spectrum_facts.py b/src/data/astroprism_spectrum_facts.py
new file mode 100644
index 0000000..da6b663
--- /dev/null
+++ b/src/data/astroprism_spectrum_facts.py
@@ -0,0 +1,651 @@
+"""Deterministic, audit-friendly facts for AstroPRISM DESI spectra.
+
+The fact extractor implements the intermediate stage used by the
+AstroSpecLM Figure 1 / Appendix A.5 recipe: numerical spectra are converted
+to a compact structured report before any language model writes a caption.
+
+This module intentionally does not generate prose. It measures a conservative
+set of galaxy features using a supplied catalog redshift, retaining that
+conditioning in every line record. The measurements are intended for caption
+candidate construction and expert review, not as a replacement for a survey
+pipeline or publication-quality line fitting.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Literal
+
+import numpy as np
+
+Nature = Literal["emission", "absorption"]
+
+
+@dataclass(frozen=True)
+class SpectralLineDefinition:
+ """One redshift-conditioned local line measurement."""
+
+ key: str
+ label: str
+ rest_wavelength_angstrom: float
+ nature: Nature
+ line_window_rest_angstrom: tuple[float, float]
+ continuum_windows_rest_angstrom: tuple[tuple[float, float], tuple[float, float]]
+ caveat: str | None = None
+
+
+# Vacuum wavelengths are used because DESI wavelengths are on a vacuum grid.
+# Blend reference wavelengths are descriptive centroids; the full integration
+# window, rather than the reference centroid, defines the measurement.
+GALAXY_LINE_DEFINITIONS: tuple[SpectralLineDefinition, ...] = (
+ SpectralLineDefinition(
+ "oii_3727_3729",
+ "[O II] 3727/3729 blend",
+ 3728.48,
+ "emission",
+ (3718.0, 3739.0),
+ ((3653.0, 3710.0), (3745.0, 3800.0)),
+ "The doublet is measured as one blended feature.",
+ ),
+ SpectralLineDefinition(
+ "ca_k_3934",
+ "Ca K",
+ 3934.78,
+ "absorption",
+ (3923.0, 3946.0),
+ ((3890.0, 3915.0), (4005.0, 4030.0)),
+ ),
+ SpectralLineDefinition(
+ "ca_h_3969",
+ "Ca H",
+ 3969.59,
+ "absorption",
+ (3957.0, 3982.0),
+ ((3890.0, 3915.0), (4005.0, 4030.0)),
+ "Ca H can be blended with H epsilon at this resolution.",
+ ),
+ SpectralLineDefinition(
+ "hbeta_4863",
+ "H-beta",
+ 4862.68,
+ "emission",
+ (4854.0, 4872.0),
+ ((4770.0, 4820.0), (4890.0, 4930.0)),
+ ),
+ SpectralLineDefinition(
+ "oiii_4960",
+ "[O III] 4960",
+ 4960.30,
+ "emission",
+ (4953.0, 4968.0),
+ ((4890.0, 4930.0), (5020.0, 5050.0)),
+ ),
+ SpectralLineDefinition(
+ "oiii_5008",
+ "[O III] 5008",
+ 5008.24,
+ "emission",
+ (5000.0, 5017.0),
+ ((4970.0, 4990.0), (5020.0, 5050.0)),
+ ),
+ SpectralLineDefinition(
+ "mgb_5175",
+ "Mg b blend",
+ 5175.0,
+ "absorption",
+ (5155.0, 5196.0),
+ ((5100.0, 5140.0), (5205.0, 5245.0)),
+ "The Mg b triplet is measured as one blended absorption feature.",
+ ),
+ SpectralLineDefinition(
+ "nad_5893",
+ "Na D blend",
+ 5893.0,
+ "absorption",
+ (5878.0, 5910.0),
+ ((5820.0, 5860.0), (5925.0, 5965.0)),
+ "The Na D doublet is measured as one blended absorption feature.",
+ ),
+ SpectralLineDefinition(
+ "nii_6550",
+ "[N II] 6550",
+ 6549.86,
+ "emission",
+ (6543.0, 6556.0),
+ ((6480.0, 6525.0), (6605.0, 6640.0)),
+ ),
+ SpectralLineDefinition(
+ "halpha_6565",
+ "H-alpha",
+ 6564.61,
+ "emission",
+ (6557.0, 6573.0),
+ ((6480.0, 6525.0), (6605.0, 6640.0)),
+ ),
+ SpectralLineDefinition(
+ "nii_6585",
+ "[N II] 6585",
+ 6585.27,
+ "emission",
+ (6578.0, 6593.0),
+ ((6480.0, 6525.0), (6605.0, 6640.0)),
+ ),
+ SpectralLineDefinition(
+ "sii_6718",
+ "[S II] 6718",
+ 6718.29,
+ "emission",
+ (6711.0, 6725.0),
+ ((6650.0, 6690.0), (6750.0, 6790.0)),
+ ),
+ SpectralLineDefinition(
+ "sii_6733",
+ "[S II] 6733",
+ 6732.67,
+ "emission",
+ (6726.0, 6740.0),
+ ((6650.0, 6690.0), (6750.0, 6790.0)),
+ ),
+)
+
+
+DESI_ARM_OVERLAPS_ANGSTROM: tuple[tuple[float, float], ...] = (
+ (5635.0, 5939.0),
+ (7445.0, 7731.0),
+)
+
+# Broad review flags, not extra mask bits. The extracted boolean mask has
+# already collapsed the original DESI bit reasons, and no sky model was
+# retained, so these intervals only request human inspection.
+REVIEW_SENSITIVE_OBSERVED_INTERVALS_ANGSTROM: tuple[tuple[str, float, float], ...] = (
+ ("telluric_b_band_neighborhood", 6860.0, 6950.0),
+ ("telluric_a_band_neighborhood", 7580.0, 7700.0),
+ ("red_sky_residual_sensitive", 7000.0, 9824.0),
+)
+
+
+def _as_1d_float(name: str, value: np.ndarray) -> np.ndarray:
+ array = np.asarray(value, dtype=np.float64)
+ if array.ndim != 1:
+ raise ValueError(f"{name} must be one-dimensional, got {array.shape}")
+ return array
+
+
+def _validated_arrays(
+ wavelength: np.ndarray,
+ flux: np.ndarray,
+ ivar: np.ndarray,
+ mask: np.ndarray,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+ wavelength = _as_1d_float("wavelength", wavelength)
+ flux = _as_1d_float("flux", flux)
+ ivar = _as_1d_float("ivar", ivar)
+ mask = np.asarray(mask, dtype=bool)
+ if mask.ndim != 1:
+ raise ValueError(f"mask must be one-dimensional, got {mask.shape}")
+ if not (wavelength.shape == flux.shape == ivar.shape == mask.shape):
+ raise ValueError("wavelength, flux, ivar, and mask shapes must agree")
+ if wavelength.size < 3:
+ raise ValueError("a spectrum must contain at least three samples")
+ if not np.isfinite(wavelength).all() or np.any(np.diff(wavelength) <= 0):
+ raise ValueError("wavelength must be finite and strictly increasing")
+ if not np.isfinite(flux).all():
+ raise ValueError("flux must be finite")
+ if not np.isfinite(ivar).all() or np.any(ivar < 0):
+ raise ValueError("ivar must be finite and nonnegative")
+ return wavelength, flux, ivar, mask
+
+
+def _interval_mask(values: np.ndarray, interval: tuple[float, float]) -> np.ndarray:
+ return (values >= interval[0]) & (values <= interval[1])
+
+
+def _robust_weighted_linear_fit(
+ x: np.ndarray,
+ y: np.ndarray,
+ weight: np.ndarray,
+ *,
+ iterations: int = 4,
+ clip_sigma: float = 4.0,
+) -> tuple[np.ndarray, np.ndarray, float, np.ndarray]:
+ """Fit ``y = intercept + slope*x`` with conservative sigma clipping."""
+
+ keep = np.ones(y.size, dtype=bool)
+ coefficients = np.array([np.nan, np.nan], dtype=np.float64)
+ covariance = np.full((2, 2), np.nan, dtype=np.float64)
+ residual_mad = np.nan
+ for _ in range(iterations):
+ if np.count_nonzero(keep) < 8:
+ raise ValueError("too few continuum samples after clipping")
+ design = np.column_stack((np.ones(np.count_nonzero(keep)), x[keep]))
+ sqrt_weight = np.sqrt(weight[keep])
+ weighted_design = design * sqrt_weight[:, None]
+ weighted_y = y[keep] * sqrt_weight
+ coefficients, _, _, _ = np.linalg.lstsq(weighted_design, weighted_y, rcond=None)
+ residual = y - (coefficients[0] + coefficients[1] * x)
+ center = float(np.median(residual[keep]))
+ residual_mad = float(1.4826 * np.median(np.abs(residual[keep] - center)))
+ if not np.isfinite(residual_mad) or residual_mad <= 0:
+ break
+ updated = np.abs(residual - center) <= clip_sigma * residual_mad
+ if np.array_equal(updated, keep):
+ break
+ keep = updated
+
+ if np.count_nonzero(keep) < 8:
+ raise ValueError("too few continuum samples after clipping")
+ design = np.column_stack((np.ones(np.count_nonzero(keep)), x[keep]))
+ sqrt_weight = np.sqrt(weight[keep])
+ weighted_design = design * sqrt_weight[:, None]
+ weighted_y = y[keep] * sqrt_weight
+ coefficients, _, _, _ = np.linalg.lstsq(weighted_design, weighted_y, rcond=None)
+ full_residual = y - (coefficients[0] + coefficients[1] * x)
+ center = float(np.median(full_residual[keep]))
+ residual_mad = float(1.4826 * np.median(np.abs(full_residual[keep] - center)))
+ normal = design.T @ (weight[keep, None] * design)
+ covariance = np.linalg.pinv(normal)
+ residual = full_residual[keep]
+ degrees_of_freedom = max(np.count_nonzero(keep) - 2, 1)
+ reduced_chi2 = float(np.sum(weight[keep] * residual**2) / degrees_of_freedom)
+ covariance *= max(reduced_chi2, 1.0)
+ return coefficients, covariance, residual_mad, keep
+
+
+def measure_line(
+ wavelength: np.ndarray,
+ flux: np.ndarray,
+ ivar: np.ndarray,
+ mask: np.ndarray,
+ *,
+ catalog_redshift: float,
+ definition: SpectralLineDefinition,
+ detection_sigma: float = 5.0,
+ tentative_sigma: float = 3.0,
+) -> dict[str, Any]:
+ """Measure one line relative to a clipped local linear continuum.
+
+ Equivalent widths use the signed convention: emission is positive and
+ absorption is negative. The line lookup is explicitly conditioned on the supplied
+ catalog redshift; it is not an independent redshift measurement.
+ """
+
+ wavelength, flux, ivar, mask = _validated_arrays(wavelength, flux, ivar, mask)
+ if not np.isfinite(catalog_redshift) or catalog_redshift < 0:
+ raise ValueError("catalog_redshift must be finite and nonnegative")
+ rest_wavelength = wavelength / (1.0 + catalog_redshift)
+ base_valid = (~mask) & (ivar > 0)
+ line_grid = _interval_mask(rest_wavelength, definition.line_window_rest_angstrom)
+ line_valid = line_grid & base_valid
+ side_grid = _interval_mask(
+ rest_wavelength, definition.continuum_windows_rest_angstrom[0]
+ ) | _interval_mask(rest_wavelength, definition.continuum_windows_rest_angstrom[1])
+ side_valid = side_grid & base_valid
+ line_coverage = float(np.count_nonzero(line_valid) / max(np.count_nonzero(line_grid), 1))
+ side_coverage = float(np.count_nonzero(side_valid) / max(np.count_nonzero(side_grid), 1))
+ required_min = min(
+ definition.line_window_rest_angstrom[0],
+ *(window[0] for window in definition.continuum_windows_rest_angstrom),
+ )
+ required_max = max(
+ definition.line_window_rest_angstrom[1],
+ *(window[1] for window in definition.continuum_windows_rest_angstrom),
+ )
+ full_window_coverage = (
+ rest_wavelength[0] <= required_min and rest_wavelength[-1] >= required_max
+ )
+
+ common: dict[str, Any] = {
+ "key": definition.key,
+ "label": definition.label,
+ "nature": definition.nature,
+ "rest_wavelength_angstrom": definition.rest_wavelength_angstrom,
+ "expected_observed_wavelength_angstrom": (
+ definition.rest_wavelength_angstrom * (1.0 + catalog_redshift)
+ ),
+ "line_window_rest_angstrom": list(definition.line_window_rest_angstrom),
+ "continuum_windows_rest_angstrom": [
+ list(window) for window in definition.continuum_windows_rest_angstrom
+ ],
+ "line_valid_fraction": line_coverage,
+ "continuum_valid_fraction": side_coverage,
+ "conditioned_on_catalog_redshift": True,
+ "caveat": definition.caveat,
+ }
+ if (
+ not full_window_coverage
+ or np.count_nonzero(line_grid) == 0
+ or np.count_nonzero(side_grid) == 0
+ ):
+ return {
+ **common,
+ "status": "out_of_coverage",
+ "quality": "unavailable",
+ "detection_class": "unavailable",
+ }
+ if line_coverage < 0.8 or side_coverage < 0.8 or np.count_nonzero(side_valid) < 16:
+ return {
+ **common,
+ "status": "masked",
+ "quality": "unavailable",
+ "detection_class": "unavailable",
+ }
+
+ center = definition.rest_wavelength_angstrom
+ side_x = rest_wavelength[side_valid] - center
+ side_y = flux[side_valid]
+ side_weight = ivar[side_valid]
+ try:
+ coefficients, covariance, residual_mad, keep = _robust_weighted_linear_fit(
+ side_x, side_y, side_weight
+ )
+ except (ValueError, np.linalg.LinAlgError):
+ return {
+ **common,
+ "status": "not_assessed",
+ "quality": "unavailable",
+ "detection_class": "unavailable",
+ }
+
+ line_x = rest_wavelength[line_valid] - center
+ line_flux = flux[line_valid]
+ line_ivar = ivar[line_valid]
+ line_wavelength = wavelength[line_valid]
+ local_continuum = coefficients[0] + coefficients[1] * line_x
+ delta_observed = np.gradient(line_wavelength)
+ direction = 1.0 if definition.nature == "emission" else -1.0
+ signed_residual = line_flux - local_continuum
+ integrated_signed_strength = float(np.sum(signed_residual * delta_observed))
+ expected_nature_strength = direction * integrated_signed_strength
+
+ pixel_variance = float(np.sum(delta_observed**2 / line_ivar))
+ continuum_vector = np.array(
+ [np.sum(delta_observed), np.sum(delta_observed * line_x)],
+ dtype=np.float64,
+ )
+ continuum_variance = float(continuum_vector @ covariance @ continuum_vector)
+ formal_uncertainty = float(np.sqrt(max(pixel_variance + continuum_variance, 0.0)))
+ empirical_uncertainty = float(
+ residual_mad * np.sqrt(np.sum(delta_observed**2)) if np.isfinite(residual_mad) else 0.0
+ )
+ strength_uncertainty = max(formal_uncertainty, empirical_uncertainty)
+ significance = (
+ expected_nature_strength / strength_uncertainty if strength_uncertainty > 0 else np.nan
+ )
+ continuum_level = float(np.median(local_continuum))
+ if continuum_level > 0:
+ equivalent_width = integrated_signed_strength / continuum_level / (1.0 + catalog_redshift)
+ equivalent_width_uncertainty = (
+ strength_uncertainty / continuum_level / (1.0 + catalog_redshift)
+ )
+ else:
+ equivalent_width = np.nan
+ equivalent_width_uncertainty = np.nan
+
+ positive_residual = np.clip(direction * signed_residual, 0.0, None)
+ if np.sum(positive_residual) > 0:
+ observed_centroid = float(
+ np.sum(line_wavelength * positive_residual) / np.sum(positive_residual)
+ )
+ else:
+ observed_centroid = np.nan
+
+ if np.isfinite(significance) and significance >= detection_sigma:
+ status = "measured"
+ quality = "valid"
+ detection_class = "detected"
+ elif np.isfinite(significance) and significance >= tentative_sigma:
+ status = "measured"
+ quality = "low_confidence"
+ detection_class = "tentative"
+ else:
+ status = "not_detected"
+ quality = "valid"
+ detection_class = "not_detected"
+
+ quality_flags: list[str] = []
+ masked_line_pixels = int(np.count_nonzero(line_grid & (~base_valid)))
+ if masked_line_pixels:
+ quality_flags.append("masked_pixels_in_line_window")
+ if quality == "valid":
+ quality = "low_confidence"
+ observed_lo = definition.line_window_rest_angstrom[0] * (1.0 + catalog_redshift)
+ observed_hi = definition.line_window_rest_angstrom[1] * (1.0 + catalog_redshift)
+ for overlap_lo, overlap_hi in DESI_ARM_OVERLAPS_ANGSTROM:
+ if observed_hi >= overlap_lo and observed_lo <= overlap_hi:
+ quality_flags.append("desi_arm_overlap_region")
+ break
+ for flag, sensitive_lo, sensitive_hi in REVIEW_SENSITIVE_OBSERVED_INTERVALS_ANGSTROM:
+ if observed_hi >= sensitive_lo and observed_lo <= sensitive_hi:
+ quality_flags.append(flag)
+
+ return {
+ **common,
+ "status": status,
+ "quality": quality,
+ "detection_class": detection_class,
+ "quality_flags": quality_flags,
+ "masked_line_pixels": masked_line_pixels,
+ "integrated_signed_flux_source_flux_angstrom": integrated_signed_strength,
+ "integrated_strength_uncertainty_source_flux_angstrom": strength_uncertainty,
+ "significance_sigma": float(significance),
+ "equivalent_width_signed_rest_angstrom": float(equivalent_width),
+ "equivalent_width_uncertainty_angstrom": float(equivalent_width_uncertainty),
+ "observed_centroid_angstrom": (
+ observed_centroid if detection_class != "not_detected" else None
+ ),
+ "local_continuum_source_flux": continuum_level,
+ "continuum_sideband_mad_source_flux": float(residual_mad),
+ "continuum_samples_used": int(np.count_nonzero(keep)),
+ }
+
+
+def _band_mean(
+ rest_wavelength: np.ndarray,
+ flux: np.ndarray,
+ ivar: np.ndarray,
+ valid: np.ndarray,
+ interval: tuple[float, float],
+ *,
+ fnu: bool,
+) -> tuple[float, float, int, float]:
+ grid = _interval_mask(rest_wavelength, interval)
+ selected = grid & valid
+ count = int(np.count_nonzero(selected))
+ coverage = float(count / max(np.count_nonzero(grid), 1))
+ if count == 0:
+ return np.nan, np.nan, 0, coverage
+ scale = rest_wavelength[selected] ** 2 if fnu else np.ones(count)
+ values = flux[selected] * scale
+ mean = float(np.mean(values))
+ uncertainty = float(np.sqrt(np.sum(scale**2 / ivar[selected])) / count)
+ return mean, uncertainty, count, coverage
+
+
+def _continuum_band_location(
+ rest_wavelength: np.ndarray,
+ flux: np.ndarray,
+ ivar: np.ndarray,
+ valid: np.ndarray,
+ interval: tuple[float, float],
+) -> tuple[float, float, int, float]:
+ grid = _interval_mask(rest_wavelength, interval)
+ # Remove the predefined diagnostic line windows from broad-band continuum.
+ excluded = np.zeros_like(valid)
+ for definition in GALAXY_LINE_DEFINITIONS:
+ lo, hi = definition.line_window_rest_angstrom
+ excluded |= (rest_wavelength >= lo - 4.0) & (rest_wavelength <= hi + 4.0)
+ selected = grid & valid & (~excluded)
+ total = int(np.count_nonzero(grid & (~excluded)))
+ coverage = float(np.count_nonzero(selected) / max(total, 1))
+ values = flux[selected]
+ variances = 1.0 / ivar[selected]
+ if values.size == 0:
+ return np.nan, np.nan, 0, coverage
+ keep = np.ones(values.size, dtype=bool)
+ for _ in range(4):
+ median = float(np.median(values[keep]))
+ mad = float(1.4826 * np.median(np.abs(values[keep] - median)))
+ if not np.isfinite(mad) or mad <= 0:
+ break
+ updated = np.abs(values - median) <= 4.0 * mad
+ if np.array_equal(updated, keep):
+ break
+ keep = updated
+ location = float(np.mean(values[keep]))
+ formal = float(np.sqrt(np.sum(variances[keep])) / np.count_nonzero(keep))
+ scatter = float(np.std(values[keep], ddof=1) / np.sqrt(np.count_nonzero(keep)))
+ return location, max(formal, scatter), int(np.count_nonzero(keep)), coverage
+
+
+def measure_spectrum_facts(
+ wavelength: np.ndarray,
+ flux: np.ndarray,
+ ivar: np.ndarray,
+ mask: np.ndarray,
+ *,
+ catalog_redshift: float,
+) -> dict[str, Any]:
+ """Return a compact grounded report for one observed-frame spectrum."""
+
+ wavelength, flux, ivar, mask = _validated_arrays(wavelength, flux, ivar, mask)
+ if not np.isfinite(catalog_redshift) or catalog_redshift < 0:
+ raise ValueError("catalog_redshift must be finite and nonnegative")
+ valid = (~mask) & (ivar > 0)
+ if np.count_nonzero(valid) == 0:
+ raise ValueError("spectrum has no unmasked pixels with positive inverse variance")
+ rest_wavelength = wavelength / (1.0 + catalog_redshift)
+
+ dn_blue, dn_blue_error, dn_blue_count, dn_blue_coverage = _band_mean(
+ rest_wavelength, flux, ivar, valid, (3850.0, 3950.0), fnu=True
+ )
+ dn_red, dn_red_error, dn_red_count, dn_red_coverage = _band_mean(
+ rest_wavelength, flux, ivar, valid, (4000.0, 4100.0), fnu=True
+ )
+ dn_bands_in_range = rest_wavelength[0] <= 3850.0 and rest_wavelength[-1] >= 4100.0
+ dn_coverage_valid = dn_bands_in_range and dn_blue_coverage >= 0.9 and dn_red_coverage >= 0.9
+ if dn_blue > 0 and dn_red > 0 and dn_coverage_valid:
+ dn4000 = float(dn_red / dn_blue)
+ dn4000_error = float(
+ dn4000 * np.sqrt((dn_red_error / dn_red) ** 2 + (dn_blue_error / dn_blue) ** 2)
+ )
+ else:
+ dn4000 = np.nan
+ dn4000_error = np.nan
+ if not np.isfinite(dn4000):
+ dn4000_label = "unavailable"
+ elif dn4000 >= 1.6:
+ dn4000_label = "strong"
+ elif dn4000 >= 1.3:
+ dn4000_label = "moderate"
+ else:
+ dn4000_label = "weak"
+
+ blue, blue_error, blue_count, blue_coverage = _continuum_band_location(
+ rest_wavelength, flux, ivar, valid, (4200.0, 4600.0)
+ )
+ red, red_error, red_count, red_coverage = _continuum_band_location(
+ rest_wavelength, flux, ivar, valid, (6000.0, 6400.0)
+ )
+ continuum_bands_in_range = rest_wavelength[0] <= 4200.0 and rest_wavelength[-1] >= 6400.0
+ continuum_coverage_valid = (
+ continuum_bands_in_range and blue_coverage >= 0.9 and red_coverage >= 0.9
+ )
+ if blue > 0 and red > 0 and continuum_coverage_valid:
+ red_blue_ratio = float(red / blue)
+ red_blue_error = float(
+ red_blue_ratio * np.sqrt((red_error / red) ** 2 + (blue_error / blue) ** 2)
+ )
+ else:
+ red_blue_ratio = np.nan
+ red_blue_error = np.nan
+ if not np.isfinite(red_blue_ratio):
+ continuum_label = "unavailable"
+ elif red_blue_ratio >= 1.2:
+ continuum_label = "red/rising"
+ elif red_blue_ratio >= 1.1:
+ continuum_label = "mildly red/rising"
+ elif red_blue_ratio <= 0.8:
+ continuum_label = "blue/declining"
+ elif red_blue_ratio <= 0.9:
+ continuum_label = "mildly blue/declining"
+ else:
+ continuum_label = "approximately flat"
+
+ optical = valid & (rest_wavelength >= 3850.0) & (rest_wavelength <= 7000.0)
+ pixel_snr = flux[optical] * np.sqrt(ivar[optical])
+ median_pixel_snr = float(np.median(pixel_snr)) if pixel_snr.size else np.nan
+
+ lines = [
+ measure_line(
+ wavelength,
+ flux,
+ ivar,
+ mask,
+ catalog_redshift=catalog_redshift,
+ definition=definition,
+ )
+ for definition in GALAXY_LINE_DEFINITIONS
+ ]
+ return {
+ "measurement_version": "astroprism-spectrum-facts-v1",
+ "wavelength_frame": "observed_vacuum",
+ "wavelength_unit": "angstrom",
+ "flux_unit": "source_array_unit_not_encoded_in_local_artifact",
+ "catalog_redshift": float(catalog_redshift),
+ "coverage": {
+ "observed_min_angstrom": float(wavelength[0]),
+ "observed_max_angstrom": float(wavelength[-1]),
+ "rest_min_angstrom": float(rest_wavelength[0]),
+ "rest_max_angstrom": float(rest_wavelength[-1]),
+ "valid_pixels": int(np.count_nonzero(valid)),
+ "total_pixels": int(wavelength.size),
+ "valid_fraction": float(np.count_nonzero(valid) / wavelength.size),
+ },
+ "quality": {
+ "median_pixel_snr_rest_3850_7000": median_pixel_snr,
+ "definition": "median flux*sqrt(ivar) over valid rest-frame pixels",
+ },
+ "continuum": {
+ "label": continuum_label,
+ "status": "measured" if np.isfinite(red_blue_ratio) else "unavailable",
+ "red_to_blue_ratio": red_blue_ratio,
+ "red_to_blue_ratio_uncertainty": red_blue_error,
+ "blue_window_rest_angstrom": [4200.0, 4600.0],
+ "red_window_rest_angstrom": [6000.0, 6400.0],
+ "blue_location_source_flux": blue,
+ "red_location_source_flux": red,
+ "blue_samples": blue_count,
+ "red_samples": red_count,
+ "blue_valid_fraction": blue_coverage,
+ "red_valid_fraction": red_coverage,
+ "method": "four-sigma-clipped mean excluding diagnostic line windows",
+ },
+ "dn4000": {
+ "value": dn4000,
+ "uncertainty": dn4000_error,
+ "strength_label": dn4000_label,
+ "status": "measured" if np.isfinite(dn4000) else "unavailable",
+ "definition": "narrow Dn4000 ratio of mean F_nu in 4000-4100 to 3850-3950 angstrom",
+ "blue_samples": dn_blue_count,
+ "red_samples": dn_red_count,
+ "blue_valid_fraction": dn_blue_coverage,
+ "red_valid_fraction": dn_red_coverage,
+ },
+ "lines": lines,
+ "limitations": [
+ "Line searches and rest-frame windows are conditioned on the external catalog redshift.",
+ "Integrated uncertainties use the larger of propagated inverse-variance and local empirical-noise estimates.",
+ "Local linear continua and fixed windows are suitable for caption screening, not publication-quality deblending.",
+ "Signed rest-frame equivalent width is positive for emission and negative for absorption.",
+ "Broad observed-frame quality flags request review because the original DESI mask bits, sky model, and resolution matrix were not retained.",
+ "The extracted local arrays do not retain an explicit calibrated flux-unit attribute.",
+ ],
+ }
+
+
+__all__ = [
+ "GALAXY_LINE_DEFINITIONS",
+ "SpectralLineDefinition",
+ "measure_line",
+ "measure_spectrum_facts",
+]
diff --git a/src/data/astroprism_stage1a.py b/src/data/astroprism_stage1a.py
new file mode 100644
index 0000000..100bf58
--- /dev/null
+++ b/src/data/astroprism_stage1a.py
@@ -0,0 +1,807 @@
+"""Prepare the pinned AstroPRISM AION feature parity reference.
+
+This module is intentionally offline. It validates the already-downloaded
+Hugging Face Parquet shards against their download metadata, selects rows by
+the locators in the committed roster, and writes released AION image embeddings
+only as a parity/regression oracle. These precomputed features are not Stage 1A
+model inputs: training starts from raw calibrated cutouts and runs the pinned
+AION codec/shared encoder. Released spectrum embeddings are never requested.
+"""
+
+from __future__ import annotations
+
+import csv
+import gzip
+import hashlib
+import json
+import math
+import os
+import shutil
+import tempfile
+import uuid
+from collections import defaultdict
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path, PurePosixPath
+from typing import Any, TextIO
+
+import numpy as np
+import pyarrow as pa
+import pyarrow.compute as pc
+import pyarrow.parquet as pq
+
+ROSTER_SHA256 = "650935faa58953c9929e7379429bfb332a7486dab54544fedef64ca385abb3bc"
+PAIRED_REPO_ID = "EiffL/desi_x_ls_aion"
+PAIRED_REVISION = "57f7c7575260811ea076163bc743787d75b9aa49"
+CAPTION_REPO_ID = "astronolan/galaxy-descriptions"
+CAPTION_REVISION = "6890dc0c8fc793867c07a6ce4fce11e51c167d6e"
+EXPECTED_CANDIDATES = 646
+IMAGE_EMBEDDING_SHAPE = (576, 768)
+
+_PAIRED_COLUMNS = (
+ "object_id_provabgs",
+ "ra_provabgs",
+ "dec_provabgs",
+ "embedding_image",
+)
+_CAPTION_COLUMNS = ("object_id", "survey", "ra", "dec", "caption", "summary")
+_REQUIRED_ROSTER_COLUMNS = {
+ "paired_targetid",
+ "caption_object_id",
+ "paired_ra",
+ "paired_dec",
+ "caption_ra",
+ "caption_dec",
+ "separation_arcsec",
+ "paired_candidate_count",
+ "caption_candidate_count",
+ "one_to_one_row_candidate",
+ "unique_source_ids",
+ "provisional_unambiguous_candidate",
+ "paired_shard",
+ "paired_row",
+ "caption_shard",
+ "caption_row",
+}
+
+
+class Stage1AValidationError(ValueError):
+ """Raised when a source asset violates the pinned Stage 1A contract."""
+
+
+@dataclass(frozen=True)
+class Stage1AConfig:
+ """Inputs and validation pins for one extraction."""
+
+ roster_path: Path
+ paired_root: Path
+ caption_root: Path
+ output_dir: Path
+ validate_only: bool = False
+ overwrite: bool = False
+ batch_size: int = 16
+ expected_count: int = EXPECTED_CANDIDATES
+ expected_roster_sha256: str | None = ROSTER_SHA256
+ paired_revision: str = PAIRED_REVISION
+ caption_revision: str = CAPTION_REVISION
+ coordinate_tolerance_deg: float = 1e-9
+ separation_tolerance_arcsec: float = 1e-6
+
+
+@dataclass(frozen=True)
+class _SourceFile:
+ kind: str
+ repo_id: str
+ revision: str
+ shard: str
+ path: Path
+ size_bytes: int
+ sha256: str
+ metadata_sha256: str
+
+ def manifest_entry(self) -> dict[str, Any]:
+ return {
+ "shard": self.shard,
+ "size_bytes": self.size_bytes,
+ "sha256": self.sha256,
+ "download_metadata_sha256": self.metadata_sha256,
+ }
+
+
+def sha256_file(path: Path, chunk_size: int = 8 * 1024 * 1024) -> str:
+ """Return the SHA-256 digest of *path* without loading it into memory."""
+
+ digest = hashlib.sha256()
+ with path.open("rb") as handle:
+ while chunk := handle.read(chunk_size):
+ digest.update(chunk)
+ return digest.hexdigest()
+
+
+def _open_roster(path: Path) -> TextIO:
+ if path.suffix == ".gz":
+ return gzip.open(path, "rt", encoding="utf-8", newline="")
+ return path.open("r", encoding="utf-8", newline="")
+
+
+def _parse_int(value: str, *, field: str, index: int) -> int:
+ try:
+ parsed = int(value)
+ except (TypeError, ValueError) as exc:
+ raise Stage1AValidationError(
+ f"roster row {index}: {field} must be an integer, got {value!r}"
+ ) from exc
+ return parsed
+
+
+def _parse_float(value: str, *, field: str, index: int) -> float:
+ try:
+ parsed = float(value)
+ except (TypeError, ValueError) as exc:
+ raise Stage1AValidationError(
+ f"roster row {index}: {field} must be a float, got {value!r}"
+ ) from exc
+ if not math.isfinite(parsed):
+ raise Stage1AValidationError(f"roster row {index}: {field} is not finite")
+ return parsed
+
+
+def _parse_true(value: str, *, field: str, index: int) -> None:
+ if str(value).strip().lower() not in {"true", "1"}:
+ raise Stage1AValidationError(f"roster row {index}: {field} must be true, got {value!r}")
+
+
+def _read_roster(config: Stage1AConfig) -> tuple[list[dict[str, str]], str]:
+ path = Path(config.roster_path)
+ if not path.is_file():
+ raise Stage1AValidationError(f"roster does not exist: {path}")
+
+ roster_sha256 = sha256_file(path)
+ if config.expected_roster_sha256 is not None and roster_sha256 != config.expected_roster_sha256:
+ raise Stage1AValidationError(
+ "roster SHA-256 mismatch: "
+ f"expected {config.expected_roster_sha256}, got {roster_sha256}"
+ )
+
+ with _open_roster(path) as handle:
+ reader = csv.DictReader(handle)
+ fieldnames = reader.fieldnames or []
+ if len(fieldnames) != len(set(fieldnames)):
+ raise Stage1AValidationError("roster contains duplicate column names")
+ missing = sorted(_REQUIRED_ROSTER_COLUMNS - set(fieldnames))
+ if missing:
+ raise Stage1AValidationError(f"roster is missing required columns: {missing}")
+ rows = list(reader)
+
+ if len(rows) != config.expected_count:
+ raise Stage1AValidationError(
+ f"roster must contain exactly {config.expected_count} rows, got {len(rows)}"
+ )
+
+ paired_ids: set[str] = set()
+ caption_ids: set[str] = set()
+ paired_locators: set[tuple[str, int]] = set()
+ caption_locators: set[tuple[str, int]] = set()
+ for index, row in enumerate(rows):
+ paired_id = row["paired_targetid"]
+ caption_id = row["caption_object_id"]
+ if not paired_id or not caption_id:
+ raise Stage1AValidationError(f"roster row {index}: identifiers must be nonempty")
+ if paired_id in paired_ids:
+ raise Stage1AValidationError(f"duplicate paired_targetid: {paired_id}")
+ if caption_id in caption_ids:
+ raise Stage1AValidationError(f"duplicate caption_object_id: {caption_id}")
+ paired_ids.add(paired_id)
+ caption_ids.add(caption_id)
+
+ paired_row = _parse_int(row["paired_row"], field="paired_row", index=index)
+ caption_row = _parse_int(row["caption_row"], field="caption_row", index=index)
+ if paired_row < 0 or caption_row < 0:
+ raise Stage1AValidationError(f"roster row {index}: source rows must be nonnegative")
+ paired_locator = (row["paired_shard"], paired_row)
+ caption_locator = (row["caption_shard"], caption_row)
+ if paired_locator in paired_locators:
+ raise Stage1AValidationError(f"duplicate paired locator: {paired_locator}")
+ if caption_locator in caption_locators:
+ raise Stage1AValidationError(f"duplicate caption locator: {caption_locator}")
+ paired_locators.add(paired_locator)
+ caption_locators.add(caption_locator)
+
+ for field in ("paired_candidate_count", "caption_candidate_count"):
+ if _parse_int(row[field], field=field, index=index) != 1:
+ raise Stage1AValidationError(f"roster row {index}: {field} must equal 1")
+ for field in (
+ "one_to_one_row_candidate",
+ "unique_source_ids",
+ "provisional_unambiguous_candidate",
+ ):
+ _parse_true(row[field], field=field, index=index)
+
+ paired_ra = _parse_float(row["paired_ra"], field="paired_ra", index=index)
+ paired_dec = _parse_float(row["paired_dec"], field="paired_dec", index=index)
+ caption_ra = _parse_float(row["caption_ra"], field="caption_ra", index=index)
+ caption_dec = _parse_float(row["caption_dec"], field="caption_dec", index=index)
+ separation = _parse_float(row["separation_arcsec"], field="separation_arcsec", index=index)
+ if not 0.0 <= paired_ra < 360.0 or not 0.0 <= caption_ra < 360.0:
+ raise Stage1AValidationError(f"roster row {index}: RA is outside [0, 360)")
+ if not -90.0 <= paired_dec <= 90.0 or not -90.0 <= caption_dec <= 90.0:
+ raise Stage1AValidationError(f"roster row {index}: declination is outside [-90, 90]")
+ if not 0.0 <= separation <= 1.0:
+ raise Stage1AValidationError(f"roster row {index}: separation_arcsec is outside [0, 1]")
+
+ return rows, roster_sha256
+
+
+def _safe_source_path(root: Path, shard: str) -> Path:
+ pure = PurePosixPath(shard)
+ if pure.is_absolute() or not pure.parts or ".." in pure.parts:
+ raise Stage1AValidationError(f"unsafe source shard path: {shard!r}")
+ root = root.resolve()
+ path = (root / Path(*pure.parts)).resolve()
+ if not path.is_relative_to(root):
+ raise Stage1AValidationError(f"source shard escapes its root: {shard!r}")
+ if not path.is_file():
+ raise Stage1AValidationError(f"source shard does not exist: {path}")
+ return path
+
+
+def _source_file(*, kind: str, repo_id: str, revision: str, root: Path, shard: str) -> _SourceFile:
+ root = Path(root)
+ path = _safe_source_path(root, shard)
+ metadata_path = root / ".cache" / "huggingface" / "download" / f"{shard}.metadata"
+ if not metadata_path.is_file():
+ raise Stage1AValidationError(f"missing Hugging Face download metadata: {metadata_path}")
+ lines = metadata_path.read_text(encoding="utf-8").splitlines()
+ if len(lines) < 2:
+ raise Stage1AValidationError(f"malformed Hugging Face download metadata: {metadata_path}")
+ observed_revision, expected_sha256 = lines[0].strip(), lines[1].strip().lower()
+ if observed_revision != revision:
+ raise Stage1AValidationError(
+ f"revision mismatch for {shard}: expected {revision}, got {observed_revision}"
+ )
+ if len(expected_sha256) != 64 or any(c not in "0123456789abcdef" for c in expected_sha256):
+ raise Stage1AValidationError(
+ f"download metadata for {shard} does not contain an LFS SHA-256"
+ )
+ observed_sha256 = sha256_file(path)
+ if observed_sha256 != expected_sha256:
+ raise Stage1AValidationError(
+ f"source SHA-256 mismatch for {shard}: expected {expected_sha256}, "
+ f"got {observed_sha256}"
+ )
+ return _SourceFile(
+ kind=kind,
+ repo_id=repo_id,
+ revision=revision,
+ shard=shard,
+ path=path,
+ size_bytes=path.stat().st_size,
+ sha256=observed_sha256,
+ metadata_sha256=sha256_file(metadata_path),
+ )
+
+
+def _collect_sources(
+ config: Stage1AConfig, rows: list[dict[str, str]]
+) -> tuple[dict[str, _SourceFile], dict[str, _SourceFile]]:
+ paired = {
+ shard: _source_file(
+ kind="paired",
+ repo_id=PAIRED_REPO_ID,
+ revision=config.paired_revision,
+ root=Path(config.paired_root),
+ shard=shard,
+ )
+ for shard in sorted({row["paired_shard"] for row in rows})
+ }
+ captions = {
+ shard: _source_file(
+ kind="caption",
+ repo_id=CAPTION_REPO_ID,
+ revision=config.caption_revision,
+ root=Path(config.caption_root),
+ shard=shard,
+ )
+ for shard in sorted({row["caption_shard"] for row in rows})
+ }
+ return paired, captions
+
+
+def _require_columns(schema: pa.Schema, columns: tuple[str, ...], *, shard: str) -> None:
+ missing = [name for name in columns if schema.get_field_index(name) < 0]
+ if missing:
+ raise Stage1AValidationError(f"{shard}: missing Parquet columns {missing}")
+
+
+def _check_coordinate(
+ actual: Any,
+ expected_text: str,
+ *,
+ field: str,
+ candidate_index: int,
+ tolerance: float,
+) -> float:
+ if actual is None:
+ raise Stage1AValidationError(f"candidate {candidate_index}: source {field} is null")
+ value = float(actual)
+ expected = float(expected_text)
+ if not math.isfinite(value) or abs(value - expected) > tolerance:
+ raise Stage1AValidationError(
+ f"candidate {candidate_index}: {field} mismatch; expected {expected}, got {value}"
+ )
+ return value
+
+
+def _fixed_embedding_type() -> pa.DataType:
+ return pa.list_(pa.list_(pa.float16(), IMAGE_EMBEDDING_SHAPE[1]), IMAGE_EMBEDDING_SHAPE[0])
+
+
+def _extract_embedding(array: pa.Array, local_index: int, *, candidate_index: int) -> np.ndarray:
+ selected = array.slice(local_index, 1)
+ if selected.null_count:
+ raise Stage1AValidationError(f"candidate {candidate_index}: image embedding is null")
+ flat = pc.list_flatten(pc.list_flatten(selected))
+ if flat.null_count:
+ raise Stage1AValidationError(
+ f"candidate {candidate_index}: image embedding contains null values"
+ )
+ values = flat.to_numpy(zero_copy_only=False)
+ if values.dtype != np.float16 or values.size != math.prod(IMAGE_EMBEDDING_SHAPE):
+ raise Stage1AValidationError(
+ f"candidate {candidate_index}: invalid image embedding buffer "
+ f"({values.dtype}, {values.size} values)"
+ )
+ embedding = values.reshape(IMAGE_EMBEDDING_SHAPE)
+ if not np.isfinite(embedding).all():
+ raise Stage1AValidationError(
+ f"candidate {candidate_index}: image embedding contains non-finite values"
+ )
+ if not np.any(embedding):
+ raise Stage1AValidationError(f"candidate {candidate_index}: image embedding is all zero")
+ return embedding
+
+
+def _scan_paired(
+ config: Stage1AConfig,
+ rows: list[dict[str, str]],
+ sources: dict[str, _SourceFile],
+ output: np.memmap | None,
+) -> tuple[list[dict[str, Any]], list[float]]:
+ grouped: dict[str, list[tuple[int, dict[str, str]]]] = defaultdict(list)
+ for index, row in enumerate(rows):
+ grouped[row["paired_shard"]].append((index, row))
+
+ found: list[dict[str, Any] | None] = [None] * len(rows)
+ norms: list[float] = []
+ for shard in sorted(grouped):
+ source = sources[shard]
+ parquet = pq.ParquetFile(source.path)
+ _require_columns(parquet.schema_arrow, _PAIRED_COLUMNS, shard=shard)
+ embedding_type = parquet.schema_arrow.field("embedding_image").type
+ if embedding_type != _fixed_embedding_type():
+ raise Stage1AValidationError(
+ f"{shard}: embedding_image must be float16{IMAGE_EMBEDDING_SHAPE}, "
+ f"got {embedding_type}"
+ )
+ if not pa.types.is_string(parquet.schema_arrow.field("object_id_provabgs").type):
+ raise Stage1AValidationError(f"{shard}: object_id_provabgs must be a string")
+
+ wanted = sorted(
+ (
+ _parse_int(row["paired_row"], field="paired_row", index=index),
+ index,
+ row,
+ )
+ for index, row in grouped[shard]
+ )
+ if wanted[-1][0] >= parquet.metadata.num_rows:
+ raise Stage1AValidationError(
+ f"{shard}: paired row {wanted[-1][0]} is outside {parquet.metadata.num_rows} rows"
+ )
+
+ pointer = 0
+ row_offset = 0
+ for batch in parquet.iter_batches(
+ batch_size=config.batch_size, columns=list(_PAIRED_COLUMNS), use_threads=False
+ ):
+ batch_end = row_offset + len(batch)
+ while pointer < len(wanted) and wanted[pointer][0] < batch_end:
+ source_row, candidate_index, roster_row = wanted[pointer]
+ if source_row < row_offset:
+ raise Stage1AValidationError(f"{shard}: duplicate or unordered row locator")
+ local_index = source_row - row_offset
+ source_id = batch.column(0)[local_index].as_py()
+ if source_id != roster_row["paired_targetid"]:
+ raise Stage1AValidationError(
+ f"candidate {candidate_index}: paired ID mismatch; expected "
+ f"{roster_row['paired_targetid']!r}, got {source_id!r}"
+ )
+ ra = _check_coordinate(
+ batch.column(1)[local_index].as_py(),
+ roster_row["paired_ra"],
+ field="paired_ra",
+ candidate_index=candidate_index,
+ tolerance=config.coordinate_tolerance_deg,
+ )
+ dec = _check_coordinate(
+ batch.column(2)[local_index].as_py(),
+ roster_row["paired_dec"],
+ field="paired_dec",
+ candidate_index=candidate_index,
+ tolerance=config.coordinate_tolerance_deg,
+ )
+ embedding = _extract_embedding(
+ batch.column(3), local_index, candidate_index=candidate_index
+ )
+ if output is not None:
+ output[candidate_index] = embedding
+ norms.append(float(np.linalg.norm(embedding.astype(np.float32))))
+ found[candidate_index] = {"ra": ra, "dec": dec}
+ pointer += 1
+ row_offset = batch_end
+ if pointer == len(wanted):
+ break
+ if pointer != len(wanted):
+ raise Stage1AValidationError(f"{shard}: not all paired locators were found")
+
+ if any(item is None for item in found):
+ raise Stage1AValidationError("not every roster candidate produced an image embedding")
+ return [item for item in found if item is not None], norms
+
+
+def _scan_captions(
+ config: Stage1AConfig,
+ rows: list[dict[str, str]],
+ sources: dict[str, _SourceFile],
+) -> list[dict[str, Any]]:
+ grouped: dict[str, list[tuple[int, dict[str, str]]]] = defaultdict(list)
+ for index, row in enumerate(rows):
+ grouped[row["caption_shard"]].append((index, row))
+
+ found: list[dict[str, Any] | None] = [None] * len(rows)
+ for shard in sorted(grouped):
+ source = sources[shard]
+ parquet = pq.ParquetFile(source.path)
+ _require_columns(parquet.schema_arrow, _CAPTION_COLUMNS, shard=shard)
+ for field in ("object_id", "survey", "caption", "summary"):
+ if not pa.types.is_string(parquet.schema_arrow.field(field).type):
+ raise Stage1AValidationError(f"{shard}: {field} must be a string")
+
+ wanted = sorted(
+ (
+ _parse_int(row["caption_row"], field="caption_row", index=index),
+ index,
+ row,
+ )
+ for index, row in grouped[shard]
+ )
+ if wanted[-1][0] >= parquet.metadata.num_rows:
+ raise Stage1AValidationError(
+ f"{shard}: caption row {wanted[-1][0]} is outside {parquet.metadata.num_rows} rows"
+ )
+
+ pointer = 0
+ row_offset = 0
+ for batch in parquet.iter_batches(
+ batch_size=max(config.batch_size, 128),
+ columns=list(_CAPTION_COLUMNS),
+ use_threads=False,
+ ):
+ batch_end = row_offset + len(batch)
+ while pointer < len(wanted) and wanted[pointer][0] < batch_end:
+ source_row, candidate_index, roster_row = wanted[pointer]
+ if source_row < row_offset:
+ raise Stage1AValidationError(f"{shard}: duplicate or unordered row locator")
+ local_index = source_row - row_offset
+ source_id = batch.column(0)[local_index].as_py()
+ survey = batch.column(1)[local_index].as_py()
+ if source_id != roster_row["caption_object_id"]:
+ raise Stage1AValidationError(
+ f"candidate {candidate_index}: caption ID mismatch; expected "
+ f"{roster_row['caption_object_id']!r}, got {source_id!r}"
+ )
+ if survey != "legacy_dr10":
+ raise Stage1AValidationError(
+ f"candidate {candidate_index}: survey must be 'legacy_dr10', got {survey!r}"
+ )
+ ra = _check_coordinate(
+ batch.column(2)[local_index].as_py(),
+ roster_row["caption_ra"],
+ field="caption_ra",
+ candidate_index=candidate_index,
+ tolerance=config.coordinate_tolerance_deg,
+ )
+ dec = _check_coordinate(
+ batch.column(3)[local_index].as_py(),
+ roster_row["caption_dec"],
+ field="caption_dec",
+ candidate_index=candidate_index,
+ tolerance=config.coordinate_tolerance_deg,
+ )
+ caption = batch.column(4)[local_index].as_py()
+ summary = batch.column(5)[local_index].as_py()
+ if not isinstance(caption, str) or not caption.strip():
+ raise Stage1AValidationError(
+ f"candidate {candidate_index}: caption is null or empty"
+ )
+ if not isinstance(summary, str) or not summary.strip():
+ raise Stage1AValidationError(
+ f"candidate {candidate_index}: summary is null or empty"
+ )
+ found[candidate_index] = {
+ "ra": ra,
+ "dec": dec,
+ "caption": caption,
+ "summary": summary,
+ }
+ pointer += 1
+ row_offset = batch_end
+ if pointer == len(wanted):
+ break
+ if pointer != len(wanted):
+ raise Stage1AValidationError(f"{shard}: not all caption locators were found")
+
+ if any(item is None for item in found):
+ raise Stage1AValidationError("not every roster candidate produced a caption")
+ return [item for item in found if item is not None]
+
+
+def _angular_separation_arcsec(ra1: float, dec1: float, ra2: float, dec2: float) -> float:
+ ra1_rad, dec1_rad, ra2_rad, dec2_rad = map(math.radians, (ra1, dec1, ra2, dec2))
+ sin_dec = math.sin((dec2_rad - dec1_rad) / 2.0)
+ sin_ra = math.sin((ra2_rad - ra1_rad) / 2.0)
+ haversine = sin_dec**2 + math.cos(dec1_rad) * math.cos(dec2_rad) * sin_ra**2
+ haversine = min(1.0, max(0.0, haversine))
+ return math.degrees(2.0 * math.asin(math.sqrt(haversine))) * 3600.0
+
+
+def _candidate_records(
+ config: Stage1AConfig,
+ rows: list[dict[str, str]],
+ paired_values: list[dict[str, Any]],
+ caption_values: list[dict[str, Any]],
+ paired_sources: dict[str, _SourceFile],
+ caption_sources: dict[str, _SourceFile],
+) -> list[dict[str, Any]]:
+ records = []
+ for index, (row, paired, caption) in enumerate(
+ zip(rows, paired_values, caption_values, strict=False)
+ ):
+ separation = _angular_separation_arcsec(
+ paired["ra"], paired["dec"], caption["ra"], caption["dec"]
+ )
+ recorded_separation = float(row["separation_arcsec"])
+ if separation > 1.0 + config.separation_tolerance_arcsec:
+ raise Stage1AValidationError(
+ f"candidate {index}: source separation exceeds 1 arcsec ({separation})"
+ )
+ if abs(separation - recorded_separation) > config.separation_tolerance_arcsec:
+ raise Stage1AValidationError(
+ f"candidate {index}: separation mismatch; expected {recorded_separation}, "
+ f"got {separation}"
+ )
+
+ paired_source = paired_sources[row["paired_shard"]]
+ caption_source = caption_sources[row["caption_shard"]]
+ records.append(
+ {
+ "schema_version": 1,
+ "candidate_index": index,
+ "paired_targetid": row["paired_targetid"],
+ "caption_object_id": row["caption_object_id"],
+ "coordinates_deg": {
+ "paired_ra": paired["ra"],
+ "paired_dec": paired["dec"],
+ "caption_ra": caption["ra"],
+ "caption_dec": caption["dec"],
+ },
+ "separation_arcsec": separation,
+ "matching": {
+ "paired_candidate_count": 1,
+ "caption_candidate_count": 1,
+ "one_to_one_row_candidate": True,
+ "unique_source_ids": True,
+ "provisional_unambiguous_candidate": True,
+ },
+ "paired_source": {
+ "repo_id": PAIRED_REPO_ID,
+ "revision": config.paired_revision,
+ "shard": row["paired_shard"],
+ "row": int(row["paired_row"]),
+ "sha256": paired_source.sha256,
+ },
+ "caption_source": {
+ "repo_id": CAPTION_REPO_ID,
+ "revision": config.caption_revision,
+ "shard": row["caption_shard"],
+ "row": int(row["caption_row"]),
+ "sha256": caption_source.sha256,
+ "survey": "legacy_dr10",
+ },
+ "image_embedding": {
+ "origin": "precomputed_release",
+ "file": "image_embeddings.f16.npy",
+ "index": index,
+ "dtype": "float16",
+ "shape": list(IMAGE_EMBEDDING_SHAPE),
+ },
+ "caption": {
+ "original": caption["caption"],
+ "summary_original": caption["summary"],
+ "status": "unreviewed",
+ },
+ "acceptance_status": "pending_review",
+ "source_group_id": None,
+ "split": None,
+ }
+ )
+ return records
+
+
+def _write_jsonl(path: Path, records: list[dict[str, Any]]) -> None:
+ with path.open("w", encoding="utf-8", newline="\n") as handle:
+ for record in records:
+ handle.write(json.dumps(record, ensure_ascii=False, sort_keys=True, allow_nan=False))
+ handle.write("\n")
+
+
+def _path_exists(path: Path) -> bool:
+ return path.exists() or path.is_symlink()
+
+
+def _commit_staging(staging: Path, output: Path, *, overwrite: bool) -> None:
+ if not _path_exists(output):
+ os.replace(staging, output)
+ return
+ if not overwrite:
+ raise FileExistsError(f"output already exists: {output}")
+
+ backup = output.with_name(f".{output.name}.backup-{uuid.uuid4().hex}")
+ os.replace(output, backup)
+ try:
+ os.replace(staging, output)
+ except BaseException:
+ os.replace(backup, output)
+ raise
+ if backup.is_dir() and not backup.is_symlink():
+ shutil.rmtree(backup)
+ else:
+ backup.unlink(missing_ok=True)
+
+
+def prepare_stage1a(config: Stage1AConfig) -> dict[str, Any]:
+ """Validate inputs and optionally write the Stage 1A candidate artifact.
+
+ The returned dictionary is a concise summary suitable for CLI output.
+ ``validate_only`` performs the full checksum/schema/content scan without
+ creating or modifying the output directory.
+ """
+
+ if config.batch_size < 1:
+ raise ValueError("batch_size must be positive")
+
+ output_dir = Path(config.output_dir)
+ if not config.validate_only and _path_exists(output_dir) and not config.overwrite:
+ raise FileExistsError(f"output already exists: {output_dir}")
+
+ rows, roster_sha256 = _read_roster(config)
+ paired_sources, caption_sources = _collect_sources(config, rows)
+
+ staging: Path | None = None
+ embedding_output: np.memmap | None = None
+ try:
+ if not config.validate_only:
+ output_dir.parent.mkdir(parents=True, exist_ok=True)
+ staging = Path(
+ tempfile.mkdtemp(prefix=f".{output_dir.name}.tmp-", dir=output_dir.parent)
+ )
+ embedding_output = np.lib.format.open_memmap(
+ staging / "image_embeddings.f16.npy",
+ mode="w+",
+ dtype=np.float16,
+ shape=(len(rows), *IMAGE_EMBEDDING_SHAPE),
+ )
+
+ paired_values, norms = _scan_paired(config, rows, paired_sources, embedding_output)
+ caption_values = _scan_captions(config, rows, caption_sources)
+ records = _candidate_records(
+ config,
+ rows,
+ paired_values,
+ caption_values,
+ paired_sources,
+ caption_sources,
+ )
+
+ if embedding_output is not None:
+ embedding_output.flush()
+ del embedding_output
+ embedding_output = None
+
+ summary = {
+ "status": "validated" if config.validate_only else "prepared",
+ "candidate_count": len(records),
+ "roster_sha256": roster_sha256,
+ "paired_revision": config.paired_revision,
+ "caption_revision": config.caption_revision,
+ "image_embedding_dtype": "float16",
+ "image_embedding_shape": [len(records), *IMAGE_EMBEDDING_SHAPE],
+ "spectra_included": False,
+ }
+ if config.validate_only:
+ return summary
+
+ assert staging is not None
+ candidates_path = staging / "candidates.jsonl"
+ embeddings_path = staging / "image_embeddings.f16.npy"
+ _write_jsonl(candidates_path, records)
+ manifest = {
+ "schema_version": 1,
+ "artifact_type": "astroprism_aion_feature_parity_reference",
+ "role": "parity_reference_only",
+ "created_at_utc": datetime.now(timezone.utc).isoformat(),
+ "pins": {
+ "roster_sha256": roster_sha256,
+ "paired": {"repo_id": PAIRED_REPO_ID, "revision": config.paired_revision},
+ "captions": {
+ "repo_id": CAPTION_REPO_ID,
+ "revision": config.caption_revision,
+ "config": "default",
+ "split": "default",
+ "survey": "legacy_dr10",
+ },
+ },
+ "roster": {
+ "path": str(Path(config.roster_path)),
+ "sha256": roster_sha256,
+ "row_count": len(rows),
+ },
+ "sources": {
+ "paired": [paired_sources[key].manifest_entry() for key in sorted(paired_sources)],
+ "captions": [
+ caption_sources[key].manifest_entry() for key in sorted(caption_sources)
+ ],
+ },
+ "outputs": {
+ "candidates": {
+ "path": candidates_path.name,
+ "row_count": len(records),
+ "size_bytes": candidates_path.stat().st_size,
+ "sha256": sha256_file(candidates_path),
+ },
+ "image_embeddings": {
+ "path": embeddings_path.name,
+ "dtype": "float16",
+ "shape": [len(records), *IMAGE_EMBEDDING_SHAPE],
+ "size_bytes": embeddings_path.stat().st_size,
+ "sha256": sha256_file(embeddings_path),
+ },
+ },
+ "validation": {
+ "candidate_count": len(records),
+ "coordinate_tolerance_deg": config.coordinate_tolerance_deg,
+ "separation_tolerance_arcsec": config.separation_tolerance_arcsec,
+ "image_embedding_l2_norm": {
+ "min": min(norms),
+ "median": float(np.median(norms)),
+ "max": max(norms),
+ },
+ "acceptance_status": "pending_review",
+ },
+ "spectra_included": False,
+ }
+ manifest_path = staging / "manifest.json"
+ manifest_path.write_text(
+ json.dumps(manifest, indent=2, sort_keys=True, allow_nan=False) + "\n",
+ encoding="utf-8",
+ )
+ _commit_staging(staging, output_dir, overwrite=config.overwrite)
+ staging = None
+ summary["output_dir"] = str(output_dir)
+ summary["manifest"] = str(output_dir / "manifest.json")
+ return summary
+ finally:
+ if embedding_output is not None:
+ del embedding_output
+ if staging is not None and staging.exists():
+ shutil.rmtree(staging)
diff --git a/src/data/astroprism_stage1a_collate.py b/src/data/astroprism_stage1a_collate.py
new file mode 100644
index 0000000..95e0759
--- /dev/null
+++ b/src/data/astroprism_stage1a_collate.py
@@ -0,0 +1,128 @@
+"""Lockstep collation for AstroPRISM Stage 1A image/chat samples."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+from torch.nn.utils.rnn import pad_sequence
+
+from .astroprism_stage1a_dataset import (
+ ASTROPRISM_IMAGE_SHAPE,
+ ASTROPRISM_IMAGE_TOKENS,
+ IGNORE_INDEX,
+)
+
+
+class AstroPRISMStage1ACollator:
+ """Pad Stage 1A chat samples without weakening their masking contract.
+
+ The generic multimodal collators use one padding value for every 1-D
+ tensor. Stage 1A requires three distinct values: the tokenizer pad ID for
+ ``text``, zero for ``text_attention_mask``, and ``-100`` for ``labels``.
+ This dedicated collator also rechecks the 576-token image-prefix budget and
+ never truncates an oversized example.
+ """
+
+ def __init__(
+ self,
+ tokenizer: Any,
+ *,
+ max_merged_seq_length: int,
+ image_token_count: int = ASTROPRISM_IMAGE_TOKENS,
+ ignore_index: int = IGNORE_INDEX,
+ ) -> None:
+ pad_token_id = getattr(tokenizer, "pad_token_id", None)
+ if pad_token_id is None:
+ raise ValueError("AstroPRISM Stage 1A requires tokenizer.pad_token_id")
+ if image_token_count != ASTROPRISM_IMAGE_TOKENS:
+ raise ValueError(
+ "AstroPRISM Stage 1A requires exactly "
+ f"{ASTROPRISM_IMAGE_TOKENS} image tokens, got {image_token_count}"
+ )
+ if max_merged_seq_length <= image_token_count:
+ raise ValueError("max_merged_seq_length must exceed the 576-token AION image prefix")
+ self.pad_token_id = int(pad_token_id)
+ self.max_merged_seq_length = int(max_merged_seq_length)
+ self.image_token_count = int(image_token_count)
+ self.ignore_index = int(ignore_index)
+
+ def __call__(self, batch: list[dict[str, Any]]) -> dict[str, Any]:
+ if not batch:
+ return {}
+ required = {
+ "image",
+ "text",
+ "text_attention_mask",
+ "labels",
+ "image_token_count",
+ "merged_length",
+ "provenance",
+ }
+ for index, sample in enumerate(batch):
+ missing = required.difference(sample)
+ if missing:
+ raise ValueError(f"Stage 1A sample {index} is missing keys: {sorted(missing)}")
+ text = sample["text"]
+ labels = sample["labels"]
+ attention = sample["text_attention_mask"]
+ if not all(isinstance(value, torch.Tensor) for value in (text, labels, attention)):
+ raise TypeError(f"Stage 1A sample {index} text/labels/mask must be tensors")
+ if text.ndim != 1 or labels.shape != text.shape or attention.shape != text.shape:
+ raise ValueError(
+ f"Stage 1A sample {index} text, labels, and attention mask must be "
+ "aligned 1-D tensors"
+ )
+ if int(sample["image_token_count"]) != self.image_token_count:
+ raise ValueError(
+ f"Stage 1A sample {index} has image_token_count="
+ f"{sample['image_token_count']}; expected {self.image_token_count}"
+ )
+ merged_length = self.image_token_count + int(text.numel())
+ if int(sample["merged_length"]) != merged_length:
+ raise ValueError(
+ f"Stage 1A sample {index} has inconsistent merged_length="
+ f"{sample['merged_length']}; expected {merged_length}"
+ )
+ if merged_length > self.max_merged_seq_length:
+ raise ValueError(
+ f"Stage 1A sample {index} merged sequence length {merged_length} exceeds "
+ f"max_merged_seq_length={self.max_merged_seq_length}; refusing to truncate"
+ )
+ image = sample["image"]
+ if not isinstance(image, torch.Tensor) or image.shape != ASTROPRISM_IMAGE_SHAPE:
+ raise ValueError(
+ f"Stage 1A sample {index} image must be {ASTROPRISM_IMAGE_SHAPE}, got "
+ f"{getattr(image, 'shape', None)}"
+ )
+ if image.dtype != torch.float32:
+ raise ValueError(
+ f"Stage 1A sample {index} image must be float32, got {image.dtype}"
+ )
+
+ return {
+ "image": torch.stack([sample["image"] for sample in batch]),
+ "text": pad_sequence(
+ [sample["text"] for sample in batch],
+ batch_first=True,
+ padding_value=self.pad_token_id,
+ ),
+ "text_attention_mask": pad_sequence(
+ [sample["text_attention_mask"] for sample in batch],
+ batch_first=True,
+ padding_value=0,
+ ),
+ "labels": pad_sequence(
+ [sample["labels"] for sample in batch],
+ batch_first=True,
+ padding_value=self.ignore_index,
+ ),
+ "image_token_count": self.image_token_count,
+ "merged_length": torch.tensor(
+ [int(sample["merged_length"]) for sample in batch], dtype=torch.long
+ ),
+ "provenance": [sample["provenance"] for sample in batch],
+ }
+
+
+__all__ = ["AstroPRISMStage1ACollator"]
diff --git a/src/data/astroprism_stage1a_dataset.py b/src/data/astroprism_stage1a_dataset.py
new file mode 100644
index 0000000..726df13
--- /dev/null
+++ b/src/data/astroprism_stage1a_dataset.py
@@ -0,0 +1,588 @@
+"""Dataset contract for AstroPRISM Stage 1A raw-image alignment.
+
+The dataset intentionally consumes the calibrated raw-image artifact rather
+than the released AION embeddings. The latter remain a parity reference and
+are only used here as the source of the matched caption records.
+
+Each sample contains one ``float32[4, 160, 160]`` Legacy Survey image and a
+Qwen chat sequence. Only tokens belonging to the final assistant response are
+supervised; the system/user prefix is masked with ``-100``. The AION image
+observation is supplied out of band under ``image`` and expands to 576 tokens
+in the model, so it is included in the merged-length safety check.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+ASTROPRISM_IMAGE_SHAPE = (4, 160, 160)
+ASTROPRISM_IMAGE_TOKENS = 576
+ASTROPRISM_SOURCE_BANDS = ("des-g", "des-r", "des-i", "des-z")
+ASTROPRISM_AION_BANDS = ("DES-G", "DES-R", "DES-I", "DES-Z")
+IGNORE_INDEX = -100
+
+DEFAULT_SYSTEM_PROMPT = (
+ "You are an astronomy assistant. Describe only morphology and visual evidence "
+ "supported by the supplied calibrated four-band galaxy image."
+)
+DEFAULT_USER_PROMPT = "Describe the galaxy in this calibrated four-band image."
+
+_CAPTION_FIELD_BY_RELEASE_STATUS = {
+ "released_reused": "original",
+ "released_revised": "revised",
+}
+
+
+class AstroPRISMStage1ADatasetError(ValueError):
+ """Raised when a Stage 1A artifact or sample violates its contract."""
+
+
+@dataclass(frozen=True)
+class AstroPRISMStage1ADatasetConfig:
+ """Paths and safety settings for :class:`AstroPRISMStage1ADataset`.
+
+ ``allow_unsafe_candidates`` exists for unit tests and data inspection. A
+ training launch should leave it false so pending/unreviewed candidates and
+ records without a source-level split cannot enter a batch.
+ """
+
+ raw_manifest_path: Path
+ candidates_manifest_path: Path
+ candidates_path: Path | None = None
+ split: str | None = None
+ max_merged_seq_length: int = 2048
+ caption_field: str = "auto"
+ system_prompt: str = DEFAULT_SYSTEM_PROMPT
+ user_prompt: str = DEFAULT_USER_PROMPT
+ allow_unsafe_candidates: bool = False
+ accepted_statuses: tuple[str, ...] = ("accepted", "approved")
+ accepted_caption_statuses: tuple[str, ...] = (
+ "released_reused",
+ "released_revised",
+ )
+ verify_file_hashes: bool = False
+
+
+def _read_json(path: Path, *, kind: str) -> dict[str, Any]:
+ if not path.is_file():
+ raise AstroPRISMStage1ADatasetError(f"{kind} does not exist: {path}")
+ try:
+ value = json.loads(path.read_text(encoding="utf-8"))
+ except (OSError, json.JSONDecodeError) as exc:
+ raise AstroPRISMStage1ADatasetError(f"cannot read {kind} {path}: {exc}") from exc
+ if not isinstance(value, dict):
+ raise AstroPRISMStage1ADatasetError(f"{kind} must contain a JSON object: {path}")
+ return value
+
+
+def _read_jsonl(path: Path) -> list[dict[str, Any]]:
+ if not path.is_file():
+ raise AstroPRISMStage1ADatasetError(f"candidate JSONL does not exist: {path}")
+ records: list[dict[str, Any]] = []
+ try:
+ with path.open("r", encoding="utf-8") as handle:
+ for line_number, line in enumerate(handle, start=1):
+ if not line.strip():
+ raise AstroPRISMStage1ADatasetError(
+ f"candidate JSONL contains a blank row at line {line_number}"
+ )
+ value = json.loads(line)
+ if not isinstance(value, dict):
+ raise AstroPRISMStage1ADatasetError(
+ f"candidate JSONL line {line_number} must be an object"
+ )
+ records.append(value)
+ except json.JSONDecodeError as exc:
+ raise AstroPRISMStage1ADatasetError(
+ f"cannot parse candidate JSONL {path} at line {exc.lineno}: {exc.msg}"
+ ) from exc
+ return records
+
+
+def _sha256_file(path: Path, chunk_size: int = 8 * 1024 * 1024) -> str:
+ digest = hashlib.sha256()
+ with path.open("rb") as handle:
+ while chunk := handle.read(chunk_size):
+ digest.update(chunk)
+ return digest.hexdigest()
+
+
+def _resolve_output_path(manifest_path: Path, output: dict[str, Any], *, name: str) -> Path:
+ relative = output.get("path")
+ if not isinstance(relative, str) or not relative:
+ raise AstroPRISMStage1ADatasetError(f"manifest output {name!r} has no path")
+ path = Path(relative)
+ return path if path.is_absolute() else manifest_path.parent / path
+
+
+def _coerce_token_ids(value: Any, *, context: str) -> list[int]:
+ """Normalize common ``apply_chat_template`` return forms to one ID list."""
+
+ if isinstance(value, dict):
+ if "input_ids" not in value:
+ raise AstroPRISMStage1ADatasetError(
+ f"tokenizer returned a mapping without input_ids for {context}"
+ )
+ value = value["input_ids"]
+ elif hasattr(value, "input_ids"):
+ value = value.input_ids
+
+ if isinstance(value, torch.Tensor):
+ value = value.detach().cpu().tolist()
+ elif isinstance(value, np.ndarray):
+ value = value.tolist()
+
+ if isinstance(value, tuple):
+ value = list(value)
+ if isinstance(value, list) and len(value) == 1 and isinstance(value[0], list | tuple):
+ value = list(value[0])
+ if not isinstance(value, list) or any(isinstance(item, list | tuple) for item in value):
+ raise AstroPRISMStage1ADatasetError(
+ f"tokenizer must return one flat token sequence for {context}"
+ )
+ try:
+ token_ids = [int(item) for item in value]
+ except (TypeError, ValueError) as exc:
+ raise AstroPRISMStage1ADatasetError(
+ f"tokenizer returned non-integer token IDs for {context}"
+ ) from exc
+ if not token_ids:
+ raise AstroPRISMStage1ADatasetError(f"tokenizer returned no token IDs for {context}")
+ return token_ids
+
+
+def _source_split_keys(record: dict[str, Any]) -> tuple[tuple[str, str], ...]:
+ """Return every stable source identity whose split must remain atomic."""
+
+ keys: list[tuple[str, str]] = []
+ for field in ("source_group_id", "paired_targetid", "caption_object_id"):
+ value = record.get(field)
+ if value is not None and str(value).strip():
+ keys.append((field, str(value)))
+ return tuple(keys)
+
+
+def _validate_source_split_integrity(records: list[dict[str, Any]]) -> None:
+ assignments: dict[tuple[str, str], str] = {}
+ for record in records:
+ split = record.get("split")
+ if split is None or not str(split).strip():
+ continue
+ split = str(split)
+ for source_key in _source_split_keys(record):
+ previous = assignments.setdefault(source_key, split)
+ if previous != split:
+ field, source_id = source_key
+ raise AstroPRISMStage1ADatasetError(
+ f"source split leakage: {field}={source_id!r} appears in both "
+ f"{previous!r} and {split!r}"
+ )
+
+
+class AstroPRISMStage1ADataset(Dataset):
+ """Map-style Stage 1A dataset backed by mmap-safe calibrated cutouts."""
+
+ def __init__(
+ self,
+ config: AstroPRISMStage1ADatasetConfig,
+ tokenizer: Any,
+ ) -> None:
+ self.config = config
+ self.active_modalities = {"image"}
+ self.tokenizer = tokenizer
+
+ if config.max_merged_seq_length <= ASTROPRISM_IMAGE_TOKENS:
+ raise AstroPRISMStage1ADatasetError(
+ "max_merged_seq_length must exceed the fixed 576-token AION image prefix"
+ )
+ if not config.caption_field.strip():
+ raise AstroPRISMStage1ADatasetError("caption_field must be non-empty")
+ if not hasattr(tokenizer, "apply_chat_template"):
+ raise AstroPRISMStage1ADatasetError(
+ "tokenizer must provide apply_chat_template for Qwen chat formatting"
+ )
+
+ self.raw_manifest_path = Path(config.raw_manifest_path)
+ self.candidates_manifest_path = Path(config.candidates_manifest_path)
+ self.raw_manifest = _read_json(self.raw_manifest_path, kind="raw-image manifest")
+ self.candidates_manifest = _read_json(
+ self.candidates_manifest_path, kind="candidate manifest"
+ )
+ self._validate_manifest_types()
+
+ candidate_output = self.candidates_manifest.get("outputs", {}).get("candidates")
+ if not isinstance(candidate_output, dict):
+ raise AstroPRISMStage1ADatasetError("candidate manifest is missing outputs.candidates")
+ self.candidates_path = (
+ Path(config.candidates_path)
+ if config.candidates_path is not None
+ else _resolve_output_path(
+ self.candidates_manifest_path, candidate_output, name="candidates"
+ )
+ )
+ self._all_records = _read_jsonl(self.candidates_path)
+ self._validate_candidate_file(candidate_output)
+
+ image_output = self.raw_manifest.get("outputs", {}).get("image_array")
+ if not isinstance(image_output, dict):
+ raise AstroPRISMStage1ADatasetError("raw-image manifest is missing outputs.image_array")
+ self.image_path = _resolve_output_path(
+ self.raw_manifest_path, image_output, name="image_array"
+ )
+ self._images = self._open_and_validate_images(image_output)
+
+ self._validate_roster_identity()
+ self._validate_candidate_alignment()
+ _validate_source_split_integrity(self._all_records)
+
+ if config.split is None:
+ self._indices = list(range(len(self._all_records)))
+ else:
+ self._indices = [
+ index
+ for index, record in enumerate(self._all_records)
+ if record.get("split") == config.split
+ ]
+ if not self._indices:
+ raise AstroPRISMStage1ADatasetError(
+ f"no candidates found for requested split {config.split!r}"
+ )
+ self._validate_training_safety()
+
+ def _validate_manifest_types(self) -> None:
+ if self.raw_manifest.get("artifact_type") != "astroprism_raw_images":
+ raise AstroPRISMStage1ADatasetError(
+ "raw manifest artifact_type must be 'astroprism_raw_images'"
+ )
+ candidate_type = self.candidates_manifest.get("artifact_type")
+ if candidate_type not in {
+ "astroprism_stage1a_candidates",
+ "astroprism_aion_feature_parity_reference",
+ }:
+ raise AstroPRISMStage1ADatasetError(
+ f"candidate manifest has an unsupported artifact_type: {candidate_type!r}"
+ )
+
+ def _validate_candidate_file(self, output: dict[str, Any]) -> None:
+ declared_count = output.get("row_count")
+ if declared_count is not None and declared_count != len(self._all_records):
+ raise AstroPRISMStage1ADatasetError(
+ f"candidate row count mismatch: manifest declares {declared_count}, "
+ f"JSONL contains {len(self._all_records)}"
+ )
+ if not self._all_records:
+ raise AstroPRISMStage1ADatasetError("candidate JSONL is empty")
+ if self.config.verify_file_hashes and output.get("sha256"):
+ actual = _sha256_file(self.candidates_path)
+ if actual != output["sha256"]:
+ raise AstroPRISMStage1ADatasetError(
+ f"candidate JSONL SHA-256 mismatch: expected {output['sha256']}, got {actual}"
+ )
+
+ def _open_and_validate_images(self, output: dict[str, Any]) -> np.ndarray:
+ if not self.image_path.is_file():
+ raise AstroPRISMStage1ADatasetError(
+ f"raw image array does not exist: {self.image_path}"
+ )
+ try:
+ images = np.load(self.image_path, mmap_mode="r", allow_pickle=False)
+ except (OSError, ValueError) as exc:
+ raise AstroPRISMStage1ADatasetError(
+ f"cannot mmap raw image array {self.image_path}: {exc}"
+ ) from exc
+
+ expected_shape = (len(self._all_records), *ASTROPRISM_IMAGE_SHAPE)
+ if images.shape != expected_shape:
+ raise AstroPRISMStage1ADatasetError(
+ f"raw image array must have shape {expected_shape}, got {images.shape}"
+ )
+ if images.dtype != np.dtype("float32"):
+ raise AstroPRISMStage1ADatasetError(
+ f"raw image array must be float32, got {images.dtype}"
+ )
+ declared_shape = output.get("shape")
+ if declared_shape is not None and list(images.shape) != declared_shape:
+ raise AstroPRISMStage1ADatasetError(
+ f"raw image shape disagrees with manifest: {images.shape} != {declared_shape}"
+ )
+ declared_dtype = output.get("dtype")
+ if declared_dtype is not None and np.dtype(declared_dtype) != images.dtype:
+ raise AstroPRISMStage1ADatasetError(
+ f"raw image dtype disagrees with manifest: {images.dtype} != {declared_dtype}"
+ )
+ if self.config.verify_file_hashes and output.get("sha256"):
+ actual = _sha256_file(self.image_path)
+ if actual != output["sha256"]:
+ raise AstroPRISMStage1ADatasetError(
+ f"raw image SHA-256 mismatch: expected {output['sha256']}, got {actual}"
+ )
+ return images
+
+ def _validate_roster_identity(self) -> None:
+ raw_roster = self.raw_manifest.get("roster", {})
+ candidate_pin = self.candidates_manifest.get("pins", {}).get("roster_sha256")
+ raw_sha = raw_roster.get("sha256") if isinstance(raw_roster, dict) else None
+ if raw_sha is not None and candidate_pin is not None and raw_sha != candidate_pin:
+ raise AstroPRISMStage1ADatasetError(
+ f"raw/candidate roster SHA-256 mismatch: {raw_sha} != {candidate_pin}"
+ )
+
+ def _validate_candidate_alignment(self) -> None:
+ raw_candidates = self.raw_manifest.get("candidates")
+ if not isinstance(raw_candidates, list):
+ raise AstroPRISMStage1ADatasetError("raw manifest is missing candidates list")
+ if len(raw_candidates) != len(self._all_records):
+ raise AstroPRISMStage1ADatasetError(
+ f"raw/caption candidate count mismatch: {len(raw_candidates)} != "
+ f"{len(self._all_records)}"
+ )
+
+ seen_indices: set[int] = set()
+ for index, (raw_record, caption_record) in enumerate(
+ zip(raw_candidates, self._all_records, strict=True)
+ ):
+ for kind, record in (("raw", raw_record), ("caption", caption_record)):
+ if not isinstance(record, dict):
+ raise AstroPRISMStage1ADatasetError(
+ f"{kind} candidate at row {index} must be an object"
+ )
+ if record.get("candidate_index") != index:
+ raise AstroPRISMStage1ADatasetError(
+ f"{kind} candidate row {index} has candidate_index="
+ f"{record.get('candidate_index')!r}; roster order must be preserved"
+ )
+ if index in seen_indices:
+ raise AstroPRISMStage1ADatasetError(f"duplicate candidate_index {index}")
+ bands = raw_record.get("bands")
+ if not isinstance(bands, dict):
+ raise AstroPRISMStage1ADatasetError(
+ f"raw candidate {index} is missing verified band metadata"
+ )
+ source_bands = tuple(bands.get("source", ()))
+ aion_bands = tuple(bands.get("aion", ()))
+ if source_bands != ASTROPRISM_SOURCE_BANDS or aion_bands != ASTROPRISM_AION_BANDS:
+ raise AstroPRISMStage1ADatasetError(
+ f"raw candidate {index} has unsupported band order: "
+ f"source={source_bands!r}, aion={aion_bands!r}"
+ )
+ seen_indices.add(index)
+ for field in ("paired_targetid", "caption_object_id"):
+ if str(raw_record.get(field)) != str(caption_record.get(field)):
+ raise AstroPRISMStage1ADatasetError(
+ f"candidate {index} {field} mismatch between raw and caption artifacts: "
+ f"{raw_record.get(field)!r} != {caption_record.get(field)!r}"
+ )
+
+ def _validate_training_safety(self) -> None:
+ """Reject unsafe records that would actually enter this dataset view.
+
+ Split-integrity checks run over the complete manifest, but review gates
+ apply to ``self._indices``. This permits a provenance-complete manifest
+ to retain explicitly rejected or unassigned rows while a split-specific
+ view consumes only accepted records. A ``split=None`` view remains strict
+ because every record would enter the dataset.
+ """
+
+ if self.config.allow_unsafe_candidates:
+ for index in self._indices:
+ self._caption_text(self._all_records[index])
+ return
+ accepted = {status.casefold() for status in self.config.accepted_statuses}
+ accepted_captions = {status.casefold() for status in self.config.accepted_caption_statuses}
+ if not accepted:
+ raise AstroPRISMStage1ADatasetError("accepted_statuses must be non-empty")
+ if not accepted_captions:
+ raise AstroPRISMStage1ADatasetError("accepted_caption_statuses must be non-empty")
+
+ for index in self._indices:
+ record = self._all_records[index]
+ status = str(record.get("acceptance_status", "")).strip()
+ if status.casefold() not in accepted:
+ raise AstroPRISMStage1ADatasetError(
+ f"candidate {index} has acceptance_status={status!r}; expected one of "
+ f"{sorted(self.config.accepted_statuses)!r}. Pass "
+ "allow_unsafe_candidates=True only for inspection/testing."
+ )
+ split = record.get("split")
+ if split is None or not str(split).strip():
+ raise AstroPRISMStage1ADatasetError(
+ f"candidate {index} has split=null; source-level splits are required. "
+ "Pass allow_unsafe_candidates=True only for inspection/testing."
+ )
+ caption = record.get("caption")
+ caption_status = (
+ str(caption.get("status", "")).strip() if isinstance(caption, dict) else ""
+ )
+ if caption_status.casefold() not in accepted_captions:
+ raise AstroPRISMStage1ADatasetError(
+ f"candidate {index} has caption.status={caption_status!r}; expected one "
+ f"of {sorted(self.config.accepted_caption_statuses)!r}. Stage 1A only "
+ "trains on captions reviewed against the exact AION crop. Pass "
+ "allow_unsafe_candidates=True only for inspection/testing."
+ )
+ self._caption_text(record)
+
+ def __len__(self) -> int:
+ return len(self._indices)
+
+ def _caption_text(self, record: dict[str, Any]) -> tuple[str, str]:
+ caption = record.get("caption")
+ if not isinstance(caption, dict):
+ raise AstroPRISMStage1ADatasetError(
+ f"candidate {record.get('candidate_index')} has no caption object"
+ )
+ requested_field = self.config.caption_field.strip()
+ caption_status = str(caption.get("status", "")).strip().casefold()
+ expected_field = _CAPTION_FIELD_BY_RELEASE_STATUS.get(caption_status)
+
+ if requested_field.casefold() == "auto":
+ if expected_field is not None:
+ selected_field = expected_field
+ elif self.config.allow_unsafe_candidates:
+ selected_field = "original"
+ else:
+ raise AstroPRISMStage1ADatasetError(
+ f"candidate {record.get('candidate_index')} has caption.status="
+ f"{caption.get('status')!r}, which has no safe automatic caption mapping"
+ )
+ else:
+ selected_field = requested_field
+ if expected_field is not None and selected_field != expected_field:
+ raise AstroPRISMStage1ADatasetError(
+ f"candidate {record.get('candidate_index')} caption.status="
+ f"{caption.get('status')!r} requires caption.{expected_field}, not "
+ f"caption.{selected_field}; use caption_field='auto'"
+ )
+ if expected_field is None and not self.config.allow_unsafe_candidates:
+ raise AstroPRISMStage1ADatasetError(
+ f"candidate {record.get('candidate_index')} has caption.status="
+ f"{caption.get('status')!r}, which has no safe caption mapping; "
+ "explicit fields are permitted only with "
+ "allow_unsafe_candidates=True for inspection/testing"
+ )
+
+ value = caption.get(selected_field)
+ if not isinstance(value, str) or not value.strip():
+ raise AstroPRISMStage1ADatasetError(
+ f"candidate {record.get('candidate_index')} has no non-empty "
+ f"caption.{selected_field} selected from caption.status="
+ f"{caption.get('status')!r}"
+ )
+ return value.strip(), selected_field
+
+ def _tokenize_chat(self, target: str) -> tuple[torch.Tensor, torch.Tensor]:
+ prefix_messages = [
+ {"role": "system", "content": self.config.system_prompt},
+ {"role": "user", "content": self.config.user_prompt},
+ ]
+ full_messages = [*prefix_messages, {"role": "assistant", "content": target}]
+
+ common = {"tokenize": True, "enable_thinking": False}
+ try:
+ prompt_value = self.tokenizer.apply_chat_template(
+ prefix_messages,
+ add_generation_prompt=True,
+ **common,
+ )
+ full_value = self.tokenizer.apply_chat_template(
+ full_messages,
+ add_generation_prompt=False,
+ **common,
+ )
+ except (TypeError, ValueError) as exc:
+ raise AstroPRISMStage1ADatasetError(
+ "tokenizer failed Qwen non-thinking chat formatting "
+ "(apply_chat_template(..., enable_thinking=False))"
+ ) from exc
+
+ prompt_ids = _coerce_token_ids(prompt_value, context="chat prompt")
+ full_ids = _coerce_token_ids(full_value, context="full chat")
+ if len(full_ids) <= len(prompt_ids) or full_ids[: len(prompt_ids)] != prompt_ids:
+ raise AstroPRISMStage1ADatasetError(
+ "Qwen full chat is not prefixed by the generation prompt; cannot construct "
+ "assistant-only labels safely"
+ )
+
+ merged_length = ASTROPRISM_IMAGE_TOKENS + len(full_ids)
+ if merged_length > self.config.max_merged_seq_length:
+ raise AstroPRISMStage1ADatasetError(
+ f"merged sequence length {merged_length} exceeds "
+ f"max_merged_seq_length={self.config.max_merged_seq_length} "
+ f"({ASTROPRISM_IMAGE_TOKENS} image tokens + {len(full_ids)} text tokens); "
+ "the dataset never truncates captions silently"
+ )
+
+ input_ids = torch.tensor(full_ids, dtype=torch.long)
+ labels = input_ids.clone()
+ labels[: len(prompt_ids)] = IGNORE_INDEX
+ if not torch.any(labels != IGNORE_INDEX):
+ raise AstroPRISMStage1ADatasetError("assistant response produced no supervised tokens")
+ return input_ids, labels
+
+ def __getitem__(self, item: int) -> dict[str, Any]:
+ source_index = self._indices[item]
+ record = self._all_records[source_index]
+ raw_record = self.raw_manifest["candidates"][source_index]
+
+ # np.array(copy=True) avoids returning a writable torch view over a
+ # read-only memmap and isolates AION's in-place preprocessing.
+ image_array = np.array(self._images[source_index], dtype=np.float32, order="C", copy=True)
+ if not np.isfinite(image_array).all():
+ raise AstroPRISMStage1ADatasetError(
+ f"candidate {source_index} raw image contains non-finite values"
+ )
+ image = torch.from_numpy(image_array)
+
+ caption_text, selected_caption_field = self._caption_text(record)
+ input_ids, labels = self._tokenize_chat(caption_text)
+ attention_mask = torch.ones_like(input_ids, dtype=torch.long)
+
+ provenance = {
+ "candidate_index": source_index,
+ "paired_targetid": str(record["paired_targetid"]),
+ "caption_object_id": str(record["caption_object_id"]),
+ "source_group_id": record.get("source_group_id"),
+ "split": record.get("split"),
+ "acceptance_status": record.get("acceptance_status"),
+ "raw_row": raw_record.get("raw_row"),
+ "bands": raw_record.get("bands"),
+ "coordinates_deg": raw_record.get("coordinates_deg"),
+ "raw_source": self.raw_manifest.get("source"),
+ "caption_source": record.get("caption_source"),
+ "paired_source": record.get("paired_source"),
+ "roster_sha256": self.candidates_manifest.get("pins", {}).get("roster_sha256"),
+ "caption_field": selected_caption_field,
+ "caption_field_config": self.config.caption_field,
+ "caption_status": record.get("caption", {}).get("status"),
+ "chat_template": "qwen_non_thinking",
+ }
+
+ return {
+ "image": image,
+ "text": input_ids,
+ "text_attention_mask": attention_mask,
+ "labels": labels,
+ "image_token_count": ASTROPRISM_IMAGE_TOKENS,
+ "merged_length": ASTROPRISM_IMAGE_TOKENS + input_ids.numel(),
+ "provenance": provenance,
+ }
+
+
+__all__ = [
+ "ASTROPRISM_AION_BANDS",
+ "ASTROPRISM_IMAGE_SHAPE",
+ "ASTROPRISM_IMAGE_TOKENS",
+ "ASTROPRISM_SOURCE_BANDS",
+ "IGNORE_INDEX",
+ "AstroPRISMStage1ADataset",
+ "AstroPRISMStage1ADatasetConfig",
+ "AstroPRISMStage1ADatasetError",
+]
diff --git a/src/data/astroprism_stage1b_spectrum_collate.py b/src/data/astroprism_stage1b_spectrum_collate.py
new file mode 100644
index 0000000..5fe4315
--- /dev/null
+++ b/src/data/astroprism_stage1b_spectrum_collate.py
@@ -0,0 +1,145 @@
+"""Lockstep collation for AstroPRISM raw-spectrum/chat samples."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+from torch.nn.utils.rnn import pad_sequence
+
+from .astroprism_spectrum_contract import ASTROPRISM_SPECTRUM_LENGTH
+from .astroprism_stage1a_dataset import IGNORE_INDEX
+from .astroprism_stage1b_spectrum_dataset import ASTROPRISM_SPECTRUM_TOKENS
+
+
+class AstroPRISMStage1BSpectrumCollator:
+ """Stack four raw DESI fields and pad assistant-only chat labels."""
+
+ def __init__(
+ self,
+ tokenizer: Any,
+ *,
+ max_merged_seq_length: int,
+ spectrum_token_count: int = ASTROPRISM_SPECTRUM_TOKENS,
+ ignore_index: int = IGNORE_INDEX,
+ ) -> None:
+ pad_token_id = getattr(tokenizer, "pad_token_id", None)
+ if pad_token_id is None:
+ raise ValueError("AstroPRISM spectrum alignment requires tokenizer.pad_token_id")
+ if spectrum_token_count != ASTROPRISM_SPECTRUM_TOKENS:
+ raise ValueError(
+ "AstroPRISM spectrum alignment requires exactly "
+ f"{ASTROPRISM_SPECTRUM_TOKENS} AION tokens, got {spectrum_token_count}"
+ )
+ if max_merged_seq_length <= spectrum_token_count:
+ raise ValueError(
+ "max_merged_seq_length must exceed the 273-token AION spectrum prefix"
+ )
+ self.pad_token_id = int(pad_token_id)
+ self.max_merged_seq_length = int(max_merged_seq_length)
+ self.spectrum_token_count = int(spectrum_token_count)
+ self.ignore_index = int(ignore_index)
+
+ def __call__(self, batch: list[dict[str, Any]]) -> dict[str, Any]:
+ if not batch:
+ return {}
+ required = {
+ "spectrum",
+ "text",
+ "text_attention_mask",
+ "labels",
+ "spectrum_token_count",
+ "merged_length",
+ "provenance",
+ }
+ spectrum_fields = {
+ "flux": torch.float32,
+ "ivar": torch.float32,
+ "mask": torch.bool,
+ "wavelength": torch.float32,
+ }
+ for index, sample in enumerate(batch):
+ missing = required.difference(sample)
+ if missing:
+ raise ValueError(
+ f"spectrum alignment sample {index} is missing keys: {sorted(missing)}"
+ )
+ text = sample["text"]
+ labels = sample["labels"]
+ attention = sample["text_attention_mask"]
+ if not all(isinstance(value, torch.Tensor) for value in (text, labels, attention)):
+ raise TypeError(
+ f"spectrum alignment sample {index} text/labels/mask must be tensors"
+ )
+ if text.ndim != 1 or labels.shape != text.shape or attention.shape != text.shape:
+ raise ValueError(
+ f"spectrum alignment sample {index} text, labels, and attention mask "
+ "must be aligned 1-D tensors"
+ )
+ if int(sample["spectrum_token_count"]) != self.spectrum_token_count:
+ raise ValueError(
+ f"spectrum alignment sample {index} has spectrum_token_count="
+ f"{sample['spectrum_token_count']}; expected {self.spectrum_token_count}"
+ )
+ merged_length = self.spectrum_token_count + int(text.numel())
+ if int(sample["merged_length"]) != merged_length:
+ raise ValueError(
+ f"spectrum alignment sample {index} has inconsistent merged_length="
+ f"{sample['merged_length']}; expected {merged_length}"
+ )
+ if merged_length > self.max_merged_seq_length:
+ raise ValueError(
+ f"spectrum alignment sample {index} merged length {merged_length} "
+ f"exceeds max_merged_seq_length={self.max_merged_seq_length}"
+ )
+ spectrum = sample["spectrum"]
+ if not isinstance(spectrum, dict) or set(spectrum) != set(spectrum_fields):
+ raise ValueError(
+ f"spectrum alignment sample {index} must contain exactly "
+ f"{sorted(spectrum_fields)}"
+ )
+ for field, dtype in spectrum_fields.items():
+ value = spectrum[field]
+ if not isinstance(value, torch.Tensor):
+ raise TypeError(f"sample {index} spectrum.{field} must be a tensor")
+ if value.shape != (ASTROPRISM_SPECTRUM_LENGTH,):
+ raise ValueError(
+ f"sample {index} spectrum.{field} must have shape "
+ f"({ASTROPRISM_SPECTRUM_LENGTH},), got {tuple(value.shape)}"
+ )
+ if value.dtype != dtype:
+ raise ValueError(
+ f"sample {index} spectrum.{field} must have dtype {dtype}, "
+ f"got {value.dtype}"
+ )
+
+ return {
+ "spectrum": {
+ field: torch.stack([sample["spectrum"][field] for sample in batch])
+ for field in spectrum_fields
+ },
+ "text": pad_sequence(
+ [sample["text"] for sample in batch],
+ batch_first=True,
+ padding_value=self.pad_token_id,
+ ),
+ "text_attention_mask": pad_sequence(
+ [sample["text_attention_mask"] for sample in batch],
+ batch_first=True,
+ padding_value=0,
+ ),
+ "labels": pad_sequence(
+ [sample["labels"] for sample in batch],
+ batch_first=True,
+ padding_value=self.ignore_index,
+ ),
+ "spectrum_token_count": self.spectrum_token_count,
+ "merged_length": torch.tensor(
+ [int(sample["merged_length"]) for sample in batch],
+ dtype=torch.long,
+ ),
+ "provenance": [sample["provenance"] for sample in batch],
+ }
+
+
+__all__ = ["AstroPRISMStage1BSpectrumCollator"]
diff --git a/src/data/astroprism_stage1b_spectrum_dataset.py b/src/data/astroprism_stage1b_spectrum_dataset.py
new file mode 100644
index 0000000..329559e
--- /dev/null
+++ b/src/data/astroprism_stage1b_spectrum_dataset.py
@@ -0,0 +1,569 @@
+"""Spectrum-only AstroPRISM alignment dataset.
+
+Stage 1B pairs raw observed-frame DESI spectra with the promoted, reviewed
+spectrum-caption artifact. Source identities and train/validation/test splits
+remain tied to the immutable 646-object overlap roster. The loader fails
+closed unless the selected records carry the accepted review decision and the
+split-specific training/evaluation policy produced by caption promotion.
+"""
+
+from __future__ import annotations
+
+from collections import Counter
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+
+from .astroprism_spectrum_contract import (
+ ASTROPRISM_SPECTRUM_LENGTH,
+ ASTROPRISM_SPECTRUM_TOKENS,
+)
+from .astroprism_stage1a_dataset import (
+ IGNORE_INDEX,
+ AstroPRISMStage1ADataset,
+ AstroPRISMStage1ADatasetError,
+ _coerce_token_ids,
+ _read_json,
+ _read_jsonl,
+ _resolve_output_path,
+ _sha256_file,
+ _validate_source_split_integrity,
+)
+
+DEFAULT_SPECTRUM_SYSTEM_PROMPT = (
+ "You are an astronomy assistant. Give a concise, objective description grounded "
+ "only in the supplied observed-frame DESI spectrum. Do not infer image morphology "
+ "or unsupported physical causes."
+)
+DEFAULT_SPECTRUM_USER_PROMPT = "Describe the supplied observed-frame DESI galaxy spectrum."
+
+_ROLE_BY_SPLIT = {
+ "train": "alignment_training",
+ "validation": "alignment_validation",
+ "test": "alignment_test",
+}
+
+
+@dataclass(frozen=True)
+class AstroPRISMStage1BSpectrumDatasetConfig:
+ raw_manifest_path: Path
+ captions_manifest_path: Path
+ captions_path: Path | None = None
+ split: str | None = "train"
+ max_merged_seq_length: int = 2048
+ system_prompt: str = DEFAULT_SPECTRUM_SYSTEM_PROMPT
+ user_prompt: str = DEFAULT_SPECTRUM_USER_PROMPT
+ accepted_caption_statuses: tuple[str, ...] = ("accepted",)
+ verify_file_hashes: bool = False
+
+
+class AstroPRISMStage1BSpectrumDataset(AstroPRISMStage1ADataset):
+ """Mmap-backed raw-spectrum dataset with assistant-only caption labels."""
+
+ def __init__(
+ self,
+ config: AstroPRISMStage1BSpectrumDatasetConfig,
+ tokenizer: Any,
+ ) -> None:
+ # Do not call the image dataset initializer; only its Dataset base and
+ # shared exception/helper contract are reused.
+ self.config = config
+ self.active_modalities = {"spectrum"}
+ self.tokenizer = tokenizer
+
+ if config.max_merged_seq_length <= ASTROPRISM_SPECTRUM_TOKENS:
+ raise AstroPRISMStage1ADatasetError(
+ "max_merged_seq_length must exceed the fixed 273-token AION spectrum prefix"
+ )
+ if not hasattr(tokenizer, "apply_chat_template"):
+ raise AstroPRISMStage1ADatasetError(
+ "tokenizer must provide apply_chat_template for Qwen chat formatting"
+ )
+
+ self.raw_manifest_path = Path(config.raw_manifest_path)
+ self.captions_manifest_path = Path(config.captions_manifest_path)
+ self.raw_manifest = _read_json(
+ self.raw_manifest_path,
+ kind="raw-spectrum manifest",
+ )
+ self.captions_manifest = _read_json(
+ self.captions_manifest_path,
+ kind="reviewed spectrum-caption manifest",
+ )
+ self._validate_manifest_types()
+
+ captions_output = self.captions_manifest.get("outputs", {}).get("captions.jsonl")
+ if not isinstance(captions_output, dict):
+ raise AstroPRISMStage1ADatasetError(
+ "spectrum-caption manifest is missing outputs['captions.jsonl']"
+ )
+ self.captions_path = (
+ Path(config.captions_path)
+ if config.captions_path is not None
+ else self.captions_manifest_path.parent / "captions.jsonl"
+ )
+ self._all_records = _read_jsonl(self.captions_path)
+ self._validate_caption_file(captions_output)
+
+ self._spectra: dict[str, np.ndarray] = {}
+ for key, dtype in (
+ ("spectrum_flux", np.dtype("float32")),
+ ("spectrum_ivar", np.dtype("float32")),
+ ("spectrum_mask", np.dtype("bool")),
+ ("spectrum_wavelength", np.dtype("float32")),
+ ):
+ output = self.raw_manifest.get("outputs", {}).get(key)
+ if not isinstance(output, dict):
+ raise AstroPRISMStage1ADatasetError(
+ f"raw-spectrum manifest is missing outputs.{key}"
+ )
+ self._spectra[key] = self._open_and_validate_array(key, output, dtype)
+
+ self._validate_roster_identity()
+ self._validate_caption_alignment()
+ _validate_source_split_integrity(self._all_records)
+
+ if config.split is None:
+ self._indices = list(range(len(self._all_records)))
+ else:
+ self._indices = [
+ index
+ for index, record in enumerate(self._all_records)
+ if record.get("split") == config.split
+ ]
+ if not self._indices:
+ raise AstroPRISMStage1ADatasetError(
+ f"no spectrum captions found for requested split {config.split!r}"
+ )
+ self._validate_training_safety()
+
+ def _validate_manifest_types(self) -> None:
+ if self.raw_manifest.get("artifact_type") != "astroprism_raw_spectra":
+ raise AstroPRISMStage1ADatasetError(
+ "raw manifest artifact_type must be 'astroprism_raw_spectra'"
+ )
+ spectrum_contract = self.raw_manifest.get("spectrum")
+ if not isinstance(spectrum_contract, dict):
+ raise AstroPRISMStage1ADatasetError(
+ "raw-spectrum manifest is missing the spectrum contract"
+ )
+ expected_contract = {
+ "length": ASTROPRISM_SPECTRUM_LENGTH,
+ "aion_token_count": ASTROPRISM_SPECTRUM_TOKENS,
+ "observed_frame": True,
+ "mask_true_means_excluded": True,
+ "aion_modality": "DESISpectrum",
+ }
+ for key, expected in expected_contract.items():
+ actual = spectrum_contract.get(key)
+ if actual != expected or type(actual) is not type(expected):
+ raise AstroPRISMStage1ADatasetError(
+ "raw-spectrum manifest contract mismatch: "
+ f"spectrum.{key} must be {expected!r}, got {actual!r}"
+ )
+
+ caption_type = self.captions_manifest.get("artifact_type")
+ if caption_type != "astroprism_spectrum_captions_accepted":
+ raise AstroPRISMStage1ADatasetError(
+ "spectrum-caption manifest artifact_type must be "
+ "'astroprism_spectrum_captions_accepted'"
+ )
+ if self.captions_manifest.get("role") != "reviewed_alignment_dataset":
+ raise AstroPRISMStage1ADatasetError(
+ "spectrum-caption manifest role must be 'reviewed_alignment_dataset'"
+ )
+ review = self.captions_manifest.get("review")
+ if not isinstance(review, dict) or review.get("status") != "accepted":
+ raise AstroPRISMStage1ADatasetError(
+ "spectrum-caption manifest review.status must be 'accepted'"
+ )
+ if review.get("decision_source") != "user_instruction":
+ raise AstroPRISMStage1ADatasetError(
+ "spectrum-caption manifest review.decision_source must be " "'user_instruction'"
+ )
+ policy = self.captions_manifest.get("training_policy")
+ if not isinstance(policy, dict) or policy.get("dataset_ready") is not True:
+ raise AstroPRISMStage1ADatasetError(
+ "spectrum-caption manifest training_policy.dataset_ready must be true"
+ )
+ if policy.get("gradient_training_splits") != ["train"]:
+ raise AstroPRISMStage1ADatasetError(
+ "spectrum-caption manifest must restrict gradient training to ['train']"
+ )
+ evaluation_splits = policy.get("evaluation_splits")
+ if not isinstance(evaluation_splits, list) or set(evaluation_splits) != {
+ "validation",
+ "test",
+ }:
+ raise AstroPRISMStage1ADatasetError(
+ "spectrum-caption manifest evaluation splits must be validation and test"
+ )
+ if policy.get("roles") != _ROLE_BY_SPLIT:
+ raise AstroPRISMStage1ADatasetError(
+ "spectrum-caption manifest split roles do not match the Stage 1B contract"
+ )
+
+ def _validate_caption_file(self, output: dict[str, Any]) -> None:
+ if not self._all_records:
+ raise AstroPRISMStage1ADatasetError("spectrum-caption JSONL is empty")
+ declared_count = self.captions_manifest.get("row_count")
+ if declared_count != len(self._all_records):
+ raise AstroPRISMStage1ADatasetError(
+ "spectrum-caption row count mismatch: manifest declares "
+ f"{declared_count!r}, JSONL contains {len(self._all_records)}"
+ )
+ declared_indices = self.captions_manifest.get("candidate_indices")
+ observed_indices = [record.get("candidate_index") for record in self._all_records]
+ if declared_indices != observed_indices:
+ raise AstroPRISMStage1ADatasetError(
+ "spectrum-caption manifest candidate_indices do not match captions.jsonl"
+ )
+ observed_splits = dict(Counter(record.get("split") for record in self._all_records))
+ declared_splits = self.captions_manifest.get("split_counts")
+ nested_splits = self.captions_manifest.get("splits", {}).get("row_counts")
+ if declared_splits != observed_splits or nested_splits != observed_splits:
+ raise AstroPRISMStage1ADatasetError(
+ "spectrum-caption split counts do not match captions.jsonl"
+ )
+ declared_size = output.get("size_bytes")
+ if declared_size is not None and declared_size != self.captions_path.stat().st_size:
+ raise AstroPRISMStage1ADatasetError(
+ "spectrum-caption JSONL size disagrees with its manifest"
+ )
+ if self.config.verify_file_hashes and output.get("sha256"):
+ actual = _sha256_file(self.captions_path)
+ if actual != output["sha256"]:
+ raise AstroPRISMStage1ADatasetError(
+ "spectrum-caption JSONL SHA-256 mismatch: expected "
+ f"{output['sha256']}, got {actual}"
+ )
+
+ def _open_and_validate_array(
+ self,
+ key: str,
+ output: dict[str, Any],
+ expected_dtype: np.dtype[Any],
+ ) -> np.ndarray:
+ path = _resolve_output_path(self.raw_manifest_path, output, name=key)
+ if not path.is_file():
+ raise AstroPRISMStage1ADatasetError(f"raw spectrum array does not exist: {path}")
+ try:
+ array = np.load(path, mmap_mode="r", allow_pickle=False)
+ except (OSError, ValueError) as exc:
+ raise AstroPRISMStage1ADatasetError(
+ f"cannot mmap raw spectrum array {path}: {exc}"
+ ) from exc
+
+ expected_shape = (len(self._all_records), ASTROPRISM_SPECTRUM_LENGTH)
+ if array.shape != expected_shape:
+ raise AstroPRISMStage1ADatasetError(
+ f"{key} must have shape {expected_shape}, got {array.shape}"
+ )
+ if array.dtype != expected_dtype:
+ raise AstroPRISMStage1ADatasetError(
+ f"{key} must have dtype {expected_dtype}, got {array.dtype}"
+ )
+ if output.get("shape") is not None and list(array.shape) != output["shape"]:
+ raise AstroPRISMStage1ADatasetError(
+ f"{key} shape disagrees with manifest: {array.shape} != {output['shape']}"
+ )
+ if output.get("dtype") is not None and np.dtype(output["dtype"]) != array.dtype:
+ raise AstroPRISMStage1ADatasetError(
+ f"{key} dtype disagrees with manifest: {array.dtype} != {output['dtype']}"
+ )
+ if self.config.verify_file_hashes and output.get("sha256"):
+ actual = _sha256_file(path)
+ if actual != output["sha256"]:
+ raise AstroPRISMStage1ADatasetError(
+ f"{key} SHA-256 mismatch: expected {output['sha256']}, got {actual}"
+ )
+ return array
+
+ def _validate_roster_identity(self) -> None:
+ raw_roster = self.raw_manifest.get("roster")
+ caption_roster = self.captions_manifest.get("roster")
+ if not isinstance(raw_roster, dict) or not isinstance(caption_roster, dict):
+ raise AstroPRISMStage1ADatasetError(
+ "raw and spectrum-caption manifests must both declare roster provenance"
+ )
+ raw_sha = raw_roster.get("sha256")
+ caption_sha = caption_roster.get("sha256")
+ if not isinstance(raw_sha, str) or not raw_sha:
+ raise AstroPRISMStage1ADatasetError("raw-spectrum roster SHA-256 is missing")
+ if not isinstance(caption_sha, str) or not caption_sha:
+ raise AstroPRISMStage1ADatasetError("spectrum-caption roster SHA-256 is missing")
+ if raw_sha != caption_sha:
+ raise AstroPRISMStage1ADatasetError(
+ f"raw/spectrum-caption roster SHA-256 mismatch: {raw_sha} != {caption_sha}"
+ )
+ expected_count = len(self._all_records)
+ for label, roster in (("raw", raw_roster), ("spectrum-caption", caption_roster)):
+ if roster.get("row_count") != expected_count:
+ raise AstroPRISMStage1ADatasetError(
+ f"{label} roster row_count must be {expected_count}, "
+ f"got {roster.get('row_count')!r}"
+ )
+
+ def _validate_caption_alignment(self) -> None:
+ raw_candidates = self.raw_manifest.get("candidates")
+ if not isinstance(raw_candidates, list):
+ raise AstroPRISMStage1ADatasetError("raw-spectrum manifest is missing candidates list")
+ if len(raw_candidates) != len(self._all_records):
+ raise AstroPRISMStage1ADatasetError(
+ f"raw/spectrum-caption row count mismatch: {len(raw_candidates)} != "
+ f"{len(self._all_records)}"
+ )
+ paired_ids: set[str] = set()
+ caption_ids: set[str] = set()
+ for index, (raw_record, caption_record) in enumerate(
+ zip(raw_candidates, self._all_records, strict=True)
+ ):
+ for kind, record in (("raw", raw_record), ("caption", caption_record)):
+ if not isinstance(record, dict) or record.get("candidate_index") != index:
+ raise AstroPRISMStage1ADatasetError(
+ f"{kind} candidate row {index} does not preserve roster order"
+ )
+ for field in ("paired_targetid", "caption_object_id"):
+ raw_value = raw_record.get(field)
+ caption_value = caption_record.get(field)
+ if not isinstance(caption_value, str) or not caption_value.strip():
+ raise AstroPRISMStage1ADatasetError(
+ f"spectrum-caption candidate {index} has no {field}"
+ )
+ if str(raw_value) != caption_value:
+ raise AstroPRISMStage1ADatasetError(
+ f"candidate {index} {field} mismatch between raw and spectrum-caption "
+ f"artifacts: {raw_value!r} != {caption_value!r}"
+ )
+ paired_id = caption_record["paired_targetid"]
+ caption_id = caption_record["caption_object_id"]
+ if paired_id in paired_ids or caption_id in caption_ids:
+ raise AstroPRISMStage1ADatasetError(
+ f"spectrum-caption candidate {index} repeats an object identity"
+ )
+ paired_ids.add(paired_id)
+ caption_ids.add(caption_id)
+ source_group = caption_record.get("source_group_id")
+ if not isinstance(source_group, str) or not source_group.strip():
+ raise AstroPRISMStage1ADatasetError(
+ f"spectrum-caption candidate {index} has no source_group_id"
+ )
+ split = caption_record.get("split")
+ if split not in _ROLE_BY_SPLIT:
+ raise AstroPRISMStage1ADatasetError(
+ f"spectrum-caption candidate {index} has unsupported split {split!r}"
+ )
+
+ def _validate_training_safety(self) -> None:
+ accepted_statuses = {
+ status.strip().casefold()
+ for status in self.config.accepted_caption_statuses
+ if status.strip()
+ }
+ if not accepted_statuses:
+ raise AstroPRISMStage1ADatasetError("accepted_caption_statuses must be non-empty")
+ for index in self._indices:
+ record = self._all_records[index]
+ caption = record.get("spectrum_caption")
+ status = str(caption.get("status", "")).strip() if isinstance(caption, dict) else ""
+ if status.casefold() not in accepted_statuses:
+ raise AstroPRISMStage1ADatasetError(
+ f"candidate {index} has spectrum_caption.status={status!r}; expected "
+ f"one of {sorted(self.config.accepted_caption_statuses)!r}"
+ )
+ review = record.get("review")
+ if not isinstance(review, dict) or review.get("status") != "accepted":
+ raise AstroPRISMStage1ADatasetError(
+ f"candidate {index} review.status must be 'accepted'"
+ )
+ if review.get("decision_source") != "user_instruction":
+ raise AstroPRISMStage1ADatasetError(
+ f"candidate {index} review.decision_source must be 'user_instruction'"
+ )
+ policy = record.get("training_policy")
+ if not isinstance(policy, dict) or policy.get("dataset_ready") is not True:
+ raise AstroPRISMStage1ADatasetError(
+ f"candidate {index} training_policy.dataset_ready must be true"
+ )
+ split = record["split"]
+ is_train = split == "train"
+ if policy.get("gradient_training_allowed") is not is_train:
+ raise AstroPRISMStage1ADatasetError(
+ f"candidate {index} gradient_training_allowed must be {is_train} "
+ f"for split {split!r}"
+ )
+ if policy.get("eligible_for_training") is not is_train:
+ raise AstroPRISMStage1ADatasetError(
+ f"candidate {index} eligible_for_training must be {is_train} "
+ f"for split {split!r}"
+ )
+ if policy.get("eligible_for_evaluation") is not (not is_train):
+ raise AstroPRISMStage1ADatasetError(
+ f"candidate {index} eligible_for_evaluation must be {not is_train} "
+ f"for split {split!r}"
+ )
+ expected_role = _ROLE_BY_SPLIT[split]
+ if policy.get("role") != expected_role:
+ raise AstroPRISMStage1ADatasetError(
+ f"candidate {index} training_policy.role must be {expected_role!r}"
+ )
+ self._caption_text(record)
+
+ def _caption_text(self, record: dict[str, Any]) -> tuple[str, str]:
+ caption = record.get("spectrum_caption")
+ if not isinstance(caption, dict):
+ raise AstroPRISMStage1ADatasetError(
+ f"candidate {record.get('candidate_index')} has no spectrum_caption object"
+ )
+ value = caption.get("original")
+ if not isinstance(value, str) or not value.strip():
+ raise AstroPRISMStage1ADatasetError(
+ f"candidate {record.get('candidate_index')} has no non-empty "
+ "spectrum_caption.original"
+ )
+ return value.strip(), "original"
+
+ def _tokenize_chat(self, target: str) -> tuple[torch.Tensor, torch.Tensor]:
+ prefix_messages = [
+ {"role": "system", "content": self.config.system_prompt},
+ {"role": "user", "content": self.config.user_prompt},
+ ]
+ full_messages = [*prefix_messages, {"role": "assistant", "content": target}]
+ common = {"tokenize": True, "enable_thinking": False}
+ try:
+ prompt_value = self.tokenizer.apply_chat_template(
+ prefix_messages,
+ add_generation_prompt=True,
+ **common,
+ )
+ full_value = self.tokenizer.apply_chat_template(
+ full_messages,
+ add_generation_prompt=False,
+ **common,
+ )
+ except (TypeError, ValueError) as exc:
+ raise AstroPRISMStage1ADatasetError(
+ "tokenizer failed Qwen non-thinking chat formatting"
+ ) from exc
+ prompt_ids = _coerce_token_ids(prompt_value, context="chat prompt")
+ full_ids = _coerce_token_ids(full_value, context="full chat")
+ if len(full_ids) <= len(prompt_ids) or full_ids[: len(prompt_ids)] != prompt_ids:
+ raise AstroPRISMStage1ADatasetError(
+ "Qwen full chat is not prefixed by the generation prompt"
+ )
+ merged_length = ASTROPRISM_SPECTRUM_TOKENS + len(full_ids)
+ if merged_length > self.config.max_merged_seq_length:
+ raise AstroPRISMStage1ADatasetError(
+ f"merged sequence length {merged_length} exceeds "
+ f"max_merged_seq_length={self.config.max_merged_seq_length} "
+ f"({ASTROPRISM_SPECTRUM_TOKENS} spectrum tokens + "
+ f"{len(full_ids)} text tokens); refusing silent truncation"
+ )
+ input_ids = torch.tensor(full_ids, dtype=torch.long)
+ labels = input_ids.clone()
+ labels[: len(prompt_ids)] = IGNORE_INDEX
+ if not torch.any(labels != IGNORE_INDEX):
+ raise AstroPRISMStage1ADatasetError("assistant response produced no supervised tokens")
+ return input_ids, labels
+
+ def __getitem__(self, item: int) -> dict[str, Any]:
+ source_index = self._indices[item]
+ record = self._all_records[source_index]
+ raw_record = self.raw_manifest["candidates"][source_index]
+
+ spectrum = {
+ "flux": torch.from_numpy(
+ np.array(
+ self._spectra["spectrum_flux"][source_index],
+ dtype=np.float32,
+ order="C",
+ copy=True,
+ )
+ ),
+ "ivar": torch.from_numpy(
+ np.array(
+ self._spectra["spectrum_ivar"][source_index],
+ dtype=np.float32,
+ order="C",
+ copy=True,
+ )
+ ),
+ "mask": torch.from_numpy(
+ np.array(
+ self._spectra["spectrum_mask"][source_index],
+ dtype=np.bool_,
+ order="C",
+ copy=True,
+ )
+ ),
+ "wavelength": torch.from_numpy(
+ np.array(
+ self._spectra["spectrum_wavelength"][source_index],
+ dtype=np.float32,
+ order="C",
+ copy=True,
+ )
+ ),
+ }
+ if not all(
+ bool(torch.isfinite(spectrum[key]).all()) for key in ("flux", "ivar", "wavelength")
+ ):
+ raise AstroPRISMStage1ADatasetError(
+ f"candidate {source_index} raw spectrum contains non-finite values"
+ )
+ if bool(torch.any(spectrum["ivar"] < 0)):
+ raise AstroPRISMStage1ADatasetError(
+ f"candidate {source_index} raw spectrum contains negative inverse variance"
+ )
+ if not bool(torch.all(torch.diff(spectrum["wavelength"]) > 0)):
+ raise AstroPRISMStage1ADatasetError(
+ f"candidate {source_index} wavelength is not strictly increasing"
+ )
+ if bool(torch.all(spectrum["mask"])):
+ raise AstroPRISMStage1ADatasetError(
+ f"candidate {source_index} spectrum has no usable samples"
+ )
+
+ caption_text, selected_caption_field = self._caption_text(record)
+ input_ids, labels = self._tokenize_chat(caption_text)
+ provenance = {
+ "candidate_index": record["candidate_index"],
+ "paired_targetid": record["paired_targetid"],
+ "caption_object_id": record["caption_object_id"],
+ "source_group_id": record["source_group_id"],
+ "split": record["split"],
+ "raw_row": raw_record.get("raw_row"),
+ "coordinates_deg": raw_record.get("coordinates_deg"),
+ "raw_source": self.raw_manifest.get("source"),
+ "roster_sha256": self.captions_manifest["roster"]["sha256"],
+ "caption_field": selected_caption_field,
+ "caption_status": record["spectrum_caption"]["status"],
+ "caption_review_status": record["review"]["status"],
+ "caption_review_decision_source": record["review"]["decision_source"],
+ "caption_training_role": record["training_policy"]["role"],
+ "caption_supervision": "reviewed_spectrum_caption",
+ "spectral_grounding_claimed": True,
+ "chat_template": "qwen_non_thinking",
+ }
+ return {
+ "spectrum": spectrum,
+ "text": input_ids,
+ "text_attention_mask": torch.ones_like(input_ids, dtype=torch.long),
+ "labels": labels,
+ "spectrum_token_count": ASTROPRISM_SPECTRUM_TOKENS,
+ "merged_length": ASTROPRISM_SPECTRUM_TOKENS + input_ids.numel(),
+ "provenance": provenance,
+ }
+
+
+__all__ = [
+ "ASTROPRISM_SPECTRUM_LENGTH",
+ "ASTROPRISM_SPECTRUM_TOKENS",
+ "DEFAULT_SPECTRUM_SYSTEM_PROMPT",
+ "DEFAULT_SPECTRUM_USER_PROMPT",
+ "AstroPRISMStage1BSpectrumDataset",
+ "AstroPRISMStage1BSpectrumDatasetConfig",
+]
diff --git a/src/data/multimodal.py b/src/data/multimodal.py
index 64edce3..7fba41b 100644
--- a/src/data/multimodal.py
+++ b/src/data/multimodal.py
@@ -8,7 +8,6 @@
import warnings
import datasets
-import h5py
import numpy as np
import requests
import torch
@@ -2134,6 +2133,14 @@ def _create_physics_stream(self, hf_id, local_path=None):
# 2. Define Generator (If Valid)
def data_generator():
+ try:
+ import h5py
+ except ImportError as exc:
+ raise RuntimeError(
+ "The Well HDF5 physics stream requires the optional h5py "
+ "runtime and HDF5 shared libraries."
+ ) from exc
+
while True:
try:
with h5py.File(file_path, "r") as f:
diff --git a/src/encoders/__init__.py b/src/encoders/__init__.py
index 0e9eb18..af6e0c0 100644
--- a/src/encoders/__init__.py
+++ b/src/encoders/__init__.py
@@ -9,6 +9,9 @@
# GraphEncoder and GeometryEncoder both swallow their own optional-dep
# ImportErrors at module level, so importing them here always succeeds. The
# loud failure happens inside their __init__ via require_modality_deps().
+from .aion_image import AIONImageEncoder
+from .aion_multimodal import AIONMultimodalEncoder
+from .aion_spectrum import AIONSpectrumEncoder
from .crystal_graph import CrystalGraphTokenEncoder
from .dna import DNAEncoder
from .geometry import GeometryEncoder
@@ -23,6 +26,9 @@
"TableEncoder",
"TimeSeriesEncoder",
"ImageEncoder",
+ "AIONImageEncoder",
+ "AIONMultimodalEncoder",
+ "AIONSpectrumEncoder",
"GeometryEncoder",
"GraphEncoder",
"CrystalGraphTokenEncoder",
diff --git a/src/encoders/aion_image.py b/src/encoders/aion_image.py
new file mode 100644
index 0000000..7689fd5
--- /dev/null
+++ b/src/encoders/aion_image.py
@@ -0,0 +1,308 @@
+"""AION raw-image encoder for AstroPRISM.
+
+This module deliberately imports AION only when :class:`AIONImageEncoder` is
+constructed. Importing PRISM therefore remains possible in environments that
+do not have the optional ``polymathic-aion`` package installed.
+
+The image path matches the released AION feature pipeline:
+
+``calibrated g/r/i/z flux -> ImageCodec -> 576 image token ids -> AION.encode``
+
+The image codec uses hard finite-scalar quantization. It is consequently kept
+in float32, frozen, and in evaluation mode. The AION image-token embedding and
+shared transformer encoder can be unfrozen later through
+``set_shared_encoder_trainable``; the codec is never made trainable by that
+hook.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+from .base import ModalityEncoder
+
+AION_REPO_ID = "polymathic-ai/aion-base"
+AION_MODEL_REVISION = "40541618104bab0fa85c8af68daeb867a720bb8c"
+AION_IMAGE_CODEC_REVISION = AION_MODEL_REVISION
+
+AION_IMAGE_TOKENS = 576
+AION_IMAGE_WIDTH = 768
+AION_RAW_IMAGE_SIZE = (160, 160)
+AION_LEGACY_SURVEY_BANDS = ("DES-G", "DES-R", "DES-I", "DES-Z")
+
+
+def _load_aion_types() -> tuple[type[Any], type[Any], type[Any]]:
+ """Import the optional AION classes only when the encoder is requested."""
+
+ try:
+ from aion import AION
+ from aion.codecs import ImageCodec
+ from aion.modalities import LegacySurveyImage
+ except ImportError as exc: # pragma: no cover - exact message is environment-dependent
+ raise ImportError(
+ "AIONImageEncoder requires the optional 'polymathic-aion' package "
+ "(and its dependencies). Install a pinned AION source checkout before "
+ "selecting the AION image encoder."
+ ) from exc
+ return AION, ImageCodec, LegacySurveyImage
+
+
+class AIONImageEncoder(ModalityEncoder):
+ """Encode calibrated Legacy Survey images with the pinned AION image path.
+
+ Args:
+ model_name_or_path: A Hugging Face repository id or a complete local
+ AION model directory.
+ model_revision: Explicit revision for the AION transformer weights.
+ codec_revision: Explicit revision for the image codec weights.
+ cache_dir: Optional Hugging Face cache containing the pinned snapshots.
+ local_files_only: Refuse network access when loading model assets.
+ bands: Input channel names. Only the verified Legacy Survey g/r/i/z
+ ordering is accepted.
+ raw_image_size: Expected calibrated cutout size before AION's official
+ 96x96 center crop.
+
+ Input:
+ Floating tensor ``(B, 4, 160, 160)`` in calibrated Legacy Survey flux
+ units and g/r/i/z channel order.
+
+ Output:
+ AION encoder features ``(B, 576, 768)``.
+ """
+
+ def __init__(
+ self,
+ model_name_or_path: str = AION_REPO_ID,
+ *,
+ model_revision: str = AION_MODEL_REVISION,
+ codec_revision: str = AION_IMAGE_CODEC_REVISION,
+ cache_dir: str | Path | None = None,
+ local_files_only: bool = True,
+ bands: Sequence[str] = AION_LEGACY_SURVEY_BANDS,
+ raw_image_size: tuple[int, int] = AION_RAW_IMAGE_SIZE,
+ ) -> None:
+ super().__init__(AION_IMAGE_WIDTH)
+
+ resolved_bands = tuple(str(band) for band in bands)
+ if resolved_bands != AION_LEGACY_SURVEY_BANDS:
+ raise ValueError(
+ "AION raw-image input must use calibrated Legacy Survey g/r/i/z channels "
+ f"in {AION_LEGACY_SURVEY_BANDS!r} order; got {resolved_bands!r}."
+ )
+ if len(raw_image_size) != 2 or any(int(size) <= 0 for size in raw_image_size):
+ raise ValueError(
+ f"raw_image_size must contain two positive integers, got {raw_image_size!r}"
+ )
+
+ self.model_name_or_path = str(model_name_or_path)
+ self.model_revision = str(model_revision)
+ self.codec_revision = str(codec_revision)
+ self.cache_dir = None if cache_dir is None else str(cache_dir)
+ self.local_files_only = bool(local_files_only)
+ self.bands = resolved_bands
+ self.raw_image_size = tuple(int(size) for size in raw_image_size)
+ self.num_image_tokens = AION_IMAGE_TOKENS
+ self.hidden_dim = AION_IMAGE_WIDTH
+
+ AION, ImageCodec, LegacySurveyImage = _load_aion_types()
+ self._legacy_survey_image_type = LegacySurveyImage
+ self.image_token_key = str(LegacySurveyImage.token_key)
+
+ common_load_kwargs = {
+ "cache_dir": self.cache_dir,
+ "local_files_only": self.local_files_only,
+ }
+ self.aion_model = AION.from_pretrained(
+ self.model_name_or_path,
+ revision=self.model_revision,
+ **common_load_kwargs,
+ )
+ self.image_codec = ImageCodec.from_pretrained(
+ self.model_name_or_path,
+ modality=LegacySurveyImage,
+ revision=self.codec_revision,
+ **common_load_kwargs,
+ )
+
+ self._validate_aion_interface()
+ self._shared_encoder_trainable = False
+ self.set_shared_encoder_trainable(False)
+
+ def _validate_aion_interface(self) -> None:
+ required = ("encoder_embeddings", "encoder", "encoder_norm", "decoder_proj_context")
+ missing = [name for name in required if not hasattr(self.aion_model, name)]
+ if missing:
+ raise RuntimeError(f"Loaded AION model is missing encoder components: {missing}")
+ if self.image_token_key not in self.aion_model.encoder_embeddings:
+ raise RuntimeError(
+ f"Loaded AION model has no encoder embedding for {self.image_token_key!r}"
+ )
+
+ def _shared_encoder_modules(self) -> tuple[nn.Module, ...]:
+ """Modules used by ``AION.encode`` for an image-only encoder pass."""
+
+ return (
+ self.aion_model.encoder_embeddings[self.image_token_key],
+ self.aion_model.encoder,
+ self.aion_model.encoder_norm,
+ self.aion_model.decoder_proj_context,
+ )
+
+ def _enforce_trainability(self) -> None:
+ # Freeze every AION component first so decoder heads and embeddings that
+ # are not used by this image-only wrapper can never enter an optimizer.
+ for parameter in self.aion_model.parameters():
+ parameter.requires_grad_(False)
+ if self._shared_encoder_trainable:
+ for module in self._shared_encoder_modules():
+ for parameter in module.parameters():
+ parameter.requires_grad_(True)
+
+ # Hard quantization breaks a pixel-to-code gradient path. The codec is
+ # always a deterministic, float32 preprocessing/tokenization component.
+ self.image_codec.float()
+ self.image_codec.eval()
+ for parameter in self.image_codec.parameters():
+ parameter.requires_grad_(False)
+
+ def _apply(self, fn, recurse=True):
+ """Keep the codec in float32 through blanket dtype casts.
+
+ ``nn.Module.to()`` reaches children via ``_apply``, so a child-level
+ ``to()`` override never sees it and the trainer's
+ ``self.model.encoders.to(dtype=target_dtype)`` rounds the codec.
+ Calling ``.float()`` afterwards restores the dtype but not the values,
+ and the codec quantizes hard: rounded weights can change token ids.
+ Device moves pass through untouched. Mirrors ``DNAEncoder._apply``.
+ """
+ is_reduced_precision = False
+ try:
+ probe = fn(torch.zeros(1, dtype=torch.float32))
+ if isinstance(probe, torch.Tensor) and probe.dtype not in (
+ torch.float32,
+ torch.float64,
+ ):
+ is_reduced_precision = True
+ except Exception:
+ pass
+
+ if is_reduced_precision:
+ for child in self.children():
+ if child is not self.image_codec:
+ child._apply(fn, recurse=recurse)
+ for key, buf in self._buffers.items():
+ if buf is not None:
+ self._buffers[key] = fn(buf)
+ return self
+
+ return super()._apply(fn, recurse=recurse)
+
+ @property
+ def shared_encoder_trainable(self) -> bool:
+ """Whether the differentiable AION image/shared-encoder path is trainable."""
+
+ return self._shared_encoder_trainable
+
+ def set_shared_encoder_trainable(self, trainable: bool) -> AIONImageEncoder:
+ """Freeze or unfreeze only the differentiable AION encoding path.
+
+ The image codec remains float32, frozen, and in evaluation mode for both
+ settings.
+ """
+
+ self._shared_encoder_trainable = bool(trainable)
+ self._enforce_trainability()
+ self.train(self.training)
+ return self
+
+ def train(self, mode: bool = True) -> AIONImageEncoder:
+ """Set runtime mode while preserving AION codec/freeze invariants."""
+
+ super().train(mode)
+ self._enforce_trainability()
+
+ # Start from deterministic evaluation behavior for the full AION model.
+ # If encoder fine-tuning is explicitly enabled, only modules on the
+ # image/shared-encoder path enter training mode.
+ self.aion_model.eval()
+ if mode and self._shared_encoder_trainable:
+ for module in self._shared_encoder_modules():
+ module.train(True)
+ self.image_codec.eval()
+ return self
+
+ def tokens_per_instance(self) -> int:
+ """Return the fixed number of image tokens the AION codec emits."""
+ return self.num_image_tokens
+
+ def _validate_inputs(self, inputs: torch.Tensor) -> None:
+ if not isinstance(inputs, torch.Tensor):
+ raise TypeError(f"AION image input must be a torch.Tensor, got {type(inputs).__name__}")
+ expected = (len(self.bands), *self.raw_image_size)
+ if inputs.ndim != 4 or tuple(inputs.shape[1:]) != expected:
+ raise ValueError(
+ "AION image input must have shape "
+ f"(B, {expected[0]}, {expected[1]}, {expected[2]}), got {tuple(inputs.shape)}"
+ )
+ if not torch.is_floating_point(inputs):
+ raise TypeError(f"AION image input must be floating point, got {inputs.dtype}")
+ if not bool(torch.isfinite(inputs).all()):
+ raise ValueError("AION image input contains non-finite calibrated flux values")
+
+ def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+ """Encode calibrated four-band images into ``(B, 576, 768)`` features."""
+ self._validate_inputs(inputs)
+ self._enforce_trainability()
+
+ # AION's Clamp operates in place on a view returned by CenterCrop. Clone
+ # here so encoding never mutates the caller's training batch. The hard
+ # codec is deliberately outside autograd and outside mixed precision.
+ with torch.no_grad(), torch.autocast(device_type=inputs.device.type, enabled=False):
+ codec_flux = inputs.detach().to(dtype=torch.float32).clone()
+ image = self._legacy_survey_image_type(flux=codec_flux, bands=list(self.bands))
+ image_token_ids = self.image_codec.encode(image)
+
+ expected_tokens = (inputs.shape[0], self.num_image_tokens)
+ if tuple(image_token_ids.shape) != expected_tokens:
+ raise RuntimeError(
+ f"AION ImageCodec must emit token ids with shape {expected_tokens}, "
+ f"got {tuple(image_token_ids.shape)}"
+ )
+
+ encoded_inputs = {self.image_token_key: image_token_ids}
+ if self._shared_encoder_trainable:
+ features = self.aion_model.encode(
+ encoded_inputs,
+ num_encoder_tokens=self.num_image_tokens,
+ )
+ else:
+ with torch.no_grad():
+ features = self.aion_model.encode(
+ encoded_inputs,
+ num_encoder_tokens=self.num_image_tokens,
+ )
+
+ expected_features = (inputs.shape[0], self.num_image_tokens, self.output_dim)
+ if tuple(features.shape) != expected_features:
+ raise RuntimeError(
+ f"AION shared encoder must emit features with shape {expected_features}, "
+ f"got {tuple(features.shape)}"
+ )
+ return features
+
+
+__all__ = [
+ "AIONImageEncoder",
+ "AION_IMAGE_CODEC_REVISION",
+ "AION_IMAGE_TOKENS",
+ "AION_IMAGE_WIDTH",
+ "AION_LEGACY_SURVEY_BANDS",
+ "AION_MODEL_REVISION",
+ "AION_RAW_IMAGE_SIZE",
+ "AION_REPO_ID",
+]
diff --git a/src/encoders/aion_multimodal.py b/src/encoders/aion_multimodal.py
new file mode 100644
index 0000000..99d8651
--- /dev/null
+++ b/src/encoders/aion_multimodal.py
@@ -0,0 +1,415 @@
+"""Joint AION image and spectrum encoder with one shared transformer owner.
+
+The standalone :class:`AIONImageEncoder` and :class:`AIONSpectrumEncoder`
+each own a complete AION model. That is the right ownership model for a
+single-modality run, but registering both wrappers in one PRISM model would
+duplicate AION's shared transformer (and, when unfrozen, let the copies
+diverge). This module provides the joint path instead:
+
+``raw image -> ImageCodec -------+``
+`` AION.encode -> 768-d tokens``
+``raw spectrum -> SpectrumCodec -+``
+
+The AION transformer is registered exactly once. Both hard-quantizing
+codecs remain float32, frozen, and in evaluation mode for the lifetime of the
+module. The shared transformer and both active modality token embeddings can
+be unfrozen together through :meth:`set_shared_encoder_trainable`.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+from .aion_image import (
+ AION_IMAGE_CODEC_REVISION,
+ AION_IMAGE_TOKENS,
+ AION_IMAGE_WIDTH,
+ AION_LEGACY_SURVEY_BANDS,
+ AION_MODEL_REVISION,
+ AION_RAW_IMAGE_SIZE,
+ AION_REPO_ID,
+)
+from .aion_spectrum import (
+ AION_RAW_SPECTRUM_LENGTH,
+ AION_SPECTRUM_CODEC_REVISION,
+ AION_SPECTRUM_FIELDS,
+ AION_SPECTRUM_TOKENS,
+ AION_SPECTRUM_WIDTH,
+)
+from .base import ModalityEncoder
+
+
+def _load_aion_types() -> tuple[type[Any], type[Any], type[Any], type[Any], type[Any]]:
+ """Import optional AION classes only when the joint encoder is requested."""
+
+ try:
+ from aion import AION
+ from aion.codecs import ImageCodec, SpectrumCodec
+ from aion.modalities import DESISpectrum, LegacySurveyImage
+ except ImportError as exc: # pragma: no cover - environment-dependent
+ raise ImportError(
+ "AIONMultimodalEncoder requires the optional 'polymathic-aion' package "
+ "(and its dependencies). Install a pinned AION source checkout before "
+ "selecting joint AION image+spectrum encoding."
+ ) from exc
+ return AION, ImageCodec, SpectrumCodec, LegacySurveyImage, DESISpectrum
+
+
+class AIONMultimodalEncoder(ModalityEncoder):
+ """Encode raw Legacy Survey images and DESI spectra through one AION model.
+
+ ``encode_modality("image", tensor)`` accepts ``float[B, 4, 160, 160]``
+ calibrated Legacy Survey flux and emits ``float[B, 576, 768]`` features.
+ ``encode_modality("spectrum", mapping)`` accepts the four raw DESI arrays
+ and emits ``float[B, 273, 768]`` features.
+ """
+
+ supported_modalities = frozenset({"image", "spectrum"})
+
+ def __init__(
+ self,
+ model_name_or_path: str = AION_REPO_ID,
+ *,
+ model_revision: str = AION_MODEL_REVISION,
+ image_codec_revision: str = AION_IMAGE_CODEC_REVISION,
+ spectrum_codec_revision: str = AION_SPECTRUM_CODEC_REVISION,
+ cache_dir: str | Path | None = None,
+ local_files_only: bool = True,
+ bands: Sequence[str] = AION_LEGACY_SURVEY_BANDS,
+ raw_image_size: tuple[int, int] = AION_RAW_IMAGE_SIZE,
+ raw_spectrum_length: int = AION_RAW_SPECTRUM_LENGTH,
+ ) -> None:
+ if AION_IMAGE_WIDTH != AION_SPECTRUM_WIDTH: # pragma: no cover - pinned constants
+ raise RuntimeError(
+ "Joint AION ownership requires image and spectrum features to share a width"
+ )
+ super().__init__(AION_IMAGE_WIDTH)
+
+ resolved_bands = tuple(str(band) for band in bands)
+ if resolved_bands != AION_LEGACY_SURVEY_BANDS:
+ raise ValueError(
+ "AION raw-image input must use calibrated Legacy Survey g/r/i/z channels "
+ f"in {AION_LEGACY_SURVEY_BANDS!r} order; got {resolved_bands!r}."
+ )
+ if len(raw_image_size) != 2 or any(int(size) <= 0 for size in raw_image_size):
+ raise ValueError(
+ f"raw_image_size must contain two positive integers, got {raw_image_size!r}"
+ )
+ if int(raw_spectrum_length) <= 0:
+ raise ValueError(
+ "raw_spectrum_length must be a positive integer, got "
+ f"{raw_spectrum_length!r}"
+ )
+
+ self.model_name_or_path = str(model_name_or_path)
+ self.model_revision = str(model_revision)
+ self.image_codec_revision = str(image_codec_revision)
+ self.spectrum_codec_revision = str(spectrum_codec_revision)
+ self.cache_dir = None if cache_dir is None else str(cache_dir)
+ self.local_files_only = bool(local_files_only)
+ self.bands = resolved_bands
+ self.raw_image_size = tuple(int(size) for size in raw_image_size)
+ self.raw_spectrum_length = int(raw_spectrum_length)
+ self.num_image_tokens = AION_IMAGE_TOKENS
+ self.num_spectrum_tokens = AION_SPECTRUM_TOKENS
+ self.hidden_dim = AION_IMAGE_WIDTH
+
+ (
+ AION,
+ ImageCodec,
+ SpectrumCodec,
+ LegacySurveyImage,
+ DESISpectrum,
+ ) = _load_aion_types()
+ self._legacy_survey_image_type = LegacySurveyImage
+ self._desi_spectrum_type = DESISpectrum
+ self.image_token_key = str(LegacySurveyImage.token_key)
+ self.spectrum_token_key = str(DESISpectrum.token_key)
+
+ common_load_kwargs = {
+ "cache_dir": self.cache_dir,
+ "local_files_only": self.local_files_only,
+ }
+ # This is the sole registered copy of the AION transformer.
+ self.aion_model = AION.from_pretrained(
+ self.model_name_or_path,
+ revision=self.model_revision,
+ **common_load_kwargs,
+ )
+ self.image_codec = ImageCodec.from_pretrained(
+ self.model_name_or_path,
+ modality=LegacySurveyImage,
+ revision=self.image_codec_revision,
+ **common_load_kwargs,
+ )
+ self.spectrum_codec = SpectrumCodec.from_pretrained(
+ self.model_name_or_path,
+ modality=DESISpectrum,
+ revision=self.spectrum_codec_revision,
+ **common_load_kwargs,
+ )
+
+ self._validate_aion_interface()
+ self._shared_encoder_trainable = False
+ self.set_shared_encoder_trainable(False)
+
+ def _validate_aion_interface(self) -> None:
+ required = ("encoder_embeddings", "encoder", "encoder_norm", "decoder_proj_context")
+ missing = [name for name in required if not hasattr(self.aion_model, name)]
+ if missing:
+ raise RuntimeError(f"Loaded AION model is missing encoder components: {missing}")
+ missing_embeddings = [
+ key
+ for key in (self.image_token_key, self.spectrum_token_key)
+ if key not in self.aion_model.encoder_embeddings
+ ]
+ if missing_embeddings:
+ raise RuntimeError(
+ "Loaded AION model is missing joint encoder embeddings: "
+ f"{missing_embeddings}"
+ )
+
+ def _shared_encoder_modules(self) -> tuple[nn.Module, ...]:
+ """Return every differentiable module used by either joint route."""
+
+ return (
+ self.aion_model.encoder_embeddings[self.image_token_key],
+ self.aion_model.encoder_embeddings[self.spectrum_token_key],
+ self.aion_model.encoder,
+ self.aion_model.encoder_norm,
+ self.aion_model.decoder_proj_context,
+ )
+
+ def _enforce_trainability(self) -> None:
+ # Decoder heads and embeddings for unrelated AION modalities never
+ # enter a PRISM optimizer.
+ for parameter in self.aion_model.parameters():
+ parameter.requires_grad_(False)
+ if self._shared_encoder_trainable:
+ for module in self._shared_encoder_modules():
+ for parameter in module.parameters():
+ parameter.requires_grad_(True)
+
+ for codec in (self.image_codec, self.spectrum_codec):
+ codec.float()
+ codec.eval()
+ for parameter in codec.parameters():
+ parameter.requires_grad_(False)
+
+ @property
+ def shared_encoder_trainable(self) -> bool:
+ """Whether the shared AION transformer is currently trainable."""
+ return self._shared_encoder_trainable
+
+ def set_shared_encoder_trainable(self, trainable: bool) -> AIONMultimodalEncoder:
+ """Toggle the one shared core and both modality embeddings together."""
+
+ self._shared_encoder_trainable = bool(trainable)
+ self._enforce_trainability()
+ self.train(self.training)
+ return self
+
+ def train(self, mode: bool = True) -> AIONMultimodalEncoder:
+ """Set runtime mode while preserving both codec invariants."""
+
+ super().train(mode)
+ self._enforce_trainability()
+ self.aion_model.eval()
+ if mode and self._shared_encoder_trainable:
+ for module in self._shared_encoder_modules():
+ module.train(True)
+ self.image_codec.eval()
+ self.spectrum_codec.eval()
+ return self
+
+ def tokens_per_instance(self, modality: str) -> int:
+ """Return the token count this shared owner emits for ``modality``."""
+ modality_key = str(modality)
+ if modality_key == "image":
+ return self.num_image_tokens
+ if modality_key == "spectrum":
+ return self.num_spectrum_tokens
+ raise ValueError(
+ f"Unsupported AION modality {modality_key!r}; expected 'image' or 'spectrum'"
+ )
+
+ def _validate_image_inputs(self, inputs: torch.Tensor) -> None:
+ if not isinstance(inputs, torch.Tensor):
+ raise TypeError(
+ f"AION image input must be a torch.Tensor, got {type(inputs).__name__}"
+ )
+ expected = (len(self.bands), *self.raw_image_size)
+ if inputs.ndim != 4 or tuple(inputs.shape[1:]) != expected:
+ raise ValueError(
+ "AION image input must have shape "
+ f"(B, {expected[0]}, {expected[1]}, {expected[2]}), got {tuple(inputs.shape)}"
+ )
+ if not torch.is_floating_point(inputs):
+ raise TypeError(f"AION image input must be floating point, got {inputs.dtype}")
+ if not bool(torch.isfinite(inputs).all()):
+ raise ValueError("AION image input contains non-finite calibrated flux values")
+
+ def _validate_spectrum_inputs(self, inputs: Mapping[str, torch.Tensor]) -> None:
+ if not isinstance(inputs, Mapping):
+ raise TypeError(
+ "AION spectrum input must be a mapping of flux/ivar/mask/wavelength "
+ f"tensors, got {type(inputs).__name__}"
+ )
+ missing = set(AION_SPECTRUM_FIELDS).difference(inputs)
+ if missing:
+ raise ValueError(f"AION spectrum input is missing fields: {sorted(missing)}")
+
+ tensors = {field: inputs[field] for field in AION_SPECTRUM_FIELDS}
+ for field, tensor in tensors.items():
+ if not isinstance(tensor, torch.Tensor):
+ raise TypeError(
+ f"AION spectrum field {field!r} must be a torch.Tensor, "
+ f"got {type(tensor).__name__}"
+ )
+ if tensor.ndim != 2 or tuple(tensor.shape[1:]) != (self.raw_spectrum_length,):
+ raise ValueError(
+ f"AION spectrum field {field!r} must have shape "
+ f"(B, {self.raw_spectrum_length}), got {tuple(tensor.shape)}"
+ )
+
+ if len({tuple(tensor.shape) for tensor in tensors.values()}) != 1:
+ raise ValueError("AION spectrum fields must have identical batch and length shapes")
+ if len({tensor.device for tensor in tensors.values()}) != 1:
+ raise ValueError("AION spectrum fields must be on the same device")
+
+ for field in ("flux", "ivar", "wavelength"):
+ tensor = tensors[field]
+ if not torch.is_floating_point(tensor):
+ raise TypeError(
+ f"AION spectrum field {field!r} must be floating point, got {tensor.dtype}"
+ )
+ if not bool(torch.isfinite(tensor).all()):
+ raise ValueError(f"AION spectrum field {field!r} contains non-finite values")
+ if tensors["mask"].dtype != torch.bool:
+ raise TypeError(
+ f"AION spectrum field 'mask' must be boolean, got {tensors['mask'].dtype}"
+ )
+ if bool((tensors["ivar"] < 0).any()):
+ raise ValueError("AION spectrum inverse variance must be non-negative")
+ if bool((torch.diff(tensors["wavelength"], dim=-1) <= 0).any()):
+ raise ValueError("AION spectrum wavelength must be strictly increasing")
+
+ def _encode_tokens(
+ self,
+ *,
+ token_key: str,
+ token_ids: torch.Tensor,
+ token_count: int,
+ ) -> torch.Tensor:
+ encoded_inputs = {token_key: token_ids}
+ if self._shared_encoder_trainable:
+ return self.aion_model.encode(
+ encoded_inputs,
+ num_encoder_tokens=token_count,
+ )
+ with torch.no_grad():
+ return self.aion_model.encode(
+ encoded_inputs,
+ num_encoder_tokens=token_count,
+ )
+
+ def _encode_image(self, inputs: torch.Tensor) -> torch.Tensor:
+ self._validate_image_inputs(inputs)
+ with torch.no_grad(), torch.autocast(device_type=inputs.device.type, enabled=False):
+ codec_flux = inputs.detach().to(dtype=torch.float32).clone()
+ image = self._legacy_survey_image_type(
+ flux=codec_flux,
+ bands=list(self.bands),
+ )
+ token_ids = self.image_codec.encode(image)
+
+ expected_tokens = (inputs.shape[0], self.num_image_tokens)
+ if tuple(token_ids.shape) != expected_tokens:
+ raise RuntimeError(
+ f"AION ImageCodec must emit token ids with shape {expected_tokens}, "
+ f"got {tuple(token_ids.shape)}"
+ )
+ features = self._encode_tokens(
+ token_key=self.image_token_key,
+ token_ids=token_ids,
+ token_count=self.num_image_tokens,
+ )
+ expected_features = (inputs.shape[0], self.num_image_tokens, self.output_dim)
+ if tuple(features.shape) != expected_features:
+ raise RuntimeError(
+ f"AION shared encoder must emit image features with shape {expected_features}, "
+ f"got {tuple(features.shape)}"
+ )
+ return features
+
+ def _encode_spectrum(self, inputs: Mapping[str, torch.Tensor]) -> torch.Tensor:
+ self._validate_spectrum_inputs(inputs)
+ reference = inputs["flux"]
+ with torch.no_grad(), torch.autocast(
+ device_type=reference.device.type,
+ enabled=False,
+ ):
+ spectrum = self._desi_spectrum_type(
+ flux=inputs["flux"].detach().to(dtype=torch.float32).clone(),
+ ivar=inputs["ivar"].detach().to(dtype=torch.float32).clone(),
+ mask=inputs["mask"].detach().clone(),
+ wavelength=inputs["wavelength"].detach().to(dtype=torch.float32).clone(),
+ )
+ token_ids = self.spectrum_codec.encode(spectrum)
+
+ expected_tokens = (reference.shape[0], self.num_spectrum_tokens)
+ if tuple(token_ids.shape) != expected_tokens:
+ raise RuntimeError(
+ f"AION SpectrumCodec must emit token ids with shape {expected_tokens}, "
+ f"got {tuple(token_ids.shape)}"
+ )
+ features = self._encode_tokens(
+ token_key=self.spectrum_token_key,
+ token_ids=token_ids,
+ token_count=self.num_spectrum_tokens,
+ )
+ expected_features = (
+ reference.shape[0],
+ self.num_spectrum_tokens,
+ self.output_dim,
+ )
+ if tuple(features.shape) != expected_features:
+ raise RuntimeError(
+ "AION shared encoder must emit spectrum features with shape "
+ f"{expected_features}, got {tuple(features.shape)}"
+ )
+ return features
+
+ def encode_modality(
+ self,
+ modality: str,
+ inputs: torch.Tensor | Mapping[str, torch.Tensor],
+ ) -> torch.Tensor:
+ """Encode one homogeneous modality batch through the shared owner."""
+
+ self._enforce_trainability()
+ modality_key = str(modality)
+ if modality_key == "image":
+ return self._encode_image(inputs) # type: ignore[arg-type]
+ if modality_key == "spectrum":
+ return self._encode_spectrum(inputs) # type: ignore[arg-type]
+ raise ValueError(
+ f"Unsupported AION modality {modality_key!r}; expected 'image' or 'spectrum'"
+ )
+
+ def forward(
+ self,
+ inputs: torch.Tensor | Mapping[str, torch.Tensor],
+ *,
+ modality: str,
+ ) -> torch.Tensor:
+ """Route one modality through the shared AION encoder."""
+ return self.encode_modality(modality, inputs)
+
+
+__all__ = ["AIONMultimodalEncoder"]
diff --git a/src/encoders/aion_spectrum.py b/src/encoders/aion_spectrum.py
new file mode 100644
index 0000000..7a76e76
--- /dev/null
+++ b/src/encoders/aion_spectrum.py
@@ -0,0 +1,303 @@
+"""AION raw-spectrum encoder for AstroPRISM.
+
+The spectrum path mirrors the released AION feature pipeline:
+
+``DESI flux/ivar/mask/wavelength -> SpectrumCodec -> 273 token ids -> AION.encode``
+
+The codec contains hard quantization, so it is always kept in float32,
+frozen, and in evaluation mode. The DESI token embedding and shared AION
+transformer can be unfrozen later with ``set_shared_encoder_trainable``.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+from .aion_image import AION_MODEL_REVISION, AION_REPO_ID
+from .base import ModalityEncoder
+
+AION_SPECTRUM_CODEC_REVISION = AION_MODEL_REVISION
+AION_SPECTRUM_TOKENS = 273
+AION_SPECTRUM_WIDTH = 768
+AION_RAW_SPECTRUM_LENGTH = 7781
+AION_SPECTRUM_FIELDS = ("flux", "ivar", "mask", "wavelength")
+
+
+def _load_aion_types() -> tuple[type[Any], type[Any], type[Any]]:
+ """Import optional AION spectrum classes only when requested."""
+
+ try:
+ from aion import AION
+ from aion.codecs import SpectrumCodec
+ from aion.modalities import DESISpectrum
+ except ImportError as exc: # pragma: no cover - environment-dependent
+ raise ImportError(
+ "AIONSpectrumEncoder requires the optional 'polymathic-aion' package "
+ "(and its dependencies). Install a pinned AION source checkout before "
+ "selecting the AION spectrum encoder."
+ ) from exc
+ return AION, SpectrumCodec, DESISpectrum
+
+
+class AIONSpectrumEncoder(ModalityEncoder):
+ """Encode observed-frame DESI spectra with the pinned AION spectrum path.
+
+ Input is a mapping with ``flux``, ``ivar``, ``mask``, and ``wavelength``
+ tensors. Every tensor must have shape ``(B, 7781)``; the first, second,
+ and fourth are floating point and ``mask`` is boolean.
+
+ Output has shape ``(B, 273, 768)``.
+ """
+
+ def __init__(
+ self,
+ model_name_or_path: str = AION_REPO_ID,
+ *,
+ model_revision: str = AION_MODEL_REVISION,
+ codec_revision: str = AION_SPECTRUM_CODEC_REVISION,
+ cache_dir: str | Path | None = None,
+ local_files_only: bool = True,
+ raw_spectrum_length: int = AION_RAW_SPECTRUM_LENGTH,
+ ) -> None:
+ super().__init__(AION_SPECTRUM_WIDTH)
+
+ if int(raw_spectrum_length) <= 0:
+ raise ValueError(
+ f"raw_spectrum_length must be a positive integer, got {raw_spectrum_length!r}"
+ )
+
+ self.model_name_or_path = str(model_name_or_path)
+ self.model_revision = str(model_revision)
+ self.codec_revision = str(codec_revision)
+ self.cache_dir = None if cache_dir is None else str(cache_dir)
+ self.local_files_only = bool(local_files_only)
+ self.raw_spectrum_length = int(raw_spectrum_length)
+ self.num_spectrum_tokens = AION_SPECTRUM_TOKENS
+ self.hidden_dim = AION_SPECTRUM_WIDTH
+
+ AION, SpectrumCodec, DESISpectrum = _load_aion_types()
+ self._desi_spectrum_type = DESISpectrum
+ self.spectrum_token_key = str(DESISpectrum.token_key)
+
+ common_load_kwargs = {
+ "cache_dir": self.cache_dir,
+ "local_files_only": self.local_files_only,
+ }
+ self.aion_model = AION.from_pretrained(
+ self.model_name_or_path,
+ revision=self.model_revision,
+ **common_load_kwargs,
+ )
+ self.spectrum_codec = SpectrumCodec.from_pretrained(
+ self.model_name_or_path,
+ modality=DESISpectrum,
+ revision=self.codec_revision,
+ **common_load_kwargs,
+ )
+
+ self._validate_aion_interface()
+ self._shared_encoder_trainable = False
+ self.set_shared_encoder_trainable(False)
+
+ def _validate_aion_interface(self) -> None:
+ required = ("encoder_embeddings", "encoder", "encoder_norm", "decoder_proj_context")
+ missing = [name for name in required if not hasattr(self.aion_model, name)]
+ if missing:
+ raise RuntimeError(f"Loaded AION model is missing encoder components: {missing}")
+ if self.spectrum_token_key not in self.aion_model.encoder_embeddings:
+ raise RuntimeError(
+ f"Loaded AION model has no encoder embedding for {self.spectrum_token_key!r}"
+ )
+
+ def _shared_encoder_modules(self) -> tuple[nn.Module, ...]:
+ """Modules used by ``AION.encode`` for a spectrum-only pass."""
+
+ return (
+ self.aion_model.encoder_embeddings[self.spectrum_token_key],
+ self.aion_model.encoder,
+ self.aion_model.encoder_norm,
+ self.aion_model.decoder_proj_context,
+ )
+
+ def _enforce_trainability(self) -> None:
+ for parameter in self.aion_model.parameters():
+ parameter.requires_grad_(False)
+ if self._shared_encoder_trainable:
+ for module in self._shared_encoder_modules():
+ for parameter in module.parameters():
+ parameter.requires_grad_(True)
+
+ self.spectrum_codec.float()
+ self.spectrum_codec.eval()
+ for parameter in self.spectrum_codec.parameters():
+ parameter.requires_grad_(False)
+
+ def _apply(self, fn, recurse=True):
+ """Keep the codec in float32 through blanket dtype casts.
+
+ ``nn.Module.to()`` reaches children via ``_apply``, so a child-level
+ ``to()`` override never sees it and the trainer's
+ ``self.model.encoders.to(dtype=target_dtype)`` rounds the codec.
+ Calling ``.float()`` afterwards restores the dtype but not the values,
+ and the codec quantizes hard: rounded weights can change token ids.
+ Device moves pass through untouched. Mirrors ``DNAEncoder._apply``.
+ """
+ is_reduced_precision = False
+ try:
+ probe = fn(torch.zeros(1, dtype=torch.float32))
+ if isinstance(probe, torch.Tensor) and probe.dtype not in (
+ torch.float32,
+ torch.float64,
+ ):
+ is_reduced_precision = True
+ except Exception:
+ pass
+
+ if is_reduced_precision:
+ for child in self.children():
+ if child is not self.spectrum_codec:
+ child._apply(fn, recurse=recurse)
+ for key, buf in self._buffers.items():
+ if buf is not None:
+ self._buffers[key] = fn(buf)
+ return self
+
+ return super()._apply(fn, recurse=recurse)
+
+ @property
+ def shared_encoder_trainable(self) -> bool:
+ """Whether the shared AION transformer is currently trainable."""
+ return self._shared_encoder_trainable
+
+ def set_shared_encoder_trainable(self, trainable: bool) -> AIONSpectrumEncoder:
+ """Toggle only the differentiable AION spectrum/shared-encoder path."""
+
+ self._shared_encoder_trainable = bool(trainable)
+ self._enforce_trainability()
+ self.train(self.training)
+ return self
+
+ def train(self, mode: bool = True) -> AIONSpectrumEncoder:
+ """Set runtime mode while preserving codec dtype/freeze invariants."""
+
+ super().train(mode)
+ self._enforce_trainability()
+ self.aion_model.eval()
+ if mode and self._shared_encoder_trainable:
+ for module in self._shared_encoder_modules():
+ module.train(True)
+ self.spectrum_codec.eval()
+ return self
+
+ def tokens_per_instance(self) -> int:
+ """Return the fixed number of spectrum tokens the AION codec emits."""
+ return self.num_spectrum_tokens
+
+ def _validate_inputs(self, inputs: Mapping[str, torch.Tensor]) -> None:
+ if not isinstance(inputs, Mapping):
+ raise TypeError(
+ "AION spectrum input must be a mapping of flux/ivar/mask/wavelength "
+ f"tensors, got {type(inputs).__name__}"
+ )
+ missing = set(AION_SPECTRUM_FIELDS).difference(inputs)
+ if missing:
+ raise ValueError(f"AION spectrum input is missing fields: {sorted(missing)}")
+
+ tensors = {field: inputs[field] for field in AION_SPECTRUM_FIELDS}
+ for field, tensor in tensors.items():
+ if not isinstance(tensor, torch.Tensor):
+ raise TypeError(
+ f"AION spectrum field {field!r} must be a torch.Tensor, "
+ f"got {type(tensor).__name__}"
+ )
+ expected_tail = (self.raw_spectrum_length,)
+ if tensor.ndim != 2 or tuple(tensor.shape[1:]) != expected_tail:
+ raise ValueError(
+ f"AION spectrum field {field!r} must have shape "
+ f"(B, {self.raw_spectrum_length}), got {tuple(tensor.shape)}"
+ )
+
+ shapes = {tuple(tensor.shape) for tensor in tensors.values()}
+ if len(shapes) != 1:
+ raise ValueError("AION spectrum fields must have identical batch and length shapes")
+ devices = {tensor.device for tensor in tensors.values()}
+ if len(devices) != 1:
+ raise ValueError("AION spectrum fields must be on the same device")
+
+ for field in ("flux", "ivar", "wavelength"):
+ tensor = tensors[field]
+ if not torch.is_floating_point(tensor):
+ raise TypeError(
+ f"AION spectrum field {field!r} must be floating point, got {tensor.dtype}"
+ )
+ if not bool(torch.isfinite(tensor).all()):
+ raise ValueError(f"AION spectrum field {field!r} contains non-finite values")
+
+ if tensors["mask"].dtype != torch.bool:
+ raise TypeError(
+ f"AION spectrum field 'mask' must be boolean, got {tensors['mask'].dtype}"
+ )
+ if bool((tensors["ivar"] < 0).any()):
+ raise ValueError("AION spectrum inverse variance must be non-negative")
+ if bool((torch.diff(tensors["wavelength"], dim=-1) <= 0).any()):
+ raise ValueError("AION spectrum wavelength must be strictly increasing")
+
+ def forward(self, inputs: Mapping[str, torch.Tensor]) -> torch.Tensor:
+ """Encode flux/ivar/mask/wavelength into ``(B, 273, 768)`` features."""
+ self._validate_inputs(inputs)
+ self._enforce_trainability()
+
+ reference = inputs["flux"]
+ with torch.no_grad(), torch.autocast(device_type=reference.device.type, enabled=False):
+ spectrum = self._desi_spectrum_type(
+ flux=inputs["flux"].detach().to(dtype=torch.float32).clone(),
+ ivar=inputs["ivar"].detach().to(dtype=torch.float32).clone(),
+ mask=inputs["mask"].detach().clone(),
+ wavelength=inputs["wavelength"].detach().to(dtype=torch.float32).clone(),
+ )
+ spectrum_token_ids = self.spectrum_codec.encode(spectrum)
+
+ expected_tokens = (reference.shape[0], self.num_spectrum_tokens)
+ if tuple(spectrum_token_ids.shape) != expected_tokens:
+ raise RuntimeError(
+ f"AION SpectrumCodec must emit token ids with shape {expected_tokens}, "
+ f"got {tuple(spectrum_token_ids.shape)}"
+ )
+
+ encoded_inputs = {self.spectrum_token_key: spectrum_token_ids}
+ if self._shared_encoder_trainable:
+ features = self.aion_model.encode(
+ encoded_inputs,
+ num_encoder_tokens=self.num_spectrum_tokens,
+ )
+ else:
+ with torch.no_grad():
+ features = self.aion_model.encode(
+ encoded_inputs,
+ num_encoder_tokens=self.num_spectrum_tokens,
+ )
+
+ expected_features = (reference.shape[0], self.num_spectrum_tokens, self.output_dim)
+ if tuple(features.shape) != expected_features:
+ raise RuntimeError(
+ f"AION shared encoder must emit features with shape {expected_features}, "
+ f"got {tuple(features.shape)}"
+ )
+ return features
+
+
+__all__ = [
+ "AIONSpectrumEncoder",
+ "AION_MODEL_REVISION",
+ "AION_RAW_SPECTRUM_LENGTH",
+ "AION_REPO_ID",
+ "AION_SPECTRUM_CODEC_REVISION",
+ "AION_SPECTRUM_FIELDS",
+ "AION_SPECTRUM_TOKENS",
+ "AION_SPECTRUM_WIDTH",
+]
diff --git a/src/modalities.py b/src/modalities.py
index 3a57b3e..7398209 100644
--- a/src/modalities.py
+++ b/src/modalities.py
@@ -41,6 +41,7 @@ class Modality(str, Enum):
TEXT = "text"
IMAGE = "image"
+ SPECTRUM = "spectrum"
TABLE = "table"
TIME_SERIES = "time_series"
GEOMETRY = "geometry"
@@ -85,6 +86,8 @@ def make_dummy_batch(
- text: handled by tokenizer downstream (returns None here)
- image: (3, 224, 224) — SigLIP2 input shape, matches `_process_image`
`default_tensor` in src/data/multimodal.py
+ - spectrum: dict of flux/ivar/mask/wavelength tensors, each with shape
+ (7781,), matching the raw DESI input expected by AION's spectrum codec
- table: (128,) int64 token ids — TAPAS input length, matches
`default_ids` in `_process_table_*`
- time_series: (64, 1) float — Moirai single-channel window (the
@@ -106,6 +109,13 @@ def make_dummy_batch(
return None
if m is Modality.IMAGE:
return torch.zeros(3, 224, 224)
+ if m is Modality.SPECTRUM:
+ return {
+ "flux": torch.zeros(7781),
+ "ivar": torch.ones(7781),
+ "mask": torch.zeros(7781, dtype=torch.bool),
+ "wavelength": torch.linspace(3600.0, 9800.0, 7781),
+ }
if m is Modality.TABLE:
return torch.zeros(128, dtype=torch.long)
if m is Modality.TIME_SERIES:
diff --git a/src/model.py b/src/model.py
index 8f76e7e..b8d767b 100644
--- a/src/model.py
+++ b/src/model.py
@@ -40,6 +40,9 @@
logger = logging.getLogger(__name__)
from .encoders import (
+ AIONImageEncoder,
+ AIONMultimodalEncoder,
+ AIONSpectrumEncoder,
DNAEncoder,
GeometryEncoder,
GraphEncoder,
@@ -192,6 +195,12 @@ class UnifiedTransformer(nn.Module):
def __init__(self, config: ModelConfig):
super().__init__()
self.config = config
+ self._joint_aion_enabled = bool(
+ Modality.IMAGE in config.modalities
+ and Modality.SPECTRUM in config.modalities
+ and config.image_encoder_type == "aion"
+ and config.spectrum_encoder_type == "aion"
+ )
self.is_vla = config.is_vla
self._backbone_load_error = None
@@ -450,6 +459,9 @@ def __init__(self, config: ModelConfig):
# --- Stage 1: Modality Encoders ---
self.encoders = nn.ModuleDict()
self.projectors = nn.ModuleDict()
+ # A modality normally owns an encoder with the same key. Joint AION
+ # maps image and spectrum to one registered owner instead.
+ self._encoder_key_by_modality: dict[str, str] = {}
# Determine target dimension for projectors
target_dim = self.backbone_dim
@@ -472,6 +484,39 @@ def __init__(self, config: ModelConfig):
proj_hidden_mult = getattr(config, "projector_hidden_mult", 1)
proj_num_layers = getattr(config, "projector_num_layers", 2)
+ projector_defaults = {
+ "norm_mode": proj_norm_mode,
+ "target_norm": proj_target_norm,
+ "modality_embed_pos": proj_embed_pos,
+ "modality_embed_scale": proj_embed_scale,
+ "text_norm_mean": proj_text_norm_mean,
+ "text_norm_std": proj_text_norm_std,
+ "text_elem_mean": proj_text_elem_mean,
+ "text_elem_std": proj_text_elem_std,
+ "norm_clip_min": proj_norm_clip_min,
+ "norm_clip_max": proj_norm_clip_max,
+ "hidden_mult": proj_hidden_mult,
+ "num_layers": proj_num_layers,
+ }
+
+ def projector_kwargs(modality: str) -> dict[str, object]:
+ """Resolve explicit modality settings over legacy global defaults."""
+
+ resolved = dict(projector_defaults)
+ if modality in {"image", "spectrum"}:
+ for field in (
+ "norm_mode",
+ "target_norm",
+ "modality_embed_pos",
+ "modality_embed_scale",
+ ):
+ value = getattr(
+ config, f"{modality}_projector_{field}", None
+ )
+ if value is not None:
+ resolved[field] = value
+ return resolved
+
logger.info(
f"Projector config: norm_mode={proj_norm_mode}, target_norm={proj_target_norm}, embed_pos={proj_embed_pos}, embed_scale={proj_embed_scale}, hidden_mult={proj_hidden_mult}, num_layers={proj_num_layers}"
)
@@ -544,27 +589,90 @@ def __init__(self, config: ModelConfig):
num_layers=proj_num_layers,
)
+ # Joint AION owns the shared transformer exactly once. The standalone
+ # wrappers remain the unchanged path for single-modality models.
+ if self._joint_aion_enabled:
+ joint_aion_encoder = AIONMultimodalEncoder(
+ model_name_or_path=config.aion_model_name_or_path,
+ model_revision=config.aion_model_revision,
+ image_codec_revision=config.aion_codec_revision,
+ spectrum_codec_revision=config.aion_spectrum_codec_revision,
+ cache_dir=config.aion_cache_dir,
+ local_files_only=config.aion_local_files_only,
+ bands=config.aion_image_bands,
+ raw_image_size=config.aion_raw_image_size,
+ raw_spectrum_length=config.aion_raw_spectrum_length,
+ )
+ self.encoders["aion"] = joint_aion_encoder
+ self._encoder_key_by_modality.update(
+ {"image": "aion", "spectrum": "aion"}
+ )
+ image_tokens = joint_aion_encoder.tokens_per_instance("image")
+ spectrum_tokens = joint_aion_encoder.tokens_per_instance("spectrum")
+ if image_tokens != config.aion_num_image_tokens:
+ raise RuntimeError(
+ "Configured AION image token count does not match the joint encoder: "
+ f"config={config.aion_num_image_tokens}, encoder={image_tokens}."
+ )
+ if spectrum_tokens != config.aion_num_spectrum_tokens:
+ raise RuntimeError(
+ "Configured AION spectrum token count does not match the joint encoder: "
+ f"config={config.aion_num_spectrum_tokens}, encoder={spectrum_tokens}."
+ )
+
# Image
if Modality.IMAGE in config.modalities:
- self.encoders["image"] = ImageEncoder(
- d_img=config.d_img,
- model_name=config.image_encoder_id,
- )
+ if not self._joint_aion_enabled:
+ if config.image_encoder_type == "aion":
+ self.encoders["image"] = AIONImageEncoder(
+ model_name_or_path=config.aion_model_name_or_path,
+ model_revision=config.aion_model_revision,
+ codec_revision=config.aion_codec_revision,
+ cache_dir=config.aion_cache_dir,
+ local_files_only=config.aion_local_files_only,
+ bands=config.aion_image_bands,
+ raw_image_size=config.aion_raw_image_size,
+ )
+ emitted_tokens = self.encoders["image"].tokens_per_instance()
+ if emitted_tokens != config.aion_num_image_tokens:
+ raise RuntimeError(
+ "Configured AION token count does not match the loaded encoder: "
+ f"config={config.aion_num_image_tokens}, encoder={emitted_tokens}."
+ )
+ else:
+ self.encoders["image"] = ImageEncoder(
+ d_img=config.d_img,
+ model_name=config.image_encoder_id,
+ )
+ self._encoder_key_by_modality["image"] = "image"
self.projectors["image"] = ModalityProjector(
config.d_img,
target_dim,
- norm_mode=proj_norm_mode,
- target_norm=proj_target_norm,
- modality_embed_pos=proj_embed_pos,
- modality_embed_scale=proj_embed_scale,
- text_norm_mean=proj_text_norm_mean,
- text_norm_std=proj_text_norm_std,
- text_elem_mean=proj_text_elem_mean,
- text_elem_std=proj_text_elem_std,
- norm_clip_min=proj_norm_clip_min,
- norm_clip_max=proj_norm_clip_max,
- hidden_mult=proj_hidden_mult,
- num_layers=proj_num_layers,
+ **projector_kwargs("image"),
+ )
+
+ # Spectrum
+ if Modality.SPECTRUM in config.modalities:
+ if not self._joint_aion_enabled:
+ self.encoders["spectrum"] = AIONSpectrumEncoder(
+ model_name_or_path=config.aion_model_name_or_path,
+ model_revision=config.aion_model_revision,
+ codec_revision=config.aion_spectrum_codec_revision,
+ cache_dir=config.aion_cache_dir,
+ local_files_only=config.aion_local_files_only,
+ raw_spectrum_length=config.aion_raw_spectrum_length,
+ )
+ self._encoder_key_by_modality["spectrum"] = "spectrum"
+ emitted_tokens = self.encoders["spectrum"].tokens_per_instance()
+ if emitted_tokens != config.aion_num_spectrum_tokens:
+ raise RuntimeError(
+ "Configured AION spectrum token count does not match the loaded encoder: "
+ f"config={config.aion_num_spectrum_tokens}, encoder={emitted_tokens}."
+ )
+ self.projectors["spectrum"] = ModalityProjector(
+ config.d_spectrum,
+ target_dim,
+ **projector_kwargs("spectrum"),
)
# Geometry
@@ -722,18 +830,20 @@ def __init__(self, config: ModelConfig):
self.action_head = RegressionDecoder(target_dim, config.action_dim)
- # Freeze Encoders if requested (Zone 1)
+ # Freeze encoders through capability-aware hooks. AION must keep its
+ # hard-quantized codec frozen even when its shared encoder is trainable;
+ # a blanket requires_grad mutation would be undone by its invariant
+ # enforcement on the next forward pass.
if config.freeze_encoders:
logger.info("Freezing Modality Encoders...")
- for param in self.encoders.parameters():
- param.requires_grad = False
-
-
- # Cache frozen state for each encoder (avoids O(params) check in forward)
- self._encoder_frozen_cache = {
- modality: not any(p.requires_grad for p in encoder.parameters())
- for modality, encoder in self.encoders.items()
- }
+ self.set_encoders_trainable(False)
+ else:
+ # Most encoders are trainable when constructed. AION deliberately
+ # starts frozen, so explicitly opt its shared encoder into training.
+ for modality, encoder in self.encoders.items():
+ if callable(getattr(encoder, "set_shared_encoder_trainable", None)):
+ self.set_encoder_trainable(modality, True)
+ self.refresh_encoder_frozen_cache()
# --- Stage 3: Unified Transformer (Legacy / Custom) ---
if self.backbone is None:
@@ -751,6 +861,99 @@ def __init__(self, config: ModelConfig):
self.head = CausalLMHead(config.d_model, config.vocab_size)
+ # Sub-modalities that are encoded by their parent's encoder. DNA splits one
+ # "dna" entry into reference/variant spans that share a single DNAEncoder.
+ _PARENT_MODALITY_KEYS = {"dna_reference": "dna", "dna_variant": "dna"}
+
+ def _encoder_key_for_modality(self, modality: str | Modality) -> str:
+ modality_key = modality.value if isinstance(modality, Modality) else str(modality)
+ # getattr: models built via __new__ (routing unit tests) never ran
+ # __init__, so the ownership map can be absent. Same-name ownership is
+ # the correct default there.
+ owners = getattr(self, "_encoder_key_by_modality", None) or {}
+ if modality_key in owners:
+ return owners[modality_key]
+ return self._PARENT_MODALITY_KEYS.get(modality_key, modality_key)
+
+ def _tokens_per_instance_for_modality(self, modality: str | Modality) -> int:
+ """Resolve a modality token count through its registered encoder owner.
+
+ Standalone encoders expose ``tokens_per_instance()`` while a shared
+ multimodal owner needs the routed modality as
+ ``tokens_per_instance(modality)``. The ownership map is the source of
+ truth for selecting between those APIs, matching feature routing.
+ """
+
+ modality_key = modality.value if isinstance(modality, Modality) else str(modality)
+ encoder_key = self._encoder_key_for_modality(modality_key)
+ if encoder_key not in self.encoders:
+ raise KeyError(
+ f"Unknown encoder for modality {modality_key!r}; resolved owner "
+ f"{encoder_key!r}, available encoders: {list(self.encoders.keys())}"
+ )
+
+ tokens_per_instance = self.encoders[encoder_key].tokens_per_instance
+ if encoder_key == modality_key:
+ return int(tokens_per_instance())
+ return int(tokens_per_instance(modality_key))
+
+ def refresh_encoder_frozen_cache(self) -> dict[str, bool]:
+ """Refresh the per-modality cache while deduplicating shared owners."""
+
+ owner_frozen = {
+ encoder_key: not any(
+ parameter.requires_grad for parameter in encoder.parameters()
+ )
+ for encoder_key, encoder in self.encoders.items()
+ }
+ cache: dict[str, bool] = {}
+ for modality in self.config.modalities:
+ modality_key = str(modality)
+ encoder_key = self._encoder_key_for_modality(modality_key)
+ if encoder_key in owner_frozen:
+ cache[modality_key] = owner_frozen[encoder_key]
+ self._encoder_frozen_cache = cache
+ return dict(cache)
+
+ def set_encoder_trainable(self, modality: str | Modality, trainable: bool) -> int:
+ """Set one modality encoder (or shared owner) trainable safely.
+
+ Joint AION image and spectrum both resolve to the same registered
+ owner. Its capability hook toggles the shared transformer and both
+ modality embeddings together while keeping both codecs protected.
+ """
+
+ modality_key = modality.value if isinstance(modality, Modality) else str(modality)
+ encoder_key = self._encoder_key_for_modality(modality_key)
+ if encoder_key not in self.encoders:
+ raise KeyError(
+ f"Unknown encoder {modality_key!r}; available encoders: "
+ f"{list(self.encoders.keys())}"
+ )
+
+ encoder = self.encoders[encoder_key]
+ specialized_hook = getattr(encoder, "set_shared_encoder_trainable", None)
+ if callable(specialized_hook):
+ specialized_hook(bool(trainable))
+ else:
+ for parameter in encoder.parameters():
+ parameter.requires_grad_(bool(trainable))
+
+ self.refresh_encoder_frozen_cache()
+ return sum(
+ parameter.numel() for parameter in encoder.parameters() if parameter.requires_grad
+ )
+
+ def set_encoders_trainable(self, trainable: bool) -> dict[str, int]:
+ """Set each registered encoder owner once, including joint AION."""
+
+ trainable_counts = {
+ encoder_key: self.set_encoder_trainable(encoder_key, trainable)
+ for encoder_key in self.encoders
+ }
+ self.refresh_encoder_frozen_cache()
+ return trainable_counts
+
def get_input_embeddings(self, inputs):
"""Helper to get embeddings from varied inputs."""
# Retrieve input_ids from 'text' if available
@@ -885,13 +1088,21 @@ def _process_multimodal_embeddings(self, inputs):
# Padding must not affect projector batch statistics or
# consume encoder work, even before the sequence is merged.
encoder_input = encoder_input.flatten(0, 1)[reference_valid]
- features = self.encoders[modality](encoder_input)
- pad_mask = getattr(self.encoders[modality], "_last_pad_mask", None)
-
+ # Joint AION routes image/spectrum through one registered
+ # owner; every other encoder keeps its direct path.
+ modality_key = modality.value if isinstance(modality, Modality) else str(modality)
+ _encoder = self.encoders[self._encoder_key_for_modality(modality_key)]
+ _encode_modality = getattr(_encoder, "encode_modality", None)
+ if callable(_encode_modality):
+ features = _encode_modality(modality_key, encoder_input)
+ else:
+ features = _encoder(encoder_input)
+ pad_mask = getattr(_encoder, "_last_pad_mask", None)
- # Detach if encoder is frozen (avoids backward through encoder)
- # Use cached frozen state to avoid O(params) check every forward
- if getattr(self, "_encoder_frozen_cache", {}).get(modality, False):
+ # Detach if encoder is frozen (avoids backward through encoder).
+ if getattr(self, "_encoder_frozen_cache", {}).get(
+ modality_key, False
+ ):
features = features.detach()
@@ -1015,8 +1226,7 @@ def _merge_text_input_ids_with_modality_embeds(
for i, modality in enumerate(modalities):
start_mask = start_modality_idx == i
if start_mask.any():
- encoder_key = _PARENT_MODALITY.get(modality, modality)
- tokens_per = int(self.encoders[encoder_key].tokens_per_instance())
+ tokens_per = self._tokens_per_instance_for_modality(modality)
slot_size[start_mask] = tokens_per
# sanity checks
@@ -1083,9 +1293,7 @@ def _merge_text_input_ids_with_modality_embeds(
final_attention_mask[batch_indices, text_positions] = 1
tokens_per_instance = {
- modality: int(
- self.encoders[_PARENT_MODALITY.get(modality, modality)].tokens_per_instance()
- )
+ modality: self._tokens_per_instance_for_modality(modality)
for modality in modalities
}
expected_tokens = {
@@ -1633,13 +1841,9 @@ def forward(
if getattr(self.backbone.config, "_attn_implementation", "sdpa") == "eager":
# Pre-computed 4D causal mask: lower-triangular 0.0, upper -inf
# Shape: (1, 1, T, T) — broadcasts over batch and heads
- causal_mask_fwd = torch.zeros(
- 1, 1, T_fwd, T_fwd, device=x.device, dtype=x.dtype
- )
+ causal_mask_fwd = torch.zeros(1, 1, T_fwd, T_fwd, device=x.device, dtype=x.dtype)
causal_mask_fwd.masked_fill_(
- ~torch.tril(
- torch.ones(T_fwd, T_fwd, device=x.device, dtype=torch.bool)
- ),
+ ~torch.tril(torch.ones(T_fwd, T_fwd, device=x.device, dtype=torch.bool)),
torch.finfo(x.dtype).min,
)
if mask_padded:
@@ -1740,7 +1944,6 @@ def forward(
# ... (Original Logic) ...
B, T_total, _ = x.shape
-
# 1. Generate RoPE frequencies
freqs_cis = self.rope(x, seq_len=T_total)
@@ -1777,11 +1980,9 @@ def forward(
x, aux_loss = block(x, mask=mask, freqs_cis=freqs_cis)
total_aux_loss += aux_loss
-
x = self.ln_f(x)
logits = self.head(x)
-
# Compute SFT loss manually if labels provided
loss = total_aux_loss
if labels is not None:
@@ -1803,9 +2004,7 @@ def forward(
# Same as Path A pad-mask (line 858).
pad_id_local = self._resolve_pad_id()
if pad_id_local is not None:
- assigned = assigned.masked_fill(
- assigned == pad_id_local, -100
- )
+ assigned = assigned.masked_fill(assigned == pad_id_local, -100)
full_labels[:, -text_len:] = assigned
shift_labels = full_labels[..., 1:].contiguous()
@@ -1829,7 +2028,7 @@ def generate(self, inputs, max_new_tokens=20, **kwargs):
but not the answer (which is what we want to generate).
"""
is_interleaved_qa = bool(self.config.is_interleaved_qa)
-
+
if "table" in inputs:
logger.warning("!!! DEBUG: Table Input Detected in Forward Pass !!!")
@@ -1853,14 +2052,8 @@ def generate(self, inputs, max_new_tokens=20, **kwargs):
# Get pad_token_id from tokenizer if available
pad_token_id = kwargs.pop("pad_token_id", None)
- if (
- pad_token_id is None
- and hasattr(self, "tokenizer")
- and self.tokenizer is not None
- ):
- pad_token_id = (
- self.tokenizer.pad_token_id or self.tokenizer.eos_token_id
- )
+ if pad_token_id is None and hasattr(self, "tokenizer") and self.tokenizer is not None:
+ pad_token_id = self.tokenizer.pad_token_id or self.tokenizer.eos_token_id
# Create attention mask, properly masking any pad tokens in the text.
# The prefix (image/modality) tokens are always valid (no padding).
@@ -1948,6 +2141,5 @@ def load_pretrained_weights(self, hf_model_id: str):
logger.info("Skipping load_pretrained_weights (Backbone already loaded).")
return
-
# ... (Previous Logic preserved if needed) ...
pass
diff --git a/src/modules/projector.py b/src/modules/projector.py
index d2db20e..934fd4a 100644
--- a/src/modules/projector.py
+++ b/src/modules/projector.py
@@ -51,6 +51,7 @@ class ModalityProjector(nn.Module):
Configurable Normalization Modes (for ablation studies):
- none: No normalization (simple MLP output)
- layernorm: Standard LayerNorm on output (original/default behavior)
+ - text_scaled_layernorm: LayerNorm initialized to the target text-token norm
- rmsnorm: RMSNorm (like LLaMA/OLMo, no centering)
- l2_token: Per-token L2 normalization to target_norm
- l2_sequence: Sequence-level L2 normalization (preserves token differences)
@@ -67,6 +68,7 @@ class ModalityProjector(nn.Module):
VALID_NORM_MODES = {
"none",
"layernorm",
+ "text_scaled_layernorm",
"rmsnorm",
"l2_token",
"l2_sequence",
@@ -181,6 +183,11 @@ def __init__(
if norm_mode == "layernorm":
self.final_norm = nn.LayerNorm(d_model)
+ elif norm_mode == "text_scaled_layernorm":
+ self.final_norm = nn.LayerNorm(d_model)
+ initial_gain = target_norm / math.sqrt(d_model)
+ nn.init.constant_(self.final_norm.weight, initial_gain)
+ nn.init.zeros_(self.final_norm.bias)
elif norm_mode == "rmsnorm":
self.final_norm = RMSNorm(d_model)
elif norm_mode in ["l2_token", "l2_sequence", "scale_only"]:
@@ -258,6 +265,11 @@ def _apply_normalization(self, x: torch.Tensor) -> torch.Tensor:
# Standard LayerNorm (original behavior)
return self.final_norm(x)
+ elif self.norm_mode == "text_scaled_layernorm":
+ # LayerNorm's unit-variance output has L2 norm sqrt(d_model), so
+ # the initialized gain makes the merge-point token norm target_norm.
+ return self.final_norm(x)
+
elif self.norm_mode == "rmsnorm":
# RMSNorm (like LLaMA/OLMo)
return self.final_norm(x)
diff --git a/src/train.py b/src/train.py
index 451175e..b572d83 100644
--- a/src/train.py
+++ b/src/train.py
@@ -2,6 +2,7 @@
import os
import random
import sys
+from pathlib import Path
from typing import Any, cast
import hydra
@@ -9,7 +10,7 @@
import torch
import torch.distributed as dist
from omegaconf import DictConfig, OmegaConf
-from torch.utils.data import DataLoader
+from torch.utils.data import DataLoader, Subset
# Path Safety
# Prioritize local source code over installed packages.
@@ -17,6 +18,25 @@
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from src.config import DYNAMIC_LENGTH_TS_PROJECTORS, ModelConfig, TrainingConfig
+from src.data.astroprism_joint_batch_sampler import AstroPRISMJointBatchSampler
+from src.data.astroprism_joint_collate import AstroPRISMJointCollator
+from src.data.astroprism_joint_dataset import AstroPRISMJointDataset
+from src.data.astroprism_stage1a_collate import AstroPRISMStage1ACollator
+from src.data.astroprism_stage1a_dataset import (
+ DEFAULT_SYSTEM_PROMPT,
+ DEFAULT_USER_PROMPT,
+ AstroPRISMStage1ADataset,
+ AstroPRISMStage1ADatasetConfig,
+)
+from src.data.astroprism_stage1b_spectrum_collate import (
+ AstroPRISMStage1BSpectrumCollator,
+)
+from src.data.astroprism_stage1b_spectrum_dataset import (
+ DEFAULT_SPECTRUM_SYSTEM_PROMPT,
+ DEFAULT_SPECTRUM_USER_PROMPT,
+ AstroPRISMStage1BSpectrumDataset,
+ AstroPRISMStage1BSpectrumDatasetConfig,
+)
from src.data.calvin_vla import CalvinVLADataset
from src.data.collate import BucketedCollator, MultimodalCollator
from src.data.multi_webdataset import (
@@ -35,13 +55,24 @@
setup_distributed,
)
from src.training.trainer_native import train_native_ddp
-from src.training.trainer_zone_a import ZoneATrainer
+from src.training.trainer_zone_a import (
+ ZoneATrainer,
+)
from src.training.trainer_zone_a_vla import ZoneAVLATrainer
from src.utils.banner import print_prism_banner
logger = logging.getLogger(__name__)
+_ASTROPRISM_MAP_STYLE_TASKS = frozenset(
+ {
+ "astroprism_stage1a",
+ "astroprism_stage1b_spectrum",
+ "astroprism_joint_aion",
+ }
+)
+
+
def _coerce_str_or_none(value):
"""Return value as a plain str, or None if the value is None or an empty
OmegaConf scalar container.
@@ -128,6 +159,269 @@ def _resolve_timeomni_strides(
return [int(x) for x in strides]
+def _build_astroprism_stage1a_data(data_cfg, tokenizer):
+ """Construct the strict raw-image Stage 1A dataset and collator."""
+ if data_cfg is None or not hasattr(data_cfg, "get"):
+ raise RuntimeError(
+ "training.task=astroprism_stage1a requires the data=astroprism_stage1a "
+ "Hydra config group"
+ )
+
+ def required_path(name: str) -> Path:
+ value = data_cfg.get(name, None)
+ if value is None or not str(value).strip():
+ raise RuntimeError(
+ f"training.task=astroprism_stage1a requires data.{name}"
+ )
+ return Path(hydra.utils.to_absolute_path(str(value)))
+
+ candidates_path_value = data_cfg.get("candidates_path", None)
+ candidates_path = (
+ Path(hydra.utils.to_absolute_path(str(candidates_path_value)))
+ if candidates_path_value is not None and str(candidates_path_value).strip()
+ else None
+ )
+ max_merged_seq_length = int(data_cfg.get("max_merged_seq_length", 2048))
+ accepted_statuses = tuple(
+ str(status)
+ for status in data_cfg.get("accepted_statuses", ("accepted", "approved"))
+ )
+ dataset_config = AstroPRISMStage1ADatasetConfig(
+ raw_manifest_path=required_path("raw_manifest_path"),
+ candidates_manifest_path=required_path("candidates_manifest_path"),
+ candidates_path=candidates_path,
+ split=_coerce_str_or_none(data_cfg.get("split", None)),
+ max_merged_seq_length=max_merged_seq_length,
+ caption_field=str(data_cfg.get("caption_field", "auto")),
+ system_prompt=str(data_cfg.get("system_prompt", DEFAULT_SYSTEM_PROMPT)),
+ user_prompt=str(data_cfg.get("user_prompt", DEFAULT_USER_PROMPT)),
+ accepted_caption_statuses=tuple(
+ str(status)
+ for status in data_cfg.get(
+ "accepted_caption_statuses", ("released_reused", "released_revised")
+ )
+ ),
+ allow_unsafe_candidates=bool(
+ data_cfg.get("allow_unsafe_candidates", False)
+ ),
+ accepted_statuses=accepted_statuses,
+ verify_file_hashes=bool(data_cfg.get("verify_file_hashes", False)),
+ )
+ dataset = AstroPRISMStage1ADataset(dataset_config, tokenizer)
+ collator = AstroPRISMStage1ACollator(
+ tokenizer,
+ max_merged_seq_length=max_merged_seq_length,
+ )
+ return dataset, collator
+
+
+def _validate_astroprism_stage1a_model_contract(
+ task: str,
+ model_config: ModelConfig,
+) -> None:
+ """Fail closed when Stage 1A is paired with a different model contract."""
+ if task != "astroprism_stage1a":
+ return
+ if model_config.image_encoder_type != "aion":
+ raise RuntimeError(
+ "training.task=astroprism_stage1a requires "
+ "model.image_encoder_type=aion"
+ )
+ modalities = {str(modality) for modality in model_config.modalities}
+ expected_modalities = {"image", "text"}
+ if modalities != expected_modalities:
+ raise RuntimeError(
+ "training.task=astroprism_stage1a requires exactly "
+ f"model.modalities={sorted(expected_modalities)}, got {sorted(modalities)}"
+ )
+
+
+def _build_astroprism_stage1b_spectrum_data(data_cfg, tokenizer):
+ """Construct reviewed spectrum-caption Stage 1B data and its collator."""
+ task = "astroprism_stage1b_spectrum"
+ if data_cfg is None or not hasattr(data_cfg, "get"):
+ raise RuntimeError(
+ f"training.task={task} requires the data={task} Hydra config group"
+ )
+
+ def required_path(name: str) -> Path:
+ value = data_cfg.get(name, None)
+ if value is None or not str(value).strip():
+ raise RuntimeError(f"training.task={task} requires data.{name}")
+ return Path(hydra.utils.to_absolute_path(str(value)))
+
+ captions_path_value = data_cfg.get("captions_path", None)
+ captions_path = (
+ Path(hydra.utils.to_absolute_path(str(captions_path_value)))
+ if captions_path_value is not None and str(captions_path_value).strip()
+ else None
+ )
+ max_merged_seq_length = int(data_cfg.get("max_merged_seq_length", 2048))
+ dataset_config = AstroPRISMStage1BSpectrumDatasetConfig(
+ raw_manifest_path=required_path("raw_manifest_path"),
+ captions_manifest_path=required_path("captions_manifest_path"),
+ captions_path=captions_path,
+ split=_coerce_str_or_none(data_cfg.get("split", "train")),
+ max_merged_seq_length=max_merged_seq_length,
+ system_prompt=str(
+ data_cfg.get("system_prompt", DEFAULT_SPECTRUM_SYSTEM_PROMPT)
+ ),
+ user_prompt=str(data_cfg.get("user_prompt", DEFAULT_SPECTRUM_USER_PROMPT)),
+ accepted_caption_statuses=tuple(
+ str(status)
+ for status in data_cfg.get("accepted_caption_statuses", ("accepted",))
+ ),
+ verify_file_hashes=bool(data_cfg.get("verify_file_hashes", False)),
+ )
+ dataset = AstroPRISMStage1BSpectrumDataset(dataset_config, tokenizer)
+ collator = AstroPRISMStage1BSpectrumCollator(
+ tokenizer,
+ max_merged_seq_length=max_merged_seq_length,
+ )
+ return dataset, collator
+
+
+def _validate_astroprism_stage1b_spectrum_model_contract(
+ task: str,
+ model_config: ModelConfig,
+) -> None:
+ """Fail closed unless the spectrum proxy receives only text and spectra."""
+ if task != "astroprism_stage1b_spectrum":
+ return
+ if model_config.spectrum_encoder_type != "aion":
+ raise RuntimeError(
+ "training.task=astroprism_stage1b_spectrum requires "
+ "model.spectrum_encoder_type=aion"
+ )
+ modalities = {str(modality) for modality in model_config.modalities}
+ expected_modalities = {"spectrum", "text"}
+ if modalities != expected_modalities:
+ raise RuntimeError(
+ "training.task=astroprism_stage1b_spectrum requires exactly "
+ f"model.modalities={sorted(expected_modalities)}, got {sorted(modalities)}"
+ )
+
+
+def _nested_data_config(data_cfg, name: str, *, split: str):
+ """Copy one joint-data subsection and enforce the requested split."""
+
+ section = data_cfg.get(name, None) if data_cfg is not None else None
+ if section is None or not hasattr(section, "get"):
+ raise RuntimeError(
+ "training.task=astroprism_joint_aion requires "
+ f"data.{name} in the data=astroprism_joint_aion config group"
+ )
+ plain = OmegaConf.to_container(section, resolve=True)
+ if not isinstance(plain, dict):
+ raise RuntimeError(f"data.{name} must be a mapping")
+ plain["split"] = split
+ plain["max_merged_seq_length"] = int(
+ data_cfg.get("max_merged_seq_length", 2048)
+ )
+ return OmegaConf.create(plain)
+
+
+def _build_astroprism_joint_data(data_cfg, tokenizer, *, split: str):
+ """Build one split of the joint, homogeneous-batch AstroPRISM dataset."""
+
+ task = "astroprism_joint_aion"
+ if data_cfg is None or not hasattr(data_cfg, "get"):
+ raise RuntimeError(
+ f"training.task={task} requires the data={task} Hydra config group"
+ )
+ if split not in {"train", "validation", "test"}:
+ raise RuntimeError(f"Unsupported AstroPRISM joint split: {split!r}")
+
+ image_dataset, image_collator = _build_astroprism_stage1a_data(
+ _nested_data_config(data_cfg, "image", split=split),
+ tokenizer,
+ )
+ spectrum_dataset, spectrum_collator = _build_astroprism_stage1b_spectrum_data(
+ _nested_data_config(data_cfg, "spectrum", split=split),
+ tokenizer,
+ )
+ dataset = AstroPRISMJointDataset(
+ image_dataset,
+ spectrum_dataset,
+ split=split,
+ )
+ collator = AstroPRISMJointCollator(
+ image_collator,
+ spectrum_collator,
+ )
+ return dataset, collator
+
+
+def _build_astroprism_joint_eval_loaders(
+ dataset,
+ collator,
+ *,
+ batch_size: int,
+ num_workers: int,
+ pin_memory: bool,
+):
+ """Build deterministic, modality-specific held-out loaders.
+
+ These loaders intentionally remain unsharded: every distributed rank runs
+ the same 65-row modality split, which avoids padding duplicates while still
+ keeping all ranks in lockstep through the wrapped model forward.
+ """
+
+ loaders = {}
+ for modality, indices in dataset.modality_indices.items():
+ subset = Subset(dataset, list(indices))
+ loaders[modality] = DataLoader(
+ subset,
+ batch_size=batch_size,
+ shuffle=False,
+ collate_fn=collator,
+ num_workers=num_workers,
+ persistent_workers=num_workers > 0,
+ prefetch_factor=2 if num_workers > 0 else None,
+ pin_memory=pin_memory,
+ )
+ return loaders
+
+
+def _validate_astroprism_joint_model_contract(
+ task: str,
+ model_config: ModelConfig,
+) -> None:
+ """Fail closed unless the joint task uses both pinned AION raw paths."""
+
+ if task != "astroprism_joint_aion":
+ return
+ if model_config.image_encoder_type != "aion":
+ raise RuntimeError(
+ "training.task=astroprism_joint_aion requires model.image_encoder_type=aion"
+ )
+ if model_config.spectrum_encoder_type != "aion":
+ raise RuntimeError(
+ "training.task=astroprism_joint_aion requires model.spectrum_encoder_type=aion"
+ )
+ modalities = {str(modality) for modality in model_config.modalities}
+ expected_modalities = {"image", "spectrum", "text"}
+ if modalities != expected_modalities:
+ raise RuntimeError(
+ "training.task=astroprism_joint_aion requires exactly "
+ f"model.modalities={sorted(expected_modalities)}, got {sorted(modalities)}"
+ )
+
+
+def _validate_astroprism_joint_runtime_environment(task: str) -> None:
+ """Reject modes that bypass the joint sampler's rank sharding."""
+
+ if task != "astroprism_joint_aion":
+ return
+ zero_stage = os.environ.get("DEEPSPEED_ZERO_STAGE", "").strip()
+ if zero_stage:
+ raise RuntimeError(
+ "training.task=astroprism_joint_aion does not support DeepSpeed: "
+ "DEEPSPEED_ZERO_STAGE must be unset so Accelerate shards the "
+ "rank-synchronous joint batch sampler"
+ )
+
+
def _local_shards_active(
local_shards_dir: str, allowed_modality: str, model_modalities
) -> bool:
@@ -321,6 +615,33 @@ def build_model_config(cfg: DictConfig, rank: int = 0) -> ModelConfig:
image_size=cfg.model.get("image_size", 224),
image_mean=tuple(cfg.model.get("image_mean", (0.5, 0.5, 0.5))),
image_std=tuple(cfg.model.get("image_std", (0.5, 0.5, 0.5))),
+ # === AION (AstroPRISM image + spectrum) ===
+ # Without these the joint path silently falls back to SigLIP, which
+ # then tries to reach the Hub for an encoder the run never wanted.
+ image_encoder_type=cfg.model.get("image_encoder_type", "siglip"),
+ spectrum_encoder_type=cfg.model.get("spectrum_encoder_type", "aion"),
+ d_spectrum=cfg.model.get("d_spectrum", 768),
+ aion_model_name_or_path=cfg.model.get(
+ "aion_model_name_or_path", "polymathic-ai/aion-base"
+ ),
+ aion_model_revision=cfg.model.get(
+ "aion_model_revision", "40541618104bab0fa85c8af68daeb867a720bb8c"
+ ),
+ aion_codec_revision=cfg.model.get(
+ "aion_codec_revision", "40541618104bab0fa85c8af68daeb867a720bb8c"
+ ),
+ aion_spectrum_codec_revision=cfg.model.get(
+ "aion_spectrum_codec_revision", "40541618104bab0fa85c8af68daeb867a720bb8c"
+ ),
+ aion_cache_dir=cfg.model.get("aion_cache_dir", None),
+ aion_local_files_only=cfg.model.get("aion_local_files_only", True),
+ aion_image_bands=tuple(
+ cfg.model.get("aion_image_bands", ("DES-G", "DES-R", "DES-I", "DES-Z"))
+ ),
+ aion_num_image_tokens=cfg.model.get("aion_num_image_tokens", 576),
+ aion_raw_image_size=tuple(cfg.model.get("aion_raw_image_size", (160, 160))),
+ aion_num_spectrum_tokens=cfg.model.get("aion_num_spectrum_tokens", 273),
+ aion_raw_spectrum_length=cfg.model.get("aion_raw_spectrum_length", 7781),
# === Timeseries ===
is_timeseries=cfg.model.get("is_timeseries", False),
ts_encoder_id=cfg.model.get(
@@ -594,6 +915,15 @@ def main(cfg: DictConfig):
# 0. Determine execution path
use_native_fsdp = os.environ.get("USE_NATIVE_FSDP", "0") == "1"
use_native_ddp = os.environ.get("USE_NATIVE_DDP", "0") == "1"
+ requested_task = cfg.training.get("task", "vlm")
+ _validate_astroprism_joint_runtime_environment(requested_task)
+ if requested_task in _ASTROPRISM_MAP_STYLE_TASKS and (
+ use_native_fsdp or use_native_ddp
+ ):
+ raise RuntimeError(
+ f"AstroPRISM task {requested_task!r} uses a finite map-style dataset and currently "
+ "supports the Accelerate path only; unset USE_NATIVE_DDP/FSDP."
+ )
# Defaults for Accelerate path (setup_distributed is native-only).
def _env_int(keys, default):
@@ -650,6 +980,9 @@ def _env_int(keys, default):
hydra_out = os.getcwd()
model_config = build_model_config(cfg, rank=rank)
+ _validate_astroprism_stage1a_model_contract(requested_task, model_config)
+ _validate_astroprism_stage1b_spectrum_model_contract(requested_task, model_config)
+ _validate_astroprism_joint_model_contract(requested_task, model_config)
# Allow override of dimensions if specified in cfg (not implemented fully in hydra yaml yet, relying on defaults)
@@ -795,8 +1128,16 @@ def _env_int(keys, default):
except Exception as e:
_rprint(f"Analysis Failed: {e}", flush=True)
- task_type = cfg.training.get("task", "vlm")
+ task_type = requested_task
is_vla_task = bool(model_config.is_vla or task_type == "vla_calvin")
+ is_astroprism_stage1a_task = task_type == "astroprism_stage1a"
+ is_astroprism_stage1b_spectrum_task = task_type == "astroprism_stage1b_spectrum"
+ is_astroprism_joint_task = task_type == "astroprism_joint_aion"
+ is_astroprism_map_style_task = task_type in _ASTROPRISM_MAP_STYLE_TASKS
+ if is_astroprism_map_style_task and model_config.is_vla:
+ raise RuntimeError(
+ f"training.task={task_type} is incompatible with model.is_vla=true"
+ )
if model_config.is_vla and task_type != "vla_calvin":
_rprint(
"Warning: model.is_vla=True while training.task!=vla_calvin. Enabling VLA path.",
@@ -920,13 +1261,19 @@ def __call__(self, text, **kwargs):
# Branches assign concrete types that don't share a common base.
dataset: Any
collator: Any
+ validation_dataset = None
+ validation_collator = None
+ test_dataset = None
+ test_collator = None
# Read once at the top so both branches (VLA and non-VLA) see the same
# value and the downstream `if use_bucketed_collator:` worker-count
- # check at line ~741 is always defined. VLA uses its own collator and
- # always wants `train_config.data_num_workers`, so we force False there.
+ # check is always defined. VLA and finite AstroPRISM tasks use dedicated
+ # collators, so force False for those tasks.
use_bucketed_collator = (
- os.environ.get("USE_BUCKETED_COLLATOR", "1") == "1" and not is_vla_task
+ os.environ.get("USE_BUCKETED_COLLATOR", "1") == "1"
+ and not is_vla_task
+ and not is_astroprism_map_style_task
)
if is_vla_task:
@@ -1008,6 +1355,55 @@ def __call__(self, text, **kwargs):
flush=True,
)
collator = VLACollator(tokenizer)
+ elif is_astroprism_stage1a_task:
+ _rprint(
+ "Initializing AstroPRISM Stage 1A raw-image dataset...",
+ flush=True,
+ )
+ dataset, collator = _build_astroprism_stage1a_data(
+ cfg.get("data", None),
+ tokenizer,
+ )
+ _rprint(f"Loaded {len(dataset):,} AstroPRISM Stage 1A samples", flush=True)
+ elif is_astroprism_stage1b_spectrum_task:
+ _rprint(
+ "Initializing AstroPRISM Stage 1B raw-spectrum proxy dataset...",
+ flush=True,
+ )
+ dataset, collator = _build_astroprism_stage1b_spectrum_data(
+ cfg.get("data", None),
+ tokenizer,
+ )
+ _rprint(f"Loaded {len(dataset):,} AstroPRISM spectrum samples", flush=True)
+ elif is_astroprism_joint_task:
+ data_cfg = cfg.get("data", None)
+ configured_split = (
+ _coerce_str_or_none(data_cfg.get("split", None))
+ if data_cfg is not None and hasattr(data_cfg, "get")
+ else None
+ )
+ if configured_split != "train":
+ raise RuntimeError(
+ "training.task=astroprism_joint_aion requires data.split=train; "
+ "validation/test are constructed as held-out, evaluation-only loaders"
+ )
+ _rprint("Initializing joint AstroPRISM image+spectrum datasets...", flush=True)
+ dataset, collator = _build_astroprism_joint_data(
+ data_cfg, tokenizer, split="train"
+ )
+ validation_dataset, validation_collator = _build_astroprism_joint_data(
+ data_cfg, tokenizer, split="validation"
+ )
+ test_dataset, test_collator = _build_astroprism_joint_data(
+ data_cfg, tokenizer, split="test"
+ )
+ _rprint(
+ "Loaded AstroPRISM joint splits: "
+ f"train={dataset.modality_lengths}, "
+ f"validation={validation_dataset.modality_lengths}, "
+ f"test={test_dataset.modality_lengths}",
+ flush=True,
+ )
else:
_rprint("Initializing Streaming Datasets...", flush=True)
@@ -1373,19 +1769,60 @@ def __call__(self, text, **kwargs):
f"using default n_workers={n_workers}",
flush=True,
)
- _rprint(
- f"[Rank {rank}] Creating DataLoader: num_workers={n_workers}, pin_memory={use_pin_memory}",
- flush=True,
- )
- train_loader = DataLoader(
- dataset,
- batch_size=train_config.batch_size,
- collate_fn=collator,
- num_workers=n_workers,
- persistent_workers=True if n_workers > 0 else False,
- prefetch_factor=2 if n_workers > 0 else None,
- pin_memory=use_pin_memory,
- )
+ validation_loaders = {}
+ test_loaders = {}
+ if is_astroprism_joint_task:
+ joint_batch_sampler = AstroPRISMJointBatchSampler(
+ dataset,
+ batch_size=train_config.batch_size,
+ seed=seed,
+ num_replicas=world_size,
+ )
+ _rprint(
+ f"[Rank {rank}] Creating rank-synchronous joint DataLoader: "
+ f"num_workers={n_workers}, pin_memory={use_pin_memory}, "
+ f"world_size={world_size}, global_batches={len(joint_batch_sampler)}",
+ flush=True,
+ )
+ train_loader = DataLoader(
+ dataset,
+ batch_sampler=joint_batch_sampler,
+ collate_fn=collator,
+ num_workers=n_workers,
+ persistent_workers=n_workers > 0,
+ prefetch_factor=2 if n_workers > 0 else None,
+ pin_memory=use_pin_memory,
+ )
+ validation_loaders = _build_astroprism_joint_eval_loaders(
+ validation_dataset,
+ validation_collator,
+ batch_size=train_config.batch_size,
+ num_workers=n_workers,
+ pin_memory=use_pin_memory,
+ )
+ test_loaders = _build_astroprism_joint_eval_loaders(
+ test_dataset,
+ test_collator,
+ batch_size=train_config.batch_size,
+ num_workers=n_workers,
+ pin_memory=use_pin_memory,
+ )
+ else:
+ _rprint(
+ f"[Rank {rank}] Creating DataLoader: num_workers={n_workers}, "
+ f"pin_memory={use_pin_memory}, shuffle={is_astroprism_map_style_task}",
+ flush=True,
+ )
+ train_loader = DataLoader(
+ dataset,
+ batch_size=train_config.batch_size,
+ shuffle=is_astroprism_map_style_task,
+ collate_fn=collator,
+ num_workers=n_workers,
+ persistent_workers=n_workers > 0,
+ prefetch_factor=2 if n_workers > 0 else None,
+ pin_memory=use_pin_memory,
+ )
_rprint(f"[Rank {rank}] DataLoader created successfully", flush=True)
# Critical: Attach tokenizer to model for Evaluators/Trainer usage
@@ -1401,7 +1838,7 @@ def __call__(self, text, **kwargs):
logger.info("\n*** LOADING WEIGHTS FROM CHECKPOINT (model-only, fresh optimizer) ***")
logger.info(f" Checkpoint: {train_config.resume_weights_only}")
prev_step = load_model_weights_only(
- model, train_config.resume_weights_only, device=device
+ model, train_config.resume_weights_only, device="cpu"
)
logger.info(f" Previous training ended at step {prev_step}")
logger.info(" Starting fresh training with new optimizer/scheduler config")
@@ -1447,7 +1884,15 @@ def __call__(self, text, **kwargs):
# here too -- ZoneATrainer handles DNA/LoRA/KEGG-eval via
# config-gated branches (lora_enabled, freeze_connector,
# bioreason_dataset), no longer needing a separate trainer class.
- trainer = ZoneATrainer(model, train_config, train_loader)
+ # AstroPRISM's joint path builds per-modality eval loaders above;
+ # every other task leaves these empty and the trainer skips eval.
+ trainer = ZoneATrainer(
+ model,
+ train_config,
+ train_loader,
+ validation_loaders=validation_loaders,
+ test_loaders=test_loaders,
+ )
# 6. Train
_rprint(f"Starting Training Loop for {train_config.max_steps} steps...")
trainer.train()
diff --git a/src/training/trainer_zone_a.py b/src/training/trainer_zone_a.py
index 22ed975..f2bda5b 100644
--- a/src/training/trainer_zone_a.py
+++ b/src/training/trainer_zone_a.py
@@ -1,4 +1,6 @@
+import json
import os
+from collections.abc import Mapping
import torch
import torch.optim as optim
@@ -28,6 +30,174 @@
wandb = None
+def _set_encoder_trainable(encoder: torch.nn.Module, trainable: bool) -> None:
+ """Apply an encoder's narrow trainability boundary when it exposes one.
+
+ AION's raw-image wrapper owns both a frozen fp32 codec and a potentially
+ trainable shared encoder. Falling back to toggling every parameter is only
+ safe for legacy encoders that do not expose that distinction.
+ """
+ set_shared_trainable = getattr(encoder, "set_shared_encoder_trainable", None)
+ if callable(set_shared_trainable):
+ set_shared_trainable(trainable)
+ return
+
+ for parameter in encoder.parameters():
+ parameter.requires_grad_(trainable)
+
+
+def _configure_encoder_for_training(
+ encoder: torch.nn.Module,
+ *,
+ trainable: bool,
+ target_dtype: torch.dtype,
+) -> None:
+ """Set dtype and trainability without recasting protected encoder state."""
+ if not callable(getattr(encoder, "set_shared_encoder_trainable", None)):
+ encoder.to(dtype=target_dtype)
+ _set_encoder_trainable(encoder, trainable)
+
+
+def _refresh_encoder_frozen_cache(model: torch.nn.Module) -> None:
+ """Refresh cached encoder freeze state through the model API when present."""
+ refresh = getattr(model, "refresh_encoder_frozen_cache", None)
+ if callable(refresh):
+ refresh()
+ return
+
+ cache = getattr(model, "_encoder_frozen_cache", None)
+ encoders = getattr(model, "encoders", None)
+ if not isinstance(cache, dict) or encoders is None:
+ return
+ for modality, encoder in encoders.items():
+ cache[modality] = not any(parameter.requires_grad for parameter in encoder.parameters())
+
+
+def _labels_for_batch(batch):
+ """Prefer collator-supplied loss labels; retain legacy text-as-labels."""
+ if "labels" in batch:
+ return batch["labels"]
+ return batch["text"]
+
+
+_ALIGNMENT_MODALITIES = ("image", "spectrum")
+_LOSS_STAT_SIZE = 2 + 2 * len(_ALIGNMENT_MODALITIES)
+
+
+def _supervised_token_count(batch, *, pad_token_id=None) -> torch.Tensor:
+ """Count tokens that contribute to the causal-language-model loss.
+
+ AstroPRISM collators provide assistant-only ``labels`` with ``-100`` for
+ prompt and padding positions. Those labels are the authoritative loss
+ mask. The text fallback keeps legacy Zone A datasets working when they do
+ not provide a separate label tensor.
+ """
+ labels = batch.get("labels")
+ if isinstance(labels, torch.Tensor):
+ return labels.ne(-100).sum(dtype=torch.float32)
+
+ text = batch.get("text")
+ if not isinstance(text, torch.Tensor):
+ return torch.tensor(0.0)
+ if pad_token_id is None:
+ return torch.tensor(float(text.numel()), device=text.device)
+ return text.ne(pad_token_id).sum(dtype=torch.float32)
+
+
+def _alignment_modality(batch) -> str | None:
+ """Return the single tracked alignment modality present in ``batch``."""
+ present = [
+ modality
+ for modality in _ALIGNMENT_MODALITIES
+ if modality in batch and batch[modality] is not None
+ ]
+ return present[0] if len(present) == 1 else None
+
+
+def _validate_provenance_split(batch, expected_split: str) -> None:
+ """Fail closed when a provenance-bearing batch crosses split roles."""
+ provenance = batch.get("provenance")
+ if isinstance(provenance, dict):
+ records = [provenance]
+ elif isinstance(provenance, list | tuple):
+ records = provenance
+ else:
+ return
+ observed = {
+ record.get("split")
+ for record in records
+ if isinstance(record, dict) and record.get("split") is not None
+ }
+ if observed and observed != {expected_split}:
+ raise RuntimeError(
+ f"expected only {expected_split!r} samples, observed splits {sorted(observed)!r}"
+ )
+
+
+def _loss_stats_for_batch(
+ loss: torch.Tensor,
+ batch,
+ *,
+ pad_token_id=None,
+ modality: str | None = None,
+) -> torch.Tensor:
+ """Build additive ``(loss numerator, supervised tokens)`` statistics.
+
+ The first pair is the aggregate. Subsequent pairs are image and spectrum,
+ respectively. A fixed-width vector lets every distributed rank execute
+ the same SUM collective even when ranks see different modalities.
+ """
+ detached_loss = loss.detach().float()
+ token_count = _supervised_token_count(
+ batch,
+ pad_token_id=pad_token_id,
+ ).to(device=detached_loss.device, dtype=torch.float32)
+ stats = torch.zeros(
+ _LOSS_STAT_SIZE,
+ device=detached_loss.device,
+ dtype=torch.float32,
+ )
+ stats[0] = detached_loss * token_count
+ stats[1] = token_count
+
+ resolved_modality = modality or _alignment_modality(batch)
+ if resolved_modality in _ALIGNMENT_MODALITIES:
+ offset = 2 + 2 * _ALIGNMENT_MODALITIES.index(resolved_modality)
+ stats[offset] = stats[0]
+ stats[offset + 1] = token_count
+ return stats
+
+
+def _loss_metrics_from_stats(
+ stats: torch.Tensor,
+ prefix: str,
+ *,
+ include_legacy_loss: bool = False,
+) -> dict[str, float]:
+ """Convert reduced additive statistics to token-weighted loss metrics."""
+ values = stats.detach().float().cpu().tolist()
+ total_tokens = float(values[1])
+ if total_tokens <= 0:
+ return {}
+
+ aggregate_loss = float(values[0]) / total_tokens
+ metrics = {
+ f"{prefix}/loss": aggregate_loss,
+ f"{prefix}/supervised_tokens": total_tokens,
+ }
+ if include_legacy_loss:
+ metrics["loss"] = aggregate_loss
+
+ for index, modality in enumerate(_ALIGNMENT_MODALITIES):
+ offset = 2 + 2 * index
+ modality_tokens = float(values[offset + 1])
+ if modality_tokens <= 0:
+ continue
+ metrics[f"{prefix}/{modality}_loss"] = float(values[offset]) / modality_tokens
+ metrics[f"{prefix}/{modality}_supervised_tokens"] = modality_tokens
+ return metrics
+
+
class _NullEvaluatorRegistry:
@staticmethod
def get(_registry_key):
@@ -86,9 +256,7 @@ def check(self, batch, step):
# Identify present modalities (excluding 'text' which is usually present as labels)
# Note: In PRISM, 'text' is the input for SFT. Multimodal inputs have other keys.
present = [
- k
- for k in batch.keys()
- if k != "text" and k != "labels" and not k.startswith("_")
+ k for k in batch.keys() if k != "text" and k != "labels" and not k.startswith("_")
]
if not present:
@@ -134,8 +302,18 @@ def _unwrap_model(self):
return inner
return self.accelerator.unwrap_model(self.model)
- def __init__(self, model: UnifiedTransformer, config: TrainingConfig, train_loader):
+ def __init__(
+ self,
+ model: UnifiedTransformer,
+ config: TrainingConfig,
+ train_loader,
+ *,
+ validation_loaders: Mapping[str, object] | None = None,
+ test_loaders: Mapping[str, object] | None = None,
+ ):
self.config = config
+ self.validation_loaders = dict(validation_loaders or {})
+ self.test_loaders = dict(test_loaders or {})
self.monitor = None # Init later
# 1. Initialize Accelerator
@@ -178,9 +356,7 @@ def __init__(self, model: UnifiedTransformer, config: TrainingConfig, train_load
from src.training.distributed import _resolve_find_unused
num_modalities = len(getattr(model.config, "modalities", []) or [])
- is_composite = (
- os.environ.get("ZE_FLAT_DEVICE_HIERARCHY", "").upper() == "COMPOSITE"
- )
+ is_composite = os.environ.get("ZE_FLAT_DEVICE_HIERARCHY", "").upper() == "COMPOSITE"
is_interleaved = bool(getattr(model.config, "is_interleaved_qa", False))
# Pass trainable_mb=0 to assume projector-only (the only regime
# this Accelerate trainer is configured for — frozen encoders +
@@ -308,22 +484,29 @@ def __init__(self, model: UnifiedTransformer, config: TrainingConfig, train_load
freeze_vit = getattr(config, "freeze_vit", True)
if freeze_vit:
self.accelerator.print("Freezing Encoders...")
- for param in self.model.encoders.parameters():
- param.requires_grad = False
else:
self.accelerator.print("Unfreezing Encoders...")
- for param in self.model.encoders.parameters():
- param.requires_grad = True
- params_vit.extend([p for p in self.model.encoders.parameters() if p.requires_grad])
-
+ for encoder in self.model.encoders.values():
+ _configure_encoder_for_training(
+ encoder,
+ trainable=not freeze_vit,
+ target_dtype=target_dtype,
+ )
+ _refresh_encoder_frozen_cache(self.model)
+ if not freeze_vit:
+ params_vit.extend(
+ parameter
+ for parameter in self.model.encoders.parameters()
+ if parameter.requires_grad
+ )
- # PATRICK: Is commenting this out breaking?
- #self.accelerator.print("Training Projectors Only (Freezing Non-Projector Params)...")
+ # PATRICK: Is commenting this out breaking?
+ # self.accelerator.print("Training Projectors Only (Freezing Non-Projector Params)...")
# Letting the freeze logic handle parameters allows experimentation with unfreezing encoders/backbones.
-
+
# Zone A is projector warmup: freeze everything except projectors.
- #for name, param in self.model.named_parameters():
+ # for name, param in self.model.named_parameters():
# param.requires_grad = name.startswith("projectors.")
# Cast projectors to target_dtype
@@ -370,7 +553,6 @@ def __init__(self, model: UnifiedTransformer, config: TrainingConfig, train_load
if params_vit:
param_groups.append({"params": params_vit, "lr": lr_vit, "name": "vit"})
-
# Group 2: LLM
if params_llm:
param_groups.append({"params": params_llm, "lr": lr_llm, "name": "llm"})
@@ -466,9 +648,7 @@ def __init__(self, model: UnifiedTransformer, config: TrainingConfig, train_load
if use_deepspeed:
ds_plugin = self.accelerator.state.deepspeed_plugin
ds_plugin.deepspeed_config["train_micro_batch_size_per_gpu"] = config.batch_size
- self.model, self.optimizer = self.accelerator.prepare(
- self.model, self.optimizer
- )
+ self.model, self.optimizer = self.accelerator.prepare(self.model, self.optimizer)
else:
self.model, self.optimizer, self.train_loader = self.accelerator.prepare(
self.model, self.optimizer, self.train_loader
@@ -515,7 +695,12 @@ def __init__(self, model: UnifiedTransformer, config: TrainingConfig, train_load
# We should instantiate them here.
# Ideally, we only run eval on main process to avoid duplication in logging.
self.evaluators = []
- if self.accelerator.is_main_process and _EVAL_IMPORT_ERROR is None:
+ if self.validation_loaders:
+ if self.accelerator.is_main_process:
+ self.accelerator.print(
+ "Skipping registry evaluators: held-out loss loaders are configured."
+ )
+ elif self.accelerator.is_main_process and _EVAL_IMPORT_ERROR is None:
self.accelerator.print("Initializing Evaluators...")
# Unwrapped model for eval generation?
# Accelerator wraps model.
@@ -574,39 +759,67 @@ def add_evaluator(name, registry_key):
f"Skipping evaluator initialization: optional dependencies unavailable ({_EVAL_IMPORT_ERROR})"
)
- # WandB Init handled by Accelerator if 'log_with' is set,
- # or we config it manually. Accelerator tracking is cleaner.
- # Explicit WandB Init (Bypassing Accelerator for reliability with Native DDP)
- if config.wandb_project and self.accelerator.is_main_process:
- if wandb is None:
- self.accelerator.print(
- "[Warning] wandb is not installed. Skipping explicit WandB initialization."
- )
- else:
- self.accelerator.print(
- f"[DEBUG] Explicitly initializing WandB: {config.wandb_project} run={config.wandb_run_name}"
- )
- try:
- wandb.init(
- project=config.wandb_project,
- name=config.wandb_run_name,
- config=config.__dict__,
- reinit=True,
- )
- self.accelerator.print("[DEBUG] WandB Explicit Init Successful!")
- except Exception as e:
- self.accelerator.print(f"[ERROR] WandB Explicit Init Failed: {e}")
- else:
- if self.accelerator.is_main_process:
- self.accelerator.print("[DEBUG] WandB Project not set - Skipping Init.")
+ self._initialize_trackers()
# IsoFLOP FLOP counter. The launcher exports `CALIBRATION_JSON` per
# cell (see `tools/isoflop_launch.py` in PR-2). When unset, the
# counter is a no-op and per-step records get `flops_per_step=None`
# — the collector treats that as missing rather than zero.
- self._flop_counter = _FlopCounter.from_calibration(
- os.environ.get("CALIBRATION_JSON")
- )
+ self._flop_counter = _FlopCounter.from_calibration(os.environ.get("CALIBRATION_JSON"))
+
+ def _initialize_trackers(self):
+ """Initialize Accelerate's W&B tracker exactly once on rank zero."""
+ if not self.accelerator.is_main_process:
+ return
+ if not self.config.wandb_project:
+ self.accelerator.print("[DEBUG] W&B project not set - skipping init.")
+ return
+ if wandb is None:
+ self.accelerator.print(
+ "[Warning] wandb is not installed. Skipping W&B tracker initialization."
+ )
+ return
+
+ wandb_kwargs = {
+ "name": self.config.wandb_run_name,
+ "reinit": True,
+ }
+ if self.config.wandb_entity:
+ wandb_kwargs["entity"] = self.config.wandb_entity
+ if self.config.wandb_mode:
+ wandb_kwargs["mode"] = self.config.wandb_mode
+ if self.config.wandb_run_id:
+ wandb_kwargs["id"] = self.config.wandb_run_id
+ wandb_kwargs["resume"] = "allow"
+
+ try:
+ self.accelerator.init_trackers(
+ project_name=self.config.wandb_project,
+ config=self.config.__dict__,
+ init_kwargs={"wandb": wandb_kwargs},
+ )
+ self.accelerator.print("[DEBUG] Accelerate W&B tracker initialized.")
+ except Exception as exc:
+ self.accelerator.print(f"[ERROR] W&B tracker initialization failed: {exc}")
+
+ def _write_loss_metrics(self, metrics, *, step, split):
+ """Append rank-zero loss metrics for offline inspection."""
+ if not self.accelerator.is_main_process or not metrics:
+ return
+ output_dir = getattr(self.config, "output_dir", None)
+ if not output_dir:
+ return
+
+ try:
+ os.makedirs(output_dir, exist_ok=True)
+ record = {"step": step, "split": split}
+ for key, value in metrics.items():
+ record[key] = float(value)
+ metrics_path = os.path.join(output_dir, "loss_metrics.jsonl")
+ with open(metrics_path, "a", encoding="utf-8") as handle:
+ handle.write(json.dumps(record, sort_keys=True) + "\n")
+ except Exception as exc: # noqa: BLE001
+ self.accelerator.print(f"[Warning] Failed to append loss metrics for {split}: {exc}")
# KEGG eval (run_kegg_eval): dataset/tokenizer are lazy-loaded on first call.
self._kegg_eval_dataset = None
@@ -864,9 +1077,7 @@ def train(self):
},
)
except Exception as _e_startup: # noqa: BLE001
- self.accelerator.print(
- f"[perf] startup_param_count skipped: {_e_startup}"
- )
+ self.accelerator.print(f"[perf] startup_param_count skipped: {_e_startup}")
data_iter = iter(self.train_loader)
@@ -919,6 +1130,11 @@ def train(self):
_pad_id_for_tokens = (
getattr(_tok_for_pad, "pad_token_id", None) if _tok_for_pad is not None else None
)
+ train_loss_window = torch.zeros(
+ _LOSS_STAT_SIZE,
+ device=self._device,
+ dtype=torch.float32,
+ )
self.optimizer.zero_grad()
@@ -943,6 +1159,9 @@ def train(self):
self.accelerator.print("Error: 'text' missing in batch for Zone A training.")
break
+ if self.validation_loaders or self.test_loaders:
+ _validate_provenance_split(batch, "train")
+
# Rank-0-only counts; scaled by world_size in the perf record below.
if self.accelerator.is_main_process:
try:
@@ -955,20 +1174,12 @@ def train(self):
text_tensor = batch.get("text")
if isinstance(text_tensor, torch.Tensor) and text_tensor.dim() == 2:
if _pad_id_for_tokens is None:
- per_sample = [int(text_tensor.shape[1])] * int(
- text_tensor.shape[0]
- )
+ per_sample = [int(text_tensor.shape[1])] * int(text_tensor.shape[0])
else:
- per_sample = (
- (text_tensor != _pad_id_for_tokens)
- .sum(dim=1)
- .tolist()
- )
+ per_sample = (text_tensor != _pad_id_for_tokens).sum(dim=1).tolist()
seq_lens_window.extend(int(n) for n in per_sample)
except Exception as _e_perf: # noqa: BLE001
- self.accelerator.print(
- f"[perf] modality/token count skipped: {_e_perf}"
- )
+ self.accelerator.print(f"[perf] modality/token count skipped: {_e_perf}")
# Monitor Data Starvation
self.monitor.check(batch, step)
@@ -987,7 +1198,7 @@ def train(self):
# --- Forward Timing ---
t1 = time.perf_counter()
with self.accelerator.autocast():
- logits, loss = self.model(batch, labels=batch["text"])
+ logits, loss = self.model(batch, labels=_labels_for_batch(batch))
# Sync for accurate timing on XPU (skipped under PRISM_PRODUCTION_MODE)
if not _production_mode and hasattr(torch, "xpu") and torch.xpu.is_available():
torch.xpu.synchronize()
@@ -998,9 +1209,7 @@ def train(self):
self.accelerator.print(f"\nCRITICAL: NaN Loss Detected at Step {step}")
if "_metadata" in batch:
# _metadata is a list of strings (collated)
- self.accelerator.print(
- f"Offending Batch Metadata: {batch['_metadata']}"
- )
+ self.accelerator.print(f"Offending Batch Metadata: {batch['_metadata']}")
else:
self.accelerator.print(
"No metadata found in batch (Was it stripped by collator?)."
@@ -1012,6 +1221,12 @@ def train(self):
# self.accelerator.print(f"Batch Keys: {batch.keys()}")
raise ValueError("Training Diverged (NaN) - See Log for Metadata")
+ train_loss_window += _loss_stats_for_batch(
+ loss,
+ batch,
+ pad_token_id=_pad_id_for_tokens,
+ )
+
# --- Backward Timing (separating compute from AllReduce) ---
# DEBUG_NO_SYNC: Run backward WITHOUT DDP AllReduce to isolate pure compute time
debug_no_sync = os.environ.get("DEBUG_NO_SYNC", "0") == "1"
@@ -1049,7 +1264,11 @@ def train(self):
# Now wait for AllReduce to complete (DDP overlaps but we measure the wait).
# Skipped under PRISM_PRODUCTION_MODE.
- if not _production_mode and hasattr(torch.distributed, "barrier") and torch.distributed.is_initialized():
+ if (
+ not _production_mode
+ and hasattr(torch.distributed, "barrier")
+ and torch.distributed.is_initialized()
+ ):
torch.distributed.barrier()
t_after_allreduce = time.perf_counter()
timing_allreduce_wait += t_after_allreduce - t_after_backward
@@ -1071,15 +1290,27 @@ def train(self):
# Separate Clipping
# 1. Projectors (Connectors)
if unwrapped_model.projectors:
- self.accelerator.clip_grad_norm_(unwrapped_model.projectors.parameters(), 1.0)
+ self.accelerator.clip_grad_norm_(
+ unwrapped_model.projectors.parameters(), 1.0
+ )
# 2. Encoders (ViT) - Only if trainable
- if not getattr(self.config, "freeze_vit", True) and unwrapped_model.encoders:
- self.accelerator.clip_grad_norm_(unwrapped_model.encoders.parameters(), 1.0)
+ if (
+ not getattr(self.config, "freeze_vit", True)
+ and unwrapped_model.encoders
+ ):
+ self.accelerator.clip_grad_norm_(
+ unwrapped_model.encoders.parameters(), 1.0
+ )
# 3. Backbone (LLM) - Only if trainable
- if not getattr(self.config, "freeze_llm", True) and unwrapped_model.backbone:
- self.accelerator.clip_grad_norm_(unwrapped_model.backbone.parameters(), 1.0)
+ if (
+ not getattr(self.config, "freeze_llm", True)
+ and unwrapped_model.backbone
+ ):
+ self.accelerator.clip_grad_norm_(
+ unwrapped_model.backbone.parameters(), 1.0
+ )
else:
# Legacy Global Clipping
max_grad_norm = getattr(self.config, "max_grad_norm", 1.0)
@@ -1180,9 +1411,7 @@ def train(self):
else 0.0
)
tokens_per_batch = (
- global_tokens_window / timing_steps
- if timing_steps > 0
- else 0.0
+ global_tokens_window / timing_steps if timing_steps > 0 else 0.0
)
# IsoFLOP per-step extensions: sequence-length distribution
# + analytic FLOP accounting + projector capacity knobs.
@@ -1290,12 +1519,30 @@ def train(self):
# 1. Loss Logging (0 disables — guard against ZeroDivisionError).
log_interval = self.config.log_every_n_steps
+ should_log_loss = False
if log_interval and step % log_interval == 0:
- avg_loss = loss.item()
+ should_log_loss = True
+ elif log_interval and (step + 1) >= self.config.max_steps:
+ should_log_loss = True
+ if should_log_loss:
+ global_loss_stats = self.accelerator.reduce(
+ train_loss_window,
+ reduction="sum",
+ )
+ loss_metrics = _loss_metrics_from_stats(
+ global_loss_stats,
+ "train",
+ include_legacy_loss=True,
+ )
+ if not loss_metrics:
+ fallback_loss = float(loss.detach().item())
+ loss_metrics = {"loss": fallback_loss, "train/loss": fallback_loss}
+ avg_loss = loss_metrics["loss"]
lr = self.scheduler.get_last_lr()[0]
- log_data.update(
- {"loss": avg_loss, "lr": lr}
- ) # "step": step is handled by accelerator kwarg
+ log_data.update(loss_metrics)
+ log_data["lr"] = lr
+ self._write_loss_metrics(loss_metrics, step=step, split="train")
+ train_loss_window.zero_()
do_log = True
if self.accelerator.is_main_process:
@@ -1356,30 +1603,118 @@ def train(self):
# Save Final Checkpoint
self.save_checkpoint("final")
+ if self.test_loaders:
+ self.run_test_evaluation(step)
self.accelerator.print("Training Complete.")
self.accelerator.print(f"Checkpoints saved to: {self.config.output_dir}")
self.accelerator.end_training()
- def run_evaluation(self, step):
- """Run all registered evaluators and log metrics."""
- if not self.accelerator.is_main_process:
- return
+ def _evaluate_loss_loaders(self, loaders, *, split, max_batches=None):
+ """Evaluate deterministic replicated loaders without gradient updates.
- self.accelerator.print(f"\n[Step {step}] Running Evaluation...")
+ The loaders intentionally remain unprepared by Accelerate. Every rank
+ therefore evaluates the same examples in the same order, which keeps
+ ZeRO/FSDP forwards collective-safe without padding the 65-row held-out
+ splits. Reduced statistics are divided by world size so token counts
+ describe the dataset once rather than the replicated compute.
+ """
+ if not loaders:
+ return {}
+
+ was_training = self.model.training
+ local_stats = torch.zeros(
+ _LOSS_STAT_SIZE,
+ device=self._device,
+ dtype=torch.float32,
+ )
self.model.eval()
+ try:
+ with torch.inference_mode():
+ for modality in sorted(loaders):
+ if modality not in _ALIGNMENT_MODALITIES:
+ raise ValueError(f"unsupported held-out loss modality: {modality!r}")
+ for batch_index, batch in enumerate(loaders[modality]):
+ if max_batches is not None and batch_index >= max_batches:
+ break
+ batch = send_to_device(batch, self._device)
+ _validate_provenance_split(batch, split)
+ observed_modality = _alignment_modality(batch)
+ if observed_modality != modality:
+ raise RuntimeError(
+ "held-out loader/batch modality mismatch: "
+ f"loader={modality!r}, batch={observed_modality!r}"
+ )
+ with self.accelerator.autocast():
+ _, loss = self.model(
+ batch,
+ labels=_labels_for_batch(batch),
+ )
+ if loss is None:
+ raise RuntimeError(f"{split} {modality} evaluation returned no loss")
+ local_stats += _loss_stats_for_batch(
+ loss,
+ batch,
+ modality=modality,
+ )
+ finally:
+ self.model.train(was_training)
+
+ reduced_stats = self.accelerator.reduce(local_stats, reduction="sum")
+ world_size = max(1, int(getattr(self.accelerator, "num_processes", 1)))
+ return _loss_metrics_from_stats(reduced_stats / world_size, split)
+
+ def run_test_evaluation(self, step):
+ """Run the full held-out test split once at the end of training."""
+ metrics = self._evaluate_loss_loaders(
+ self.test_loaders,
+ split="test",
+ max_batches=None,
+ )
+ if self.accelerator.is_main_process and metrics:
+ self.accelerator.log(metrics, step=step)
+ self._write_loss_metrics(metrics, step=step, split="test")
+ self.accelerator.print(f"Test loss metrics: {metrics}")
+ return metrics
- total_metrics = {}
- for evaluator in self.evaluators:
- try:
- name = evaluator.__class__.__name__
- metrics = evaluator.evaluate(limit=10) # Small batch for sanity
- # Prefix metrics
- for k, v in metrics.items():
- total_metrics[f"eval/{name}/{k}"] = v
- except Exception as e:
- self.accelerator.print(f"Eval Error ({name}): {e}")
+ def run_evaluation(self, step):
+ """Run held-out CE loaders, or the legacy evaluator suite."""
+ self.accelerator.print(f"\n[Step {step}] Running Evaluation...")
- self.accelerator.log(total_metrics, step=step)
+ if self.validation_loaders:
+ configured_batches = int(getattr(self.config, "validation_batches", 0) or 0)
+ max_batches = configured_batches if configured_batches > 0 else None
+ total_metrics = self._evaluate_loss_loaders(
+ self.validation_loaders,
+ split="validation",
+ max_batches=max_batches,
+ )
+ if not self.accelerator.is_main_process:
+ return total_metrics
+ self._write_loss_metrics(
+ total_metrics,
+ step=step,
+ split="validation",
+ )
+ else:
+ if not self.accelerator.is_main_process:
+ return {}
+ was_training = self.model.training
+ total_metrics = {}
+ self.model.eval()
+ try:
+ for evaluator in self.evaluators:
+ name = evaluator.__class__.__name__
+ try:
+ metrics = evaluator.evaluate(limit=10)
+ for key, value in metrics.items():
+ total_metrics[f"eval/{name}/{key}"] = value
+ except Exception as exc:
+ self.accelerator.print(f"Eval Error ({name}): {exc}")
+ finally:
+ self.model.train(was_training)
+
+ if total_metrics:
+ self.accelerator.log(total_metrics, step=step)
self.accelerator.print(f"Eval Metrics: {total_metrics}")
# IsoFLOP per-family eval-loss records. The fitter keys off
@@ -1388,9 +1723,7 @@ def run_evaluation(self, step):
# evaluator reports a per-family loss, the collector treats the
# row's `loss_` as missing.
try:
- modalities = (
- model_modalities(self._unwrap_model()) or []
- )
+ modalities = model_modalities(self._unwrap_model()) or []
for family in modalities:
# Pick the first metric whose lowercased key contains the
# modality name AND the substring "loss" — accommodates
@@ -1428,9 +1761,9 @@ def run_evaluation(self, step):
},
)
except Exception as _e_eval_perf: # noqa: BLE001
- self.accelerator.print(
- f"[perf] eval per-family loss record skipped: {_e_eval_perf}"
- )
+ self.accelerator.print(f"[perf] eval per-family loss record skipped: {_e_eval_perf}")
+
+ return total_metrics
def run_kegg_eval(self, step):
"""
diff --git a/tests/multimodal/test_make_dummy_batch.py b/tests/multimodal/test_make_dummy_batch.py
index fee1640..c67c8ad 100644
--- a/tests/multimodal/test_make_dummy_batch.py
+++ b/tests/multimodal/test_make_dummy_batch.py
@@ -10,6 +10,15 @@ def test_image_dummy_shape():
assert t.shape == (3, 224, 224)
+def test_spectrum_dummy_has_raw_desi_fields():
+ spectrum = make_dummy_batch(Modality.SPECTRUM)
+ assert isinstance(spectrum, dict)
+ assert set(spectrum) == {"flux", "ivar", "mask", "wavelength"}
+ assert all(tensor.shape == (7781,) for tensor in spectrum.values())
+ assert spectrum["mask"].dtype == torch.bool
+ assert torch.all(torch.diff(spectrum["wavelength"]) > 0)
+
+
def test_table_dummy_is_long():
t = make_dummy_batch(Modality.TABLE)
assert t.dtype == torch.long
diff --git a/tests/multimodal/test_modality_enum.py b/tests/multimodal/test_modality_enum.py
index b4605ee..3f37bea 100644
--- a/tests/multimodal/test_modality_enum.py
+++ b/tests/multimodal/test_modality_enum.py
@@ -6,6 +6,7 @@
def test_modality_str_equality():
"""Modality.X == 'x' so existing string-keyed dicts and configs work."""
assert Modality.IMAGE == "image"
+ assert Modality.SPECTRUM == "spectrum"
assert Modality.TEXT == "text"
assert Modality.TIME_SERIES == "time_series"
assert "image" == Modality.IMAGE
@@ -28,11 +29,12 @@ def test_dict_lookup_round_trip():
assert d2["graph"] == "graph_enc"
-def test_all_modalities_has_seven():
- assert len(ALL_MODALITIES) == 7
+def test_all_modalities_has_eight():
+ assert len(ALL_MODALITIES) == 8
assert set(ALL_MODALITIES) == {
Modality.TEXT,
Modality.IMAGE,
+ Modality.SPECTRUM,
Modality.TABLE,
Modality.TIME_SERIES,
Modality.GEOMETRY,
@@ -43,6 +45,7 @@ def test_all_modalities_has_seven():
def test_parse_modality_coerces_string():
assert parse_modality("image") is Modality.IMAGE
+ assert parse_modality("spectrum") is Modality.SPECTRUM
assert parse_modality(Modality.GRAPH) is Modality.GRAPH
diff --git a/tests/test_accept_astroprism_stage1a.py b/tests/test_accept_astroprism_stage1a.py
new file mode 100644
index 0000000..266e184
--- /dev/null
+++ b/tests/test_accept_astroprism_stage1a.py
@@ -0,0 +1,179 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import pytest
+from tools.accept_astroprism_stage1a import (
+ AcceptanceConfig,
+ AcceptanceValidationError,
+ accept_stage1a,
+ sha256_file,
+)
+
+pytestmark = [pytest.mark.unit, pytest.mark.multimodal]
+
+
+def _write_source(tmp_path: Path) -> tuple[Path, list[dict]]:
+ source = tmp_path / "source"
+ source.mkdir()
+ records = []
+ for index in range(646):
+ ra = index * 0.1
+ if index == 536:
+ ra = 535 * 0.1 + 1.0 / 3600.0
+ records.append(
+ {
+ "schema_version": 1,
+ "candidate_index": index,
+ "paired_targetid": str(10_000 + index),
+ "caption_object_id": f"caption-{index}",
+ "coordinates_deg": {
+ "paired_ra": ra,
+ "paired_dec": 0.0,
+ "caption_ra": ra,
+ "caption_dec": 0.0,
+ },
+ "caption": {
+ "original": f" Original caption {index} — preserved.\n",
+ "summary_original": f"Summary {index}",
+ "status": "unreviewed",
+ },
+ "image_embedding": {
+ "file": "image_embeddings.f16.npy",
+ "index": index,
+ "dtype": "float16",
+ "shape": [576, 768],
+ },
+ "acceptance_status": "pending_review",
+ "source_group_id": None,
+ "split": None,
+ }
+ )
+ candidates = source / "candidates.jsonl"
+ candidates.write_text(
+ "".join(
+ json.dumps(record, ensure_ascii=False, sort_keys=True, allow_nan=False) + "\n"
+ for record in records
+ ),
+ encoding="utf-8",
+ )
+ embeddings = source / "image_embeddings.f16.npy"
+ embeddings.write_bytes(b"parity-reference")
+ manifest = {
+ "schema_version": 1,
+ "artifact_type": "astroprism_stage1a_candidates",
+ "pins": {"roster_sha256": "roster-digest"},
+ "roster": {"row_count": 646, "sha256": "roster-digest"},
+ "outputs": {
+ "candidates": {
+ "path": candidates.name,
+ "row_count": 646,
+ "size_bytes": candidates.stat().st_size,
+ "sha256": sha256_file(candidates),
+ },
+ "image_embeddings": {
+ "path": embeddings.name,
+ "dtype": "float16",
+ "shape": [646, 576, 768],
+ "size_bytes": embeddings.stat().st_size,
+ "sha256": sha256_file(embeddings),
+ },
+ },
+ "validation": {"acceptance_status": "pending_review"},
+ }
+ manifest_path = source / "manifest.json"
+ manifest_path.write_text(json.dumps(manifest), encoding="utf-8")
+ return manifest_path, records
+
+
+def _read_jsonl(path: Path) -> list[dict]:
+ return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines()]
+
+
+def test_promotes_deterministically_with_exact_grouped_splits(tmp_path: Path) -> None:
+ source_manifest, originals = _write_source(tmp_path)
+ source_candidates = source_manifest.parent / "candidates.jsonl"
+ original_source_digest = sha256_file(source_candidates)
+
+ first = tmp_path / "accepted-a"
+ second = tmp_path / "accepted-b"
+ first_summary = accept_stage1a(AcceptanceConfig(source_manifest, first))
+ second_summary = accept_stage1a(AcceptanceConfig(source_manifest, second))
+
+ assert (first / "candidates.jsonl").read_bytes() == (second / "candidates.jsonl").read_bytes()
+ assert sha256_file(source_candidates) == original_source_digest
+ assert first_summary["candidate_count"] == 646
+ assert first_summary["group_count"] == 645
+ assert first_summary["split_counts"] == {"test": 65, "train": 516, "validation": 65}
+ assert first_summary["mapping_sha256"] == second_summary["mapping_sha256"]
+
+ accepted = _read_jsonl(first / "candidates.jsonl")
+ assert [record["candidate_index"] for record in accepted] == list(range(646))
+ assert [record["caption"]["original"] for record in accepted] == [
+ record["caption"]["original"] for record in originals
+ ]
+ assert all(record["acceptance_status"] == "accepted" for record in accepted)
+ assert all(record["caption"]["status"] == "released_reused" for record in accepted)
+ assert accepted[535]["source_group_id"] == accepted[536]["source_group_id"]
+ assert accepted[535]["split"] == accepted[536]["split"]
+ assert all(
+ record["image_embedding"]["file"] == "../source/image_embeddings.f16.npy"
+ for record in accepted
+ )
+
+ group_splits: dict[str, set[str]] = {}
+ for record in accepted:
+ group_splits.setdefault(record["source_group_id"], set()).add(record["split"])
+ assert all(len(splits) == 1 for splits in group_splits.values())
+
+ manifest = json.loads((first / "manifest.json").read_text(encoding="utf-8"))
+ output = manifest["outputs"]["candidates"]
+ assert output["sha256"] == sha256_file(first / "candidates.jsonl")
+ assert output["size_bytes"] == (first / "candidates.jsonl").stat().st_size
+ assert manifest["splits"]["grouping_policy"] == "skycc60-v1"
+ assert manifest["splits"]["overlap_radius_arcsec"] == 60.0
+ assert manifest["splits"]["row_counts"] == {
+ "test": 65,
+ "train": 516,
+ "validation": 65,
+ }
+ assert manifest["lineage"]["source_candidates"]["sha256"] == original_source_digest
+ assert "validation" not in manifest
+
+
+def test_refuses_source_hash_mismatch_without_writing_output(tmp_path: Path) -> None:
+ source_manifest, _records = _write_source(tmp_path)
+ manifest = json.loads(source_manifest.read_text(encoding="utf-8"))
+ manifest["outputs"]["candidates"]["sha256"] = "0" * 64
+ source_manifest.write_text(json.dumps(manifest), encoding="utf-8")
+ output = tmp_path / "accepted"
+
+ with pytest.raises(AcceptanceValidationError, match="SHA-256"):
+ accept_stage1a(AcceptanceConfig(source_manifest, output))
+
+ assert not output.exists()
+ assert not list(tmp_path.glob(".accepted.tmp-*"))
+
+
+def test_refuses_to_overwrite_or_repromote(tmp_path: Path) -> None:
+ source_manifest, _records = _write_source(tmp_path)
+ output = tmp_path / "accepted"
+ accept_stage1a(AcceptanceConfig(source_manifest, output))
+
+ with pytest.raises(FileExistsError, match="already exists"):
+ accept_stage1a(AcceptanceConfig(source_manifest, output))
+
+ source_candidates = source_manifest.parent / "candidates.jsonl"
+ records = _read_jsonl(source_candidates)
+ records[0]["acceptance_status"] = "accepted"
+ source_candidates.write_text(
+ "".join(json.dumps(record, sort_keys=True) + "\n" for record in records),
+ encoding="utf-8",
+ )
+ manifest = json.loads(source_manifest.read_text(encoding="utf-8"))
+ manifest["outputs"]["candidates"]["sha256"] = sha256_file(source_candidates)
+ source_manifest.write_text(json.dumps(manifest), encoding="utf-8")
+
+ with pytest.raises(AcceptanceValidationError, match="not pending_review"):
+ accept_stage1a(AcceptanceConfig(source_manifest, tmp_path / "second"))
diff --git a/tests/test_aion_image_encoder.py b/tests/test_aion_image_encoder.py
new file mode 100644
index 0000000..c1f2a40
--- /dev/null
+++ b/tests/test_aion_image_encoder.py
@@ -0,0 +1,246 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import pytest
+import src.encoders.aion_image as aion_image_module
+import torch
+import torch.nn as nn
+from src.encoders.aion_image import (
+ AION_IMAGE_CODEC_REVISION,
+ AION_IMAGE_TOKENS,
+ AION_IMAGE_WIDTH,
+ AION_LEGACY_SURVEY_BANDS,
+ AION_MODEL_REVISION,
+ AIONImageEncoder,
+)
+
+pytestmark = pytest.mark.unit
+
+
+@dataclass
+class _FakeLegacySurveyImage:
+ flux: torch.Tensor
+ bands: list[str]
+
+ token_key = "tok_image"
+
+
+class _FakeImageCodec(nn.Module):
+ from_pretrained_calls: list[tuple[tuple, dict]] = []
+
+ def __init__(self) -> None:
+ super().__init__()
+ self.scale = nn.Parameter(torch.ones((), dtype=torch.float64))
+ self.last_image: _FakeLegacySurveyImage | None = None
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ cls.from_pretrained_calls.append((args, kwargs))
+ return cls()
+
+ def encode(self, image: _FakeLegacySurveyImage) -> torch.Tensor:
+ self.last_image = image
+ batch_size = image.flux.shape[0]
+ # Deliberately mimic AION's in-place preprocessing to prove the wrapper
+ # cloned the caller's tensor.
+ image.flux.zero_()
+ return torch.zeros(
+ batch_size, AION_IMAGE_TOKENS, dtype=torch.long, device=image.flux.device
+ )
+
+
+class _FakeAION(nn.Module):
+ from_pretrained_calls: list[tuple[tuple, dict]] = []
+
+ def __init__(self) -> None:
+ super().__init__()
+ self.encoder_embeddings = nn.ModuleDict({"tok_image": nn.Embedding(8, 1)})
+ self.encoder = nn.ModuleList([nn.Linear(1, 1, bias=False)])
+ self.encoder_norm = nn.LayerNorm(1)
+ self.decoder_proj_context = nn.Linear(1, 1, bias=False)
+ self.unused_decoder = nn.Linear(1, 1)
+ self.encode_calls: list[tuple[dict[str, torch.Tensor], int]] = []
+ self.output_shape: tuple[int, int] = (AION_IMAGE_TOKENS, AION_IMAGE_WIDTH)
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ cls.from_pretrained_calls.append((args, kwargs))
+ return cls()
+
+ def encode(self, inputs: dict[str, torch.Tensor], *, num_encoder_tokens: int):
+ self.encode_calls.append((inputs, num_encoder_tokens))
+ batch_size = next(iter(inputs.values())).shape[0]
+ token_count, width = self.output_shape
+ # Keep the output connected to the shared encoder for the unfreeze test.
+ scale = self.encoder[0].weight.reshape(1, 1, 1)
+ return scale * torch.ones(
+ batch_size,
+ token_count,
+ width,
+ dtype=scale.dtype,
+ device=scale.device,
+ )
+
+
+@pytest.fixture(autouse=True)
+def _mock_aion_types(monkeypatch):
+ _FakeAION.from_pretrained_calls.clear()
+ _FakeImageCodec.from_pretrained_calls.clear()
+ monkeypatch.setattr(
+ aion_image_module,
+ "_load_aion_types",
+ lambda: (_FakeAION, _FakeImageCodec, _FakeLegacySurveyImage),
+ )
+
+
+def _make_encoder(**kwargs) -> AIONImageEncoder:
+ return AIONImageEncoder(model_name_or_path="local/aion-base", **kwargs)
+
+
+def test_loads_model_and_codec_with_explicit_pinned_offline_controls(tmp_path: Path):
+ encoder = _make_encoder(cache_dir=tmp_path)
+
+ assert _FakeAION.from_pretrained_calls == [
+ (
+ ("local/aion-base",),
+ {
+ "revision": AION_MODEL_REVISION,
+ "cache_dir": str(tmp_path),
+ "local_files_only": True,
+ },
+ )
+ ]
+ assert _FakeImageCodec.from_pretrained_calls == [
+ (
+ ("local/aion-base",),
+ {
+ "modality": _FakeLegacySurveyImage,
+ "revision": AION_IMAGE_CODEC_REVISION,
+ "cache_dir": str(tmp_path),
+ "local_files_only": True,
+ },
+ )
+ ]
+ assert encoder.tokens_per_instance() == AION_IMAGE_TOKENS
+ assert encoder.output_dim == AION_IMAGE_WIDTH
+
+
+def test_raw_image_is_cloned_and_encoded_as_exactly_576_aion_tokens():
+ encoder = _make_encoder()
+ raw = torch.randn(2, 4, 160, 160)
+ original = raw.clone()
+
+ features = encoder(raw)
+
+ assert features.shape == (2, AION_IMAGE_TOKENS, AION_IMAGE_WIDTH)
+ assert torch.equal(raw, original)
+ assert encoder.image_codec.last_image is not None
+ assert encoder.image_codec.last_image.bands == list(AION_LEGACY_SURVEY_BANDS)
+ assert torch.count_nonzero(encoder.image_codec.last_image.flux) == 0
+
+ [(encoded_inputs, token_budget)] = encoder.aion_model.encode_calls
+ assert set(encoded_inputs) == {"tok_image"}
+ assert encoded_inputs["tok_image"].shape == (2, AION_IMAGE_TOKENS)
+ assert token_budget == AION_IMAGE_TOKENS
+
+
+def test_codec_stays_float32_frozen_and_eval_when_encoder_is_unfrozen():
+ encoder = _make_encoder().to(dtype=torch.bfloat16)
+ encoder.set_shared_encoder_trainable(True)
+ encoder.train()
+
+ codec_parameters = list(encoder.image_codec.parameters())
+ assert codec_parameters
+ assert all(parameter.dtype == torch.float32 for parameter in codec_parameters)
+ assert all(not parameter.requires_grad for parameter in codec_parameters)
+ assert encoder.image_codec.training is False
+
+ shared_modules = encoder._shared_encoder_modules()
+ assert all(
+ any(parameter.requires_grad for parameter in module.parameters())
+ for module in shared_modules
+ )
+ assert all(module.training for module in shared_modules)
+ assert all(
+ not parameter.requires_grad for parameter in encoder.aion_model.unused_decoder.parameters()
+ )
+
+ features = encoder(torch.randn(1, 4, 160, 160))
+ features.sum().backward()
+ assert encoder.aion_model.encoder[0].weight.grad is not None
+ assert all(parameter.grad is None for parameter in codec_parameters)
+
+ encoder.set_shared_encoder_trainable(False)
+ encoder.train()
+ assert all(not parameter.requires_grad for parameter in encoder.aion_model.parameters())
+ assert all(not module.training for module in shared_modules)
+ assert encoder.image_codec.training is False
+
+
+@pytest.mark.parametrize(
+ "bad_input,exception,message",
+ [
+ (torch.zeros(4, 160, 160), ValueError, "shape"),
+ (torch.zeros(1, 3, 160, 160), ValueError, "shape"),
+ (torch.zeros(1, 4, 96, 96), ValueError, "shape"),
+ (torch.zeros(1, 4, 160, 160, dtype=torch.long), TypeError, "floating point"),
+ (
+ torch.full((1, 4, 160, 160), float("nan")),
+ ValueError,
+ "non-finite",
+ ),
+ ],
+)
+def test_rejects_invalid_raw_image_inputs(bad_input, exception, message):
+ encoder = _make_encoder()
+ with pytest.raises(exception, match=message):
+ encoder(bad_input)
+
+
+def test_rejects_unverified_band_order():
+ with pytest.raises(ValueError, match="g/r/i/z"):
+ _make_encoder(bands=("DES-R", "DES-G", "DES-I", "DES-Z"))
+
+
+def test_rejects_unexpected_codec_token_shape():
+ encoder = _make_encoder()
+
+ def wrong_shape(_image):
+ return torch.zeros(1, AION_IMAGE_TOKENS - 1, dtype=torch.long)
+
+ encoder.image_codec.encode = wrong_shape
+ with pytest.raises(RuntimeError, match="ImageCodec must emit token ids"):
+ encoder(torch.randn(1, 4, 160, 160))
+
+
+def test_rejects_unexpected_shared_encoder_feature_shape():
+ encoder = _make_encoder()
+ encoder.aion_model.output_shape = (AION_IMAGE_TOKENS, AION_IMAGE_WIDTH - 1)
+
+ with pytest.raises(RuntimeError, match="shared encoder must emit features"):
+ encoder(torch.randn(1, 4, 160, 160))
+
+
+def test_codec_values_survive_a_blanket_reduced_precision_cast():
+ """Dtype alone is not enough: rounded weights can change token ids.
+
+ ZoneATrainer casts the whole encoder ModuleDict to the backbone dtype
+ (src/training/trainer_zone_a.py), which reaches this encoder through
+ nn.Module._apply. Restoring float32 afterwards would recover the dtype but
+ not the values, and the codec quantizes hard.
+ """
+ encoder = _make_encoder()
+ before = [p.detach().clone() for p in encoder.image_codec.parameters()]
+ assert before
+
+ import torch.nn as nn
+
+ bundle = nn.ModuleDict({"aion": encoder})
+ bundle.to(dtype=torch.bfloat16)
+
+ after = list(encoder.image_codec.parameters())
+ assert all(p.dtype == torch.float32 for p in after)
+ for original, current in zip(before, after, strict=True):
+ assert torch.equal(original, current)
diff --git a/tests/test_aion_model_wiring.py b/tests/test_aion_model_wiring.py
new file mode 100644
index 0000000..26aa1b3
--- /dev/null
+++ b/tests/test_aion_model_wiring.py
@@ -0,0 +1,174 @@
+from pathlib import Path
+from unittest.mock import patch
+
+import pytest
+import torch
+import torch.nn as nn
+import yaml
+from src.config import ModelConfig
+from src.encoders.aion_image import (
+ AION_IMAGE_CODEC_REVISION,
+ AION_IMAGE_TOKENS,
+ AION_IMAGE_WIDTH,
+ AION_LEGACY_SURVEY_BANDS,
+ AION_MODEL_REVISION,
+ AION_REPO_ID,
+)
+from src.model import UnifiedTransformer
+
+pytestmark = pytest.mark.unit
+
+
+class _FakeAIONImageEncoder(nn.Module):
+ output_dim = AION_IMAGE_WIDTH
+
+ def __init__(self, *, token_count: int = AION_IMAGE_TOKENS) -> None:
+ super().__init__()
+ self.shared_weight = nn.Parameter(torch.ones(1))
+ self.codec_weight = nn.Parameter(torch.ones(1), requires_grad=False)
+ self.token_count = token_count
+ self.trainability_calls: list[bool] = []
+
+ def tokens_per_instance(self) -> int:
+ return self.token_count
+
+ def set_shared_encoder_trainable(self, trainable: bool):
+ self.trainability_calls.append(bool(trainable))
+ self.shared_weight.requires_grad_(bool(trainable))
+ self.codec_weight.requires_grad_(False)
+ return self
+
+ def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+ return self.shared_weight * torch.ones(
+ inputs.shape[0],
+ self.token_count,
+ self.output_dim,
+ device=inputs.device,
+ )
+
+
+def _aion_config(*, freeze_encoders: bool = True, **kwargs) -> ModelConfig:
+ values = {
+ "d_model": 32,
+ "vocab_size": 64,
+ "num_layers": 1,
+ "num_heads": 4,
+ "num_experts": 2,
+ "modalities": ["image"],
+ "d_img": AION_IMAGE_WIDTH,
+ "image_encoder_type": "aion",
+ "freeze_encoders": freeze_encoders,
+ "aion_model_name_or_path": "local/aion-base",
+ "aion_cache_dir": "local/cache",
+ }
+ values.update(kwargs)
+ return ModelConfig(**values)
+
+
+def test_siglip_remains_the_default_image_encoder_type():
+ assert ModelConfig().image_encoder_type == "siglip"
+
+
+@pytest.mark.parametrize(
+ ("override", "message"),
+ [
+ ({"d_img": AION_IMAGE_WIDTH - 1}, "d_img=768"),
+ ({"aion_num_image_tokens": AION_IMAGE_TOKENS - 1}, "exactly 576"),
+ (
+ {"aion_image_bands": ("DES-R", "DES-G", "DES-I", "DES-Z")},
+ "Legacy Survey channels",
+ ),
+ ({"aion_raw_image_size": (96, 96)}, "160x160"),
+ ],
+)
+def test_aion_config_rejects_incompatible_feature_contract(override, message):
+ with pytest.raises(ValueError, match=message):
+ _aion_config(**override)
+
+
+def test_model_selects_aion_with_pinned_offline_controls_and_768_adapter():
+ encoder = _FakeAIONImageEncoder()
+
+ with patch("src.model.AIONImageEncoder", return_value=encoder) as factory:
+ model = UnifiedTransformer(_aion_config())
+
+ factory.assert_called_once_with(
+ model_name_or_path="local/aion-base",
+ model_revision=AION_MODEL_REVISION,
+ codec_revision=AION_IMAGE_CODEC_REVISION,
+ cache_dir="local/cache",
+ local_files_only=True,
+ bands=AION_LEGACY_SURVEY_BANDS,
+ raw_image_size=(160, 160),
+ )
+ assert model.encoders["image"] is encoder
+ assert model.projectors["image"].fc1.in_features == AION_IMAGE_WIDTH
+ assert encoder.trainability_calls == [False]
+ assert model._encoder_frozen_cache == {"image": True}
+
+
+def test_model_trainability_hook_never_unfreezes_the_aion_codec():
+ encoder = _FakeAIONImageEncoder()
+
+ with patch("src.model.AIONImageEncoder", return_value=encoder):
+ model = UnifiedTransformer(_aion_config(freeze_encoders=False))
+
+ assert encoder.trainability_calls == [True]
+ assert encoder.shared_weight.requires_grad
+ assert not encoder.codec_weight.requires_grad
+ assert model._encoder_frozen_cache == {"image": False}
+
+ assert model.set_encoder_trainable("image", False) == 0
+ assert not encoder.shared_weight.requires_grad
+ assert not encoder.codec_weight.requires_grad
+ assert model._encoder_frozen_cache == {"image": True}
+
+ assert model.set_encoder_trainable("image", True) == encoder.shared_weight.numel()
+ assert encoder.shared_weight.requires_grad
+ assert not encoder.codec_weight.requires_grad
+ assert model._encoder_frozen_cache == {"image": False}
+
+
+def test_refresh_encoder_frozen_cache_tracks_external_hook_changes():
+ encoder = _FakeAIONImageEncoder()
+ with patch("src.model.AIONImageEncoder", return_value=encoder):
+ model = UnifiedTransformer(_aion_config())
+
+ encoder.set_shared_encoder_trainable(True)
+ assert model._encoder_frozen_cache == {"image": True}
+ assert model.refresh_encoder_frozen_cache() == {"image": False}
+
+
+def test_model_rejects_aion_encoder_with_unexpected_token_count():
+ encoder = _FakeAIONImageEncoder(token_count=AION_IMAGE_TOKENS - 1)
+ with (
+ patch("src.model.AIONImageEncoder", return_value=encoder),
+ pytest.raises(RuntimeError, match="token count"),
+ ):
+ UnifiedTransformer(_aion_config())
+
+
+def test_astroprism_hydra_preset_is_raw_aion_qwen3_4b():
+ config_path = Path("src/conf/model/prism_astroprism_qwen3_4b.yaml")
+ cfg = yaml.safe_load(config_path.read_text())
+
+ asset_root = "${oc.env:ASTROPRISM_ROOT,_local/astroprism}"
+ assert cfg["backbone_id"] == f"{asset_root}/hf/models/qwen3-4b"
+ assert cfg["tokenizer_id"] == f"{asset_root}/hf/models/qwen3-4b"
+ assert cfg["image_encoder_type"] == "aion"
+ assert cfg["d_img"] == AION_IMAGE_WIDTH
+ assert cfg["aion_model_name_or_path"] == AION_REPO_ID
+ # null means the ambient HF cache. huggingface_hub's cache_dir is the hub
+ # directory itself, unlike HF_HOME which has hub/ appended, so a path here
+ # resolved one level too high and never found the weights. The released
+ # checkpoint's training_config.yaml records null for the same reason.
+ assert cfg["aion_cache_dir"] is None
+ assert cfg["aion_model_revision"] == AION_MODEL_REVISION
+ assert cfg["aion_codec_revision"] == AION_IMAGE_CODEC_REVISION
+ assert cfg["aion_local_files_only"] is True
+ assert tuple(cfg["aion_image_bands"]) == AION_LEGACY_SURVEY_BANDS
+ assert cfg["aion_num_image_tokens"] == AION_IMAGE_TOKENS
+ assert cfg["aion_raw_image_size"] == [160, 160]
+ assert cfg["image_processor_id"] is None
+ assert cfg["projector_norm_mode"] == "none"
+ assert cfg["projector_modality_embed_pos"] == "after_norm"
diff --git a/tests/test_aion_multimodal_encoder.py b/tests/test_aion_multimodal_encoder.py
new file mode 100644
index 0000000..a3e7510
--- /dev/null
+++ b/tests/test_aion_multimodal_encoder.py
@@ -0,0 +1,404 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import pytest
+import src.encoders.aion_multimodal as aion_multimodal_module
+import torch
+import torch.nn as nn
+from src.config import ModelConfig
+from src.encoders.aion_image import (
+ AION_IMAGE_CODEC_REVISION,
+ AION_IMAGE_TOKENS,
+ AION_IMAGE_WIDTH,
+ AION_MODEL_REVISION,
+)
+from src.encoders.aion_multimodal import AIONMultimodalEncoder
+from src.encoders.aion_spectrum import (
+ AION_RAW_SPECTRUM_LENGTH,
+ AION_SPECTRUM_CODEC_REVISION,
+ AION_SPECTRUM_TOKENS,
+)
+from src.model import UnifiedTransformer
+
+pytestmark = pytest.mark.unit
+
+
+@dataclass
+class _FakeLegacySurveyImage:
+ flux: torch.Tensor
+ bands: list[str]
+
+ token_key = "tok_image"
+
+
+@dataclass
+class _FakeDESISpectrum:
+ flux: torch.Tensor
+ ivar: torch.Tensor
+ mask: torch.Tensor
+ wavelength: torch.Tensor
+
+ token_key = "tok_spectrum_desi"
+
+
+class _FakeImageCodec(nn.Module):
+ from_pretrained_calls: list[tuple[tuple, dict]] = []
+
+ def __init__(self) -> None:
+ super().__init__()
+ self.scale = nn.Parameter(torch.ones((), dtype=torch.float64))
+ self.last_image: _FakeLegacySurveyImage | None = None
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ cls.from_pretrained_calls.append((args, kwargs))
+ return cls()
+
+ def encode(self, image: _FakeLegacySurveyImage) -> torch.Tensor:
+ self.last_image = image
+ image.flux.zero_()
+ return torch.zeros(
+ image.flux.shape[0],
+ AION_IMAGE_TOKENS,
+ dtype=torch.long,
+ device=image.flux.device,
+ )
+
+
+class _FakeSpectrumCodec(nn.Module):
+ from_pretrained_calls: list[tuple[tuple, dict]] = []
+
+ def __init__(self) -> None:
+ super().__init__()
+ self.scale = nn.Parameter(torch.ones((), dtype=torch.float64))
+ self.last_spectrum: _FakeDESISpectrum | None = None
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ cls.from_pretrained_calls.append((args, kwargs))
+ return cls()
+
+ def encode(self, spectrum: _FakeDESISpectrum) -> torch.Tensor:
+ self.last_spectrum = spectrum
+ spectrum.flux.zero_()
+ return torch.zeros(
+ spectrum.flux.shape[0],
+ AION_SPECTRUM_TOKENS,
+ dtype=torch.long,
+ device=spectrum.flux.device,
+ )
+
+
+class _FakeAION(nn.Module):
+ from_pretrained_calls: list[tuple[tuple, dict]] = []
+
+ def __init__(self) -> None:
+ super().__init__()
+ self.encoder_embeddings = nn.ModuleDict(
+ {
+ "tok_image": nn.Embedding(8, 1),
+ "tok_spectrum_desi": nn.Embedding(8, 1),
+ "tok_unused": nn.Embedding(8, 1),
+ }
+ )
+ self.encoder = nn.ModuleList([nn.Linear(1, 1, bias=False)])
+ self.encoder_norm = nn.LayerNorm(1)
+ self.decoder_proj_context = nn.Linear(1, 1, bias=False)
+ self.unused_decoder = nn.Linear(1, 1)
+ self.encode_calls: list[tuple[dict[str, torch.Tensor], int]] = []
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ cls.from_pretrained_calls.append((args, kwargs))
+ return cls()
+
+ def encode(self, inputs: dict[str, torch.Tensor], *, num_encoder_tokens: int):
+ self.encode_calls.append((inputs, num_encoder_tokens))
+ batch_size = next(iter(inputs.values())).shape[0]
+ scale = self.encoder[0].weight.reshape(1, 1, 1)
+ return scale * torch.ones(
+ batch_size,
+ num_encoder_tokens,
+ AION_IMAGE_WIDTH,
+ dtype=scale.dtype,
+ device=scale.device,
+ )
+
+
+@pytest.fixture(autouse=True)
+def _mock_joint_aion_types(monkeypatch):
+ _FakeAION.from_pretrained_calls.clear()
+ _FakeImageCodec.from_pretrained_calls.clear()
+ _FakeSpectrumCodec.from_pretrained_calls.clear()
+ monkeypatch.setattr(
+ aion_multimodal_module,
+ "_load_aion_types",
+ lambda: (
+ _FakeAION,
+ _FakeImageCodec,
+ _FakeSpectrumCodec,
+ _FakeLegacySurveyImage,
+ _FakeDESISpectrum,
+ ),
+ )
+
+
+def _make_spectrum_inputs(batch_size: int = 2) -> dict[str, torch.Tensor]:
+ wavelength = torch.linspace(3600.0, 9800.0, AION_RAW_SPECTRUM_LENGTH)
+ return {
+ "flux": torch.randn(batch_size, AION_RAW_SPECTRUM_LENGTH),
+ "ivar": torch.rand(batch_size, AION_RAW_SPECTRUM_LENGTH),
+ "mask": torch.zeros(batch_size, AION_RAW_SPECTRUM_LENGTH, dtype=torch.bool),
+ "wavelength": wavelength.repeat(batch_size, 1),
+ }
+
+
+def _make_encoder(**kwargs) -> AIONMultimodalEncoder:
+ return AIONMultimodalEncoder(model_name_or_path="local/aion-base", **kwargs)
+
+
+def _joint_config(**kwargs) -> ModelConfig:
+ values = {
+ "d_model": 32,
+ "vocab_size": 64,
+ "num_layers": 1,
+ "num_heads": 4,
+ "num_experts": 2,
+ "modalities": ["image", "spectrum"],
+ "d_img": AION_IMAGE_WIDTH,
+ "d_spectrum": AION_IMAGE_WIDTH,
+ "image_encoder_type": "aion",
+ "spectrum_encoder_type": "aion",
+ "freeze_encoders": True,
+ "aion_model_name_or_path": "local/aion-base",
+ "aion_cache_dir": "local/cache",
+ "image_projector_norm_mode": "none",
+ "image_projector_modality_embed_pos": "after_norm",
+ "spectrum_projector_norm_mode": "text_scaled_layernorm",
+ "spectrum_projector_target_norm": 1.04,
+ "spectrum_projector_modality_embed_pos": "before_norm",
+ }
+ values.update(kwargs)
+ return ModelConfig(**values)
+
+
+def test_joint_encoder_loads_one_transformer_and_two_pinned_codecs(tmp_path: Path):
+ encoder = _make_encoder(cache_dir=tmp_path)
+
+ assert _FakeAION.from_pretrained_calls == [
+ (
+ ("local/aion-base",),
+ {
+ "revision": AION_MODEL_REVISION,
+ "cache_dir": str(tmp_path),
+ "local_files_only": True,
+ },
+ )
+ ]
+ assert _FakeImageCodec.from_pretrained_calls[0][1]["revision"] == (
+ AION_IMAGE_CODEC_REVISION
+ )
+ assert _FakeImageCodec.from_pretrained_calls[0][1]["modality"] is (
+ _FakeLegacySurveyImage
+ )
+ assert _FakeSpectrumCodec.from_pretrained_calls[0][1]["revision"] == (
+ AION_SPECTRUM_CODEC_REVISION
+ )
+ assert _FakeSpectrumCodec.from_pretrained_calls[0][1]["modality"] is (
+ _FakeDESISpectrum
+ )
+ assert encoder.tokens_per_instance("image") == AION_IMAGE_TOKENS
+ assert encoder.tokens_per_instance("spectrum") == AION_SPECTRUM_TOKENS
+
+
+def test_joint_encoder_routes_both_modalities_without_mutating_raw_inputs():
+ encoder = _make_encoder()
+ image = torch.randn(2, 4, 160, 160)
+ image_original = image.clone()
+ spectrum = _make_spectrum_inputs()
+ spectrum_original = {key: value.clone() for key, value in spectrum.items()}
+
+ image_features = encoder.encode_modality("image", image)
+ spectrum_features = encoder.encode_modality("spectrum", spectrum)
+
+ assert image_features.shape == (2, AION_IMAGE_TOKENS, AION_IMAGE_WIDTH)
+ assert spectrum_features.shape == (2, AION_SPECTRUM_TOKENS, AION_IMAGE_WIDTH)
+ assert torch.equal(image, image_original)
+ assert all(torch.equal(spectrum[key], spectrum_original[key]) for key in spectrum)
+ assert [set(call[0]) for call in encoder.aion_model.encode_calls] == [
+ {"tok_image"},
+ {"tok_spectrum_desi"},
+ ]
+ assert [call[1] for call in encoder.aion_model.encode_calls] == [
+ AION_IMAGE_TOKENS,
+ AION_SPECTRUM_TOKENS,
+ ]
+
+
+def test_joint_encoder_registers_shared_transformer_once_and_protects_codecs():
+ encoder = _make_encoder().to(dtype=torch.bfloat16)
+ encoder.set_shared_encoder_trainable(True).train()
+
+ state_keys = set(encoder.state_dict())
+ assert "aion_model.encoder.0.weight" in state_keys
+ assert sum(key.endswith("aion_model.encoder.0.weight") for key in state_keys) == 1
+ assert not any("image_encoder.aion_model" in key for key in state_keys)
+ assert not any("spectrum_encoder.aion_model" in key for key in state_keys)
+
+ for codec in (encoder.image_codec, encoder.spectrum_codec):
+ parameters = list(codec.parameters())
+ assert parameters
+ assert all(parameter.dtype == torch.float32 for parameter in parameters)
+ assert all(not parameter.requires_grad for parameter in parameters)
+ assert codec.training is False
+
+ assert all(
+ any(parameter.requires_grad for parameter in module.parameters())
+ for module in encoder._shared_encoder_modules()
+ )
+ assert all(module.training for module in encoder._shared_encoder_modules())
+ assert all(
+ not parameter.requires_grad
+ for parameter in encoder.aion_model.encoder_embeddings["tok_unused"].parameters()
+ )
+ assert all(
+ not parameter.requires_grad
+ for parameter in encoder.aion_model.unused_decoder.parameters()
+ )
+
+ encoder.encode_modality("image", torch.randn(1, 4, 160, 160)).sum().backward()
+ assert encoder.aion_model.encoder[0].weight.grad is not None
+ assert all(parameter.grad is None for parameter in encoder.image_codec.parameters())
+ assert all(parameter.grad is None for parameter in encoder.spectrum_codec.parameters())
+
+ encoder.set_shared_encoder_trainable(False).train()
+ assert all(not parameter.requires_grad for parameter in encoder.parameters())
+ assert encoder.image_codec.training is False
+ assert encoder.spectrum_codec.training is False
+
+
+def test_joint_model_owns_one_aion_and_keeps_adapter_contracts_separate():
+ model = UnifiedTransformer(_joint_config())
+
+ assert list(model.encoders) == ["aion"]
+ assert model._encoder_key_by_modality == {
+ "image": "aion",
+ "spectrum": "aion",
+ }
+ assert model._encoder_frozen_cache == {"image": True, "spectrum": True}
+ assert len(_FakeAION.from_pretrained_calls) == 1
+
+ image_projector = model.projectors["image"]
+ spectrum_projector = model.projectors["spectrum"]
+ assert image_projector.norm_mode == "none"
+ assert image_projector.modality_embed_pos == "after_norm"
+ assert image_projector.final_norm is None
+ assert spectrum_projector.norm_mode == "text_scaled_layernorm"
+ assert spectrum_projector.target_norm == pytest.approx(1.04)
+ assert spectrum_projector.modality_embed_pos == "before_norm"
+ assert spectrum_projector.final_norm is not None
+
+ [(image_name, image_features)] = model._process_multimodal_embeddings(
+ {"image": torch.randn(1, 4, 160, 160)}
+ )
+ [(spectrum_name, spectrum_features)] = model._process_multimodal_embeddings(
+ {"spectrum": _make_spectrum_inputs(batch_size=1)}
+ )
+ assert str(image_name) == "image"
+ assert image_features.shape == (1, AION_IMAGE_TOKENS, 32)
+ assert str(spectrum_name) == "spectrum"
+ assert spectrum_features.shape == (1, AION_SPECTRUM_TOKENS, 32)
+
+ state_keys = set(model.state_dict())
+ assert sum(key.endswith("aion_model.encoder.0.weight") for key in state_keys) == 1
+ assert "encoders.aion.aion_model.encoder.0.weight" in state_keys
+
+
+def test_joint_model_trainability_aliases_toggle_the_one_owner():
+ model = UnifiedTransformer(_joint_config())
+ encoder = model.encoders["aion"]
+
+ trainable_count = model.set_encoder_trainable("image", True)
+ assert trainable_count > 0
+ assert encoder.shared_encoder_trainable is True
+ assert model._encoder_frozen_cache == {"image": False, "spectrum": False}
+ assert all(not parameter.requires_grad for parameter in encoder.image_codec.parameters())
+ assert all(
+ not parameter.requires_grad for parameter in encoder.spectrum_codec.parameters()
+ )
+
+ assert model.set_encoder_trainable("spectrum", False) == 0
+ assert encoder.shared_encoder_trainable is False
+ assert model._encoder_frozen_cache == {"image": True, "spectrum": True}
+
+
+def test_per_modality_projector_fields_fall_back_to_global_defaults():
+ config = _joint_config(
+ projector_norm_mode="layernorm",
+ projector_target_norm=0.5,
+ projector_modality_embed_pos="none",
+ image_projector_norm_mode=None,
+ image_projector_modality_embed_pos=None,
+ spectrum_projector_norm_mode=None,
+ spectrum_projector_target_norm=None,
+ spectrum_projector_modality_embed_pos=None,
+ )
+ model = UnifiedTransformer(config)
+
+ for modality in ("image", "spectrum"):
+ projector = model.projectors[modality]
+ assert projector.norm_mode == "layernorm"
+ assert projector.target_norm == pytest.approx(0.5)
+ assert projector.modality_embed_pos == "none"
+
+@pytest.mark.parametrize(
+ ("modality", "start_idx", "end_idx", "token_count"),
+ [
+ ("image", 60, 61, AION_IMAGE_TOKENS),
+ ("spectrum", 62, 63, AION_SPECTRUM_TOKENS),
+ ],
+)
+def test_joint_merge_routes_token_count_through_shared_owner(
+ monkeypatch, modality: str, start_idx: int, end_idx: int, token_count: int
+):
+ model = UnifiedTransformer(_joint_config())
+ model.config.modality_start_end_token_indices = {
+ modality: (start_idx, end_idx)
+ }
+ owner = model.encoders["aion"]
+ calls: list[str] = []
+ original_tokens_per_instance = owner.tokens_per_instance
+
+ def record_tokens_per_instance(routed_modality: str) -> int:
+ calls.append(routed_modality)
+ return original_tokens_per_instance(routed_modality)
+
+ monkeypatch.setattr(owner, "tokens_per_instance", record_tokens_per_instance)
+
+ input_ids = torch.tensor([[5, start_idx, end_idx, 6]], dtype=torch.long)
+ input_embeds = torch.arange(4 * 32, dtype=torch.float32).reshape(1, 4, 32)
+ modality_features = torch.full((1, token_count, 32), 7.0)
+
+ final_embedding, attention_mask, position_ids, labels = (
+ model._merge_text_input_ids_with_modality_embeds(
+ input_ids=input_ids,
+ input_embeds=input_embeds,
+ embeddings_list=[
+ ("text", input_embeds),
+ (modality, modality_features),
+ ],
+ metadata=["4 0"],
+ )
+ )
+
+ assert calls == [modality, modality]
+ assert final_embedding.shape == (1, token_count + 2, 32)
+ assert torch.equal(final_embedding[:, 1 : token_count + 1], modality_features)
+ assert torch.equal(final_embedding[:, 0], input_embeds[:, 0])
+ assert torch.equal(final_embedding[:, -1], input_embeds[:, -1])
+ assert attention_mask.tolist() == [[1] * (token_count + 2)]
+ assert position_ids.tolist() == [list(range(token_count + 2))]
+ assert torch.all(labels == -100)
+
diff --git a/tests/test_aion_spectrum_encoder.py b/tests/test_aion_spectrum_encoder.py
new file mode 100644
index 0000000..5b1ee96
--- /dev/null
+++ b/tests/test_aion_spectrum_encoder.py
@@ -0,0 +1,280 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import pytest
+import src.encoders.aion_spectrum as aion_spectrum_module
+import torch
+import torch.nn as nn
+from src.encoders.aion_image import AION_MODEL_REVISION
+from src.encoders.aion_spectrum import (
+ AION_RAW_SPECTRUM_LENGTH,
+ AION_SPECTRUM_CODEC_REVISION,
+ AION_SPECTRUM_TOKENS,
+ AION_SPECTRUM_WIDTH,
+ AIONSpectrumEncoder,
+)
+
+pytestmark = pytest.mark.unit
+
+
+@dataclass
+class _FakeDESISpectrum:
+ flux: torch.Tensor
+ ivar: torch.Tensor
+ mask: torch.Tensor
+ wavelength: torch.Tensor
+
+ token_key = "tok_spectrum_desi"
+
+
+class _FakeSpectrumCodec(nn.Module):
+ from_pretrained_calls: list[tuple[tuple, dict]] = []
+
+ def __init__(self) -> None:
+ super().__init__()
+ self.scale = nn.Parameter(torch.ones((), dtype=torch.float64))
+ self.last_spectrum: _FakeDESISpectrum | None = None
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ cls.from_pretrained_calls.append((args, kwargs))
+ return cls()
+
+ def encode(self, spectrum: _FakeDESISpectrum) -> torch.Tensor:
+ self.last_spectrum = spectrum
+ batch_size = spectrum.flux.shape[0]
+ # Mimic in-place preprocessing to prove the wrapper isolates callers.
+ spectrum.flux.zero_()
+ spectrum.ivar.zero_()
+ spectrum.mask.fill_(True)
+ spectrum.wavelength.zero_()
+ return torch.zeros(
+ batch_size,
+ AION_SPECTRUM_TOKENS,
+ dtype=torch.long,
+ device=spectrum.flux.device,
+ )
+
+
+class _FakeAION(nn.Module):
+ from_pretrained_calls: list[tuple[tuple, dict]] = []
+
+ def __init__(self) -> None:
+ super().__init__()
+ self.encoder_embeddings = nn.ModuleDict({"tok_spectrum_desi": nn.Embedding(8, 1)})
+ self.encoder = nn.ModuleList([nn.Linear(1, 1, bias=False)])
+ self.encoder_norm = nn.LayerNorm(1)
+ self.decoder_proj_context = nn.Linear(1, 1, bias=False)
+ self.unused_decoder = nn.Linear(1, 1)
+ self.encode_calls: list[tuple[dict[str, torch.Tensor], int]] = []
+ self.output_shape = (AION_SPECTRUM_TOKENS, AION_SPECTRUM_WIDTH)
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ cls.from_pretrained_calls.append((args, kwargs))
+ return cls()
+
+ def encode(self, inputs: dict[str, torch.Tensor], *, num_encoder_tokens: int):
+ self.encode_calls.append((inputs, num_encoder_tokens))
+ batch_size = next(iter(inputs.values())).shape[0]
+ token_count, width = self.output_shape
+ scale = self.encoder[0].weight.reshape(1, 1, 1)
+ return scale * torch.ones(
+ batch_size,
+ token_count,
+ width,
+ dtype=scale.dtype,
+ device=scale.device,
+ )
+
+
+@pytest.fixture(autouse=True)
+def _mock_aion_types(monkeypatch):
+ _FakeAION.from_pretrained_calls.clear()
+ _FakeSpectrumCodec.from_pretrained_calls.clear()
+ monkeypatch.setattr(
+ aion_spectrum_module,
+ "_load_aion_types",
+ lambda: (_FakeAION, _FakeSpectrumCodec, _FakeDESISpectrum),
+ )
+
+
+def _make_encoder(**kwargs) -> AIONSpectrumEncoder:
+ return AIONSpectrumEncoder(model_name_or_path="local/aion-base", **kwargs)
+
+
+def _make_inputs(batch_size: int = 2) -> dict[str, torch.Tensor]:
+ wavelength = torch.linspace(3600.0, 9800.0, AION_RAW_SPECTRUM_LENGTH)
+ return {
+ "flux": torch.randn(batch_size, AION_RAW_SPECTRUM_LENGTH),
+ "ivar": torch.rand(batch_size, AION_RAW_SPECTRUM_LENGTH),
+ "mask": torch.zeros(batch_size, AION_RAW_SPECTRUM_LENGTH, dtype=torch.bool),
+ "wavelength": wavelength.repeat(batch_size, 1),
+ }
+
+
+def test_loads_model_and_spectrum_codec_with_pinned_offline_controls(tmp_path: Path):
+ encoder = _make_encoder(cache_dir=tmp_path)
+
+ assert _FakeAION.from_pretrained_calls == [
+ (
+ ("local/aion-base",),
+ {
+ "revision": AION_MODEL_REVISION,
+ "cache_dir": str(tmp_path),
+ "local_files_only": True,
+ },
+ )
+ ]
+ assert _FakeSpectrumCodec.from_pretrained_calls == [
+ (
+ ("local/aion-base",),
+ {
+ "modality": _FakeDESISpectrum,
+ "revision": AION_SPECTRUM_CODEC_REVISION,
+ "cache_dir": str(tmp_path),
+ "local_files_only": True,
+ },
+ )
+ ]
+ assert encoder.tokens_per_instance() == AION_SPECTRUM_TOKENS
+ assert encoder.output_dim == AION_SPECTRUM_WIDTH
+
+
+def test_raw_spectrum_is_cloned_and_encoded_as_exactly_273_aion_tokens():
+ encoder = _make_encoder()
+ inputs = _make_inputs()
+ originals = {name: tensor.clone() for name, tensor in inputs.items()}
+
+ features = encoder(inputs)
+
+ assert features.shape == (2, AION_SPECTRUM_TOKENS, AION_SPECTRUM_WIDTH)
+ for name in inputs:
+ assert torch.equal(inputs[name], originals[name])
+
+ spectrum = encoder.spectrum_codec.last_spectrum
+ assert spectrum is not None
+ assert spectrum.flux.dtype == torch.float32
+ assert spectrum.ivar.dtype == torch.float32
+ assert spectrum.wavelength.dtype == torch.float32
+ assert spectrum.mask.dtype == torch.bool
+ assert torch.count_nonzero(spectrum.flux) == 0
+
+ [(encoded_inputs, token_budget)] = encoder.aion_model.encode_calls
+ assert set(encoded_inputs) == {"tok_spectrum_desi"}
+ assert encoded_inputs["tok_spectrum_desi"].shape == (2, AION_SPECTRUM_TOKENS)
+ assert token_budget == AION_SPECTRUM_TOKENS
+
+
+def test_codec_stays_float32_frozen_and_eval_when_shared_encoder_is_unfrozen():
+ encoder = _make_encoder().to(dtype=torch.bfloat16)
+ encoder.set_shared_encoder_trainable(True)
+ encoder.train()
+
+ codec_parameters = list(encoder.spectrum_codec.parameters())
+ assert codec_parameters
+ assert all(parameter.dtype == torch.float32 for parameter in codec_parameters)
+ assert all(not parameter.requires_grad for parameter in codec_parameters)
+ assert encoder.spectrum_codec.training is False
+
+ shared_modules = encoder._shared_encoder_modules()
+ assert all(
+ any(parameter.requires_grad for parameter in module.parameters())
+ for module in shared_modules
+ )
+ assert all(module.training for module in shared_modules)
+ assert all(
+ not parameter.requires_grad for parameter in encoder.aion_model.unused_decoder.parameters()
+ )
+
+ encoder(_make_inputs(batch_size=1)).sum().backward()
+ assert encoder.aion_model.encoder[0].weight.grad is not None
+ assert all(parameter.grad is None for parameter in codec_parameters)
+
+ encoder.set_shared_encoder_trainable(False)
+ encoder.train()
+ assert all(not parameter.requires_grad for parameter in encoder.aion_model.parameters())
+ assert all(not module.training for module in shared_modules)
+ assert encoder.spectrum_codec.training is False
+
+
+def test_rejects_missing_or_malformed_spectrum_fields():
+ encoder = _make_encoder()
+ inputs = _make_inputs(batch_size=1)
+
+ missing = dict(inputs)
+ missing.pop("ivar")
+ with pytest.raises(ValueError, match="missing fields"):
+ encoder(missing)
+
+ wrong_shape = dict(inputs)
+ wrong_shape["flux"] = torch.zeros(1, AION_RAW_SPECTRUM_LENGTH - 1)
+ with pytest.raises(ValueError, match="must have shape"):
+ encoder(wrong_shape)
+
+ wrong_mask = dict(inputs)
+ wrong_mask["mask"] = inputs["mask"].float()
+ with pytest.raises(TypeError, match="mask.*boolean"):
+ encoder(wrong_mask)
+
+ negative_ivar = dict(inputs)
+ negative_ivar["ivar"] = -torch.ones_like(inputs["ivar"])
+ with pytest.raises(ValueError, match="non-negative"):
+ encoder(negative_ivar)
+
+ nonmonotonic = dict(inputs)
+ nonmonotonic["wavelength"] = inputs["wavelength"].flip(-1)
+ with pytest.raises(ValueError, match="strictly increasing"):
+ encoder(nonmonotonic)
+
+ nonfinite = dict(inputs)
+ nonfinite["flux"] = inputs["flux"].clone()
+ nonfinite["flux"][0, 0] = float("nan")
+ with pytest.raises(ValueError, match="non-finite"):
+ encoder(nonfinite)
+
+
+def test_rejects_unexpected_codec_token_shape():
+ encoder = _make_encoder()
+ encoder.spectrum_codec.encode = lambda _spectrum: torch.zeros(
+ 1, AION_SPECTRUM_TOKENS - 1, dtype=torch.long
+ )
+
+ with pytest.raises(RuntimeError, match="SpectrumCodec must emit token ids"):
+ encoder(_make_inputs(batch_size=1))
+
+
+def test_rejects_unexpected_shared_encoder_feature_shape():
+ encoder = _make_encoder()
+ encoder.aion_model.output_shape = (
+ AION_SPECTRUM_TOKENS,
+ AION_SPECTRUM_WIDTH - 1,
+ )
+
+ with pytest.raises(RuntimeError, match="shared encoder must emit features"):
+ encoder(_make_inputs(batch_size=1))
+
+
+def test_codec_values_survive_a_blanket_reduced_precision_cast():
+ """Dtype alone is not enough: rounded weights can change token ids.
+
+ ZoneATrainer casts the whole encoder ModuleDict to the backbone dtype
+ (src/training/trainer_zone_a.py), which reaches this encoder through
+ nn.Module._apply. Restoring float32 afterwards would recover the dtype but
+ not the values, and the codec quantizes hard.
+ """
+ encoder = _make_encoder()
+ before = [p.detach().clone() for p in encoder.spectrum_codec.parameters()]
+ assert before
+
+ import torch.nn as nn
+
+ bundle = nn.ModuleDict({"aion": encoder})
+ bundle.to(dtype=torch.bfloat16)
+
+ after = list(encoder.spectrum_codec.parameters())
+ assert all(p.dtype == torch.float32 for p in after)
+ for original, current in zip(before, after, strict=True):
+ assert torch.equal(original, current)
diff --git a/tests/test_aion_spectrum_model_wiring.py b/tests/test_aion_spectrum_model_wiring.py
new file mode 100644
index 0000000..24129dc
--- /dev/null
+++ b/tests/test_aion_spectrum_model_wiring.py
@@ -0,0 +1,155 @@
+from unittest.mock import patch
+
+import pytest
+import torch
+import torch.nn as nn
+from src.config import ModelConfig
+from src.encoders.aion_image import AION_MODEL_REVISION
+from src.encoders.aion_spectrum import (
+ AION_RAW_SPECTRUM_LENGTH,
+ AION_SPECTRUM_CODEC_REVISION,
+ AION_SPECTRUM_TOKENS,
+ AION_SPECTRUM_WIDTH,
+)
+from src.modalities import Modality
+from src.model import UnifiedTransformer
+
+pytestmark = pytest.mark.unit
+
+
+class _FakeAIONSpectrumEncoder(nn.Module):
+ output_dim = AION_SPECTRUM_WIDTH
+
+ def __init__(self, *, token_count: int = AION_SPECTRUM_TOKENS) -> None:
+ super().__init__()
+ self.shared_weight = nn.Parameter(torch.ones(1))
+ self.codec_weight = nn.Parameter(torch.ones(1), requires_grad=False)
+ self.token_count = token_count
+ self.trainability_calls: list[bool] = []
+
+ def tokens_per_instance(self) -> int:
+ return self.token_count
+
+ def set_shared_encoder_trainable(self, trainable: bool):
+ self.trainability_calls.append(bool(trainable))
+ self.shared_weight.requires_grad_(bool(trainable))
+ self.codec_weight.requires_grad_(False)
+ return self
+
+ def forward(self, inputs: dict[str, torch.Tensor]) -> torch.Tensor:
+ return self.shared_weight * torch.ones(
+ inputs["flux"].shape[0],
+ self.token_count,
+ self.output_dim,
+ device=inputs["flux"].device,
+ )
+
+
+def _spectrum_config(*, freeze_encoders: bool = True, **kwargs) -> ModelConfig:
+ values = {
+ "d_model": 32,
+ "vocab_size": 64,
+ "num_layers": 1,
+ "num_heads": 4,
+ "num_experts": 2,
+ "modalities": ["spectrum"],
+ "d_spectrum": AION_SPECTRUM_WIDTH,
+ "spectrum_encoder_type": "aion",
+ "freeze_encoders": freeze_encoders,
+ "aion_model_name_or_path": "local/aion-base",
+ "aion_cache_dir": "local/cache",
+ }
+ values.update(kwargs)
+ return ModelConfig(**values)
+
+
+@pytest.mark.parametrize(
+ ("override", "message"),
+ [
+ ({"d_spectrum": AION_SPECTRUM_WIDTH - 1}, "d_spectrum=768"),
+ (
+ {"aion_num_spectrum_tokens": AION_SPECTRUM_TOKENS - 1},
+ "exactly 273",
+ ),
+ (
+ {"aion_raw_spectrum_length": AION_RAW_SPECTRUM_LENGTH - 1},
+ "7781 samples",
+ ),
+ ({"spectrum_encoder_type": "generic"}, "must be 'aion'"),
+ ],
+)
+def test_spectrum_config_rejects_incompatible_feature_contract(override, message):
+ with pytest.raises(ValueError, match=message):
+ _spectrum_config(**override)
+
+
+def test_model_selects_aion_spectrum_with_its_own_768_adapter():
+ encoder = _FakeAIONSpectrumEncoder()
+
+ with patch("src.model.AIONSpectrumEncoder", return_value=encoder) as factory:
+ model = UnifiedTransformer(_spectrum_config())
+
+ factory.assert_called_once_with(
+ model_name_or_path="local/aion-base",
+ model_revision=AION_MODEL_REVISION,
+ codec_revision=AION_SPECTRUM_CODEC_REVISION,
+ cache_dir="local/cache",
+ local_files_only=True,
+ raw_spectrum_length=AION_RAW_SPECTRUM_LENGTH,
+ )
+ assert model.encoders["spectrum"] is encoder
+ assert model.projectors["spectrum"].fc1.in_features == AION_SPECTRUM_WIDTH
+ assert encoder.trainability_calls == [False]
+ assert model._encoder_frozen_cache == {"spectrum": True}
+
+
+def test_spectrum_encoder_hook_never_unfreezes_the_codec():
+ encoder = _FakeAIONSpectrumEncoder()
+
+ with patch("src.model.AIONSpectrumEncoder", return_value=encoder):
+ model = UnifiedTransformer(_spectrum_config(freeze_encoders=False))
+
+ assert encoder.trainability_calls == [True]
+ assert encoder.shared_weight.requires_grad
+ assert not encoder.codec_weight.requires_grad
+ assert model._encoder_frozen_cache == {"spectrum": False}
+
+ assert model.set_encoder_trainable("spectrum", False) == 0
+ assert not encoder.shared_weight.requires_grad
+ assert not encoder.codec_weight.requires_grad
+ assert model._encoder_frozen_cache == {"spectrum": True}
+
+ assert model.set_encoder_trainable("spectrum", True) == encoder.shared_weight.numel()
+ assert encoder.shared_weight.requires_grad
+ assert not encoder.codec_weight.requires_grad
+ assert model._encoder_frozen_cache == {"spectrum": False}
+
+
+def test_model_rejects_aion_spectrum_encoder_with_unexpected_token_count():
+ encoder = _FakeAIONSpectrumEncoder(token_count=AION_SPECTRUM_TOKENS - 1)
+ with (
+ patch("src.model.AIONSpectrumEncoder", return_value=encoder),
+ pytest.raises(RuntimeError, match="spectrum token count"),
+ ):
+ UnifiedTransformer(_spectrum_config())
+
+
+def test_config_accepts_joint_aion_with_explicit_projector_overrides():
+ config = ModelConfig(
+ modalities=["image", "spectrum"],
+ image_encoder_type="aion",
+ spectrum_encoder_type="aion",
+ d_img=768,
+ image_projector_norm_mode="none",
+ image_projector_modality_embed_pos="after_norm",
+ spectrum_projector_norm_mode="text_scaled_layernorm",
+ spectrum_projector_target_norm=1.04,
+ spectrum_projector_modality_embed_pos="before_norm",
+ )
+
+ assert config.modalities == [Modality.IMAGE, Modality.SPECTRUM]
+ assert config.image_projector_norm_mode == "none"
+ assert config.image_projector_modality_embed_pos == "after_norm"
+ assert config.spectrum_projector_norm_mode == "text_scaled_layernorm"
+ assert config.spectrum_projector_target_norm == pytest.approx(1.04)
+ assert config.spectrum_projector_modality_embed_pos == "before_norm"
diff --git a/tests/test_astroprism_joint_data.py b/tests/test_astroprism_joint_data.py
new file mode 100644
index 0000000..91df3de
--- /dev/null
+++ b/tests/test_astroprism_joint_data.py
@@ -0,0 +1,314 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import pytest
+import src.data.astroprism_joint_dataset as joint_dataset_module
+from accelerate.data_loader import BatchSamplerShard, prepare_data_loader
+from src.data.astroprism_joint_batch_sampler import AstroPRISMJointBatchSampler
+from src.data.astroprism_joint_collate import AstroPRISMJointCollator
+from src.data.astroprism_joint_dataset import (
+ AstroPRISMJointDataset,
+ AstroPRISMJointDatasetConfig,
+ AstroPRISMJointDatasetError,
+)
+from src.data.astroprism_stage1a_dataset import AstroPRISMStage1ADatasetConfig
+from src.data.astroprism_stage1b_spectrum_dataset import (
+ AstroPRISMStage1BSpectrumDatasetConfig,
+)
+from torch.utils.data import DataLoader, Dataset
+
+pytestmark = [pytest.mark.unit, pytest.mark.multimodal]
+
+
+@dataclass(frozen=True)
+class _ChildConfig:
+ split: str
+
+
+class _ChildDataset(Dataset):
+ def __init__(
+ self,
+ modality: str,
+ count: int,
+ *,
+ split: str,
+ config: Any | None = None,
+ ) -> None:
+ self.active_modalities = {modality}
+ self.modality = modality
+ self.config = config if config is not None else _ChildConfig(split=split)
+ self.samples = [
+ {modality: f"{modality}-{index}", "sample_id": f"{modality}-{index}"}
+ for index in range(count)
+ ]
+
+ def __len__(self) -> int:
+ return len(self.samples)
+
+ def __getitem__(self, index: int) -> dict[str, Any]:
+ return self.samples[index]
+
+
+def _joint(
+ split: str,
+ *,
+ image_count: int = 7,
+ spectrum_count: int = 5,
+) -> AstroPRISMJointDataset:
+ return AstroPRISMJointDataset(
+ _ChildDataset("image", image_count, split=split),
+ _ChildDataset("spectrum", spectrum_count, split=split),
+ split=split,
+ )
+
+
+def _batch_modality(dataset: AstroPRISMJointDataset, batch: list[int]) -> str:
+ modalities = {dataset.modality_for_index(index) for index in batch}
+ assert len(modalities) == 1
+ return modalities.pop()
+
+
+def test_joint_dataset_exposes_both_modalities_and_tags_copies() -> None:
+ dataset = _joint("train", image_count=3, spectrum_count=2)
+
+ assert dataset.active_modalities == {"image", "spectrum"}
+ assert dataset.modality_lengths == {"image": 3, "spectrum": 2}
+ assert list(dataset.modality_indices["image"]) == [0, 1, 2]
+ assert list(dataset.modality_indices["spectrum"]) == [3, 4]
+ assert dataset[0]["batch_modality"] == "image"
+ assert dataset[3]["batch_modality"] == "spectrum"
+ assert dataset[-1]["sample_id"] == "spectrum-1"
+ assert "batch_modality" not in dataset.image_dataset.samples[0]
+ assert "batch_modality" not in dataset.spectrum_dataset.samples[0]
+
+
+def test_joint_dataset_requires_children_on_the_configured_split() -> None:
+ with pytest.raises(AstroPRISMJointDatasetError, match="spectrum child split"):
+ AstroPRISMJointDataset(
+ _ChildDataset("image", 2, split="train"),
+ _ChildDataset("spectrum", 2, split="validation"),
+ split="train",
+ )
+
+
+def test_from_config_applies_one_split_to_both_source_datasets(monkeypatch) -> None:
+ image_config = AstroPRISMStage1ADatasetConfig(
+ raw_manifest_path=Path("image-raw.json"),
+ candidates_manifest_path=Path("image-captions.json"),
+ split="test",
+ )
+ spectrum_config = AstroPRISMStage1BSpectrumDatasetConfig(
+ raw_manifest_path=Path("spectrum-raw.json"),
+ captions_manifest_path=Path("spectrum-captions.json"),
+ split="validation",
+ )
+ captured: dict[str, Any] = {}
+
+ def build_image(config, tokenizer):
+ captured["image"] = (config, tokenizer)
+ return _ChildDataset("image", 2, split=config.split, config=config)
+
+ def build_spectrum(config, tokenizer):
+ captured["spectrum"] = (config, tokenizer)
+ return _ChildDataset("spectrum", 2, split=config.split, config=config)
+
+ monkeypatch.setattr(joint_dataset_module, "AstroPRISMStage1ADataset", build_image)
+ monkeypatch.setattr(
+ joint_dataset_module,
+ "AstroPRISMStage1BSpectrumDataset",
+ build_spectrum,
+ )
+ tokenizer = object()
+
+ dataset = AstroPRISMJointDataset.from_config(
+ AstroPRISMJointDatasetConfig(
+ image=image_config,
+ spectrum=spectrum_config,
+ split="train",
+ ),
+ tokenizer,
+ )
+
+ assert dataset.split == "train"
+ assert captured["image"][0].split == "train"
+ assert captured["spectrum"][0].split == "train"
+ assert captured["image"][1] is tokenizer
+ assert captured["spectrum"][1] is tokenizer
+
+
+def test_joint_collator_routes_and_tags_homogeneous_batches() -> None:
+ calls: list[tuple[str, list[dict[str, Any]]]] = []
+
+ def image_collator(batch):
+ calls.append(("image", batch))
+ return {"sample_ids": [sample["sample_id"] for sample in batch]}
+
+ def spectrum_collator(batch):
+ calls.append(("spectrum", batch))
+ return {"sample_ids": [sample["sample_id"] for sample in batch]}
+
+ collator = AstroPRISMJointCollator(image_collator, spectrum_collator)
+ output = collator(
+ [
+ {"image": "image-0", "sample_id": "image-0", "batch_modality": "image"},
+ {"image": "image-1", "sample_id": "image-1", "batch_modality": "image"},
+ ]
+ )
+
+ assert output == {
+ "sample_ids": ["image-0", "image-1"],
+ "batch_modality": "image",
+ }
+ assert calls[0][0] == "image"
+ assert all("batch_modality" not in sample for sample in calls[0][1])
+
+
+def test_joint_collator_fails_closed_on_mixed_modalities() -> None:
+ collator = AstroPRISMJointCollator(lambda batch: {}, lambda batch: {})
+ with pytest.raises(ValueError, match="modality-homogeneous"):
+ collator(
+ [
+ {"image": "image-0", "batch_modality": "image"},
+ {"spectrum": "spectrum-0", "batch_modality": "spectrum"},
+ ]
+ )
+
+
+def test_train_sampler_is_balanced_deterministic_and_epoch_shuffled() -> None:
+ dataset = _joint("train", image_count=61, spectrum_count=37)
+ sampler = AstroPRISMJointBatchSampler(
+ dataset,
+ batch_size=2,
+ seed=17,
+ num_replicas=12,
+ )
+
+ epoch_zero = list(sampler)
+ assert epoch_zero == list(sampler)
+ assert len(epoch_zero) == len(sampler) == 72
+ assert all(len(batch) == 2 for batch in epoch_zero)
+ assert all(
+ len({_batch_modality(dataset, batch) for batch in epoch_zero[start : start + 12]})
+ == 1
+ for start in range(0, len(epoch_zero), 12)
+ )
+ counts = {"image": 0, "spectrum": 0}
+ for batch in epoch_zero:
+ counts[_batch_modality(dataset, batch)] += 1
+ assert counts == {"image": 36, "spectrum": 36}
+
+ sampler.set_epoch(1)
+ epoch_one = list(sampler)
+ assert epoch_one == list(sampler)
+ assert epoch_one != epoch_zero
+
+
+def test_train_sampler_is_rank_synchronous_under_12_rank_accelerate_sharding() -> None:
+ dataset = _joint("train", image_count=61, spectrum_count=37)
+ sampler = AstroPRISMJointBatchSampler(
+ dataset,
+ batch_size=2,
+ seed=29,
+ num_replicas=12,
+ )
+ rank_batches = [
+ list(
+ BatchSamplerShard(
+ sampler,
+ num_processes=12,
+ process_index=rank,
+ split_batches=False,
+ even_batches=True,
+ )
+ )
+ for rank in range(12)
+ ]
+
+ assert {len(batches) for batches in rank_batches} == {6}
+ step_modalities: list[str] = []
+ for optimizer_step in range(6):
+ modalities = {
+ _batch_modality(dataset, rank_batches[rank][optimizer_step])
+ for rank in range(12)
+ }
+ assert len(modalities) == 1
+ step_modalities.append(modalities.pop())
+ assert step_modalities.count("image") == 3
+ assert step_modalities.count("spectrum") == 3
+
+
+
+def test_prepared_accelerate_loader_propagates_epoch_to_joint_sampler() -> None:
+ dataset = _joint("train", image_count=61, spectrum_count=37)
+ sampler = AstroPRISMJointBatchSampler(
+ dataset,
+ batch_size=2,
+ seed=43,
+ num_replicas=12,
+ )
+ loader = DataLoader(
+ dataset,
+ batch_sampler=sampler,
+ collate_fn=lambda samples: samples,
+ )
+ prepared = prepare_data_loader(
+ loader,
+ device=None,
+ num_processes=12,
+ process_index=0,
+ split_batches=False,
+ put_on_device=False,
+ )
+
+ prepared.set_epoch(0)
+ epoch_zero = [
+ [sample["sample_id"] for sample in batch]
+ for batch in prepared
+ ]
+ assert sampler.epoch == 0
+
+ prepared.set_epoch(4)
+ epoch_four = [
+ [sample["sample_id"] for sample in batch]
+ for batch in prepared
+ ]
+ assert sampler.epoch == 4
+ assert epoch_four != epoch_zero
+
+def test_sampler_resolves_world_size_from_environment(monkeypatch) -> None:
+ for variable in (
+ "WORLD_SIZE",
+ "ACCELERATE_NUM_PROCESSES",
+ "OMPI_COMM_WORLD_SIZE",
+ "PMI_SIZE",
+ ):
+ monkeypatch.delenv(variable, raising=False)
+ monkeypatch.setenv("WORLD_SIZE", "12")
+
+ sampler = AstroPRISMJointBatchSampler(_joint("train"), batch_size=2)
+
+ assert sampler.num_replicas == 12
+
+
+@pytest.mark.parametrize("split", ["validation", "test"])
+def test_evaluation_sampler_is_deterministic_and_never_repeats(split: str) -> None:
+ dataset = _joint(split, image_count=5, spectrum_count=3)
+ sampler = AstroPRISMJointBatchSampler(
+ dataset,
+ batch_size=2,
+ seed=31,
+ num_replicas=12,
+ )
+
+ batches = list(sampler)
+ sampler.set_epoch(9)
+
+ assert list(sampler) == batches
+ assert len(batches) == len(sampler) == 5
+ assert all(len({_batch_modality(dataset, batch)}) == 1 for batch in batches)
+ flattened = [index for batch in batches for index in batch]
+ assert sorted(flattened) == list(range(len(dataset)))
+
diff --git a/tests/test_astroprism_joint_train_entrypoint.py b/tests/test_astroprism_joint_train_entrypoint.py
new file mode 100644
index 0000000..d422a6a
--- /dev/null
+++ b/tests/test_astroprism_joint_train_entrypoint.py
@@ -0,0 +1,143 @@
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import src.train as train_module
+import yaml
+from hydra import compose, initialize_config_dir
+from omegaconf import OmegaConf
+
+
+def test_joint_hydra_presets_preserve_both_adapter_contracts():
+ config_dir = str((Path(__file__).parents[1] / "src" / "conf").resolve())
+ with initialize_config_dir(config_dir=config_dir, version_base="1.2"):
+ cfg = compose(
+ config_name="config",
+ overrides=[
+ "model=prism_astroprism_joint_qwen3_4b",
+ "training=astroprism_joint_aion",
+ "+data=astroprism_joint_aion",
+ ],
+ )
+
+ assert cfg.training.task == "astroprism_joint_aion"
+ assert cfg.training.max_steps == 200
+ assert cfg.training.freeze_llm is True
+ assert cfg.training.freeze_vit is True
+ assert cfg.training.eval_enabled is True
+ assert cfg.training.eval_every_n_steps == 50
+ assert cfg.training.validation_batches == 0
+ assert set(cfg.model.modalities) == {"text", "image", "spectrum"}
+ assert cfg.model.image_projector_norm_mode == "none"
+ assert cfg.model.image_projector_modality_embed_pos == "after_norm"
+ assert cfg.model.spectrum_projector_norm_mode == "text_scaled_layernorm"
+ assert cfg.model.spectrum_projector_target_norm == pytest.approx(1.04)
+ assert cfg.model.spectrum_projector_modality_embed_pos == "before_norm"
+ assert cfg.data.split == "train"
+
+
+def test_joint_builder_forces_one_split_into_both_children(monkeypatch):
+ captured = {}
+
+ class _Child:
+ def __init__(self, modality, split):
+ self.active_modalities = {modality}
+ self.config = SimpleNamespace(split=split)
+
+ def __len__(self):
+ return 2
+
+ def __getitem__(self, index):
+ return {"index": index}
+
+ def image_builder(config, tokenizer):
+ captured["image"] = config
+ return _Child("image", config.split), object()
+
+ def spectrum_builder(config, tokenizer):
+ captured["spectrum"] = config
+ return _Child("spectrum", config.split), object()
+
+ class _JointCollator:
+ def __init__(self, image, spectrum):
+ captured["collators"] = (image, spectrum)
+
+ monkeypatch.setattr(train_module, "_build_astroprism_stage1a_data", image_builder)
+ monkeypatch.setattr(
+ train_module,
+ "_build_astroprism_stage1b_spectrum_data",
+ spectrum_builder,
+ )
+ monkeypatch.setattr(train_module, "AstroPRISMJointCollator", _JointCollator)
+
+ data_cfg = OmegaConf.create(
+ {
+ "split": "train",
+ "max_merged_seq_length": 2048,
+ "image": {"raw_manifest_path": "image.json"},
+ "spectrum": {"raw_manifest_path": "spectrum.json"},
+ }
+ )
+ dataset, _ = train_module._build_astroprism_joint_data(
+ data_cfg,
+ object(),
+ split="validation",
+ )
+
+ assert dataset.split == "validation"
+ assert captured["image"].split == "validation"
+ assert captured["spectrum"].split == "validation"
+ assert captured["image"].max_merged_seq_length == 2048
+ assert captured["spectrum"].max_merged_seq_length == 2048
+
+
+def test_joint_model_contract_is_fail_closed():
+ valid = SimpleNamespace(
+ image_encoder_type="aion",
+ spectrum_encoder_type="aion",
+ modalities=["text", "image", "spectrum"],
+ )
+ train_module._validate_astroprism_joint_model_contract(
+ "astroprism_joint_aion", valid
+ )
+
+ with pytest.raises(RuntimeError, match="requires exactly"):
+ train_module._validate_astroprism_joint_model_contract(
+ "astroprism_joint_aion",
+ SimpleNamespace(
+ image_encoder_type="aion",
+ spectrum_encoder_type="aion",
+ modalities=["text", "image"],
+ ),
+ )
+
+
+def test_joint_runtime_rejects_deepspeed_loader_bypass(monkeypatch):
+ monkeypatch.setenv("DEEPSPEED_ZERO_STAGE", "2")
+ with pytest.raises(RuntimeError, match="does not support DeepSpeed"):
+ train_module._validate_astroprism_joint_runtime_environment(
+ "astroprism_joint_aion"
+ )
+
+ train_module._validate_astroprism_joint_runtime_environment("vlm")
+ monkeypatch.delenv("DEEPSPEED_ZERO_STAGE")
+ train_module._validate_astroprism_joint_runtime_environment(
+ "astroprism_joint_aion"
+ )
+
+
+def test_joint_experiment_tracks_validation_and_final_test_offline():
+ [experiment] = yaml.safe_load(
+ Path("experiments/astroprism_joint_aion.yaml").read_text(encoding="utf-8")
+ )["experiments"]
+ overrides = experiment["common_overrides"]
+
+ assert experiment["id"] == "ASTROPRISM-JOINT-AION"
+ assert overrides["training.task"] == "astroprism_joint_aion"
+ assert overrides["training.max_steps"] == 200
+ assert overrides["training.eval_enabled"] is True
+ assert overrides["training.validation_batches"] == 0
+ assert overrides["training.freeze_llm"] is True
+ assert overrides["training.freeze_vit"] is True
+ assert overrides["wandb.mode"] == "offline"
+ assert "astroprism_joint_aion" in train_module._ASTROPRISM_MAP_STYLE_TASKS
diff --git a/tests/test_astroprism_raw_images.py b/tests/test_astroprism_raw_images.py
new file mode 100644
index 0000000..39100bc
--- /dev/null
+++ b/tests/test_astroprism_raw_images.py
@@ -0,0 +1,443 @@
+from __future__ import annotations
+
+import csv
+import gzip
+import json
+from dataclasses import replace
+from pathlib import Path
+from typing import Any
+
+import h5py
+import numpy as np
+import pytest
+import src.data.astroprism_raw_images as raw_images
+from src.data.astroprism_raw_images import (
+ RawImageConfig,
+ RawImageValidationError,
+ extract_raw_images,
+ sha256_file,
+)
+
+_ROSTER_FIELDS = [
+ "paired_targetid",
+ "caption_object_id",
+ "paired_ra",
+ "paired_dec",
+ "caption_ra",
+ "caption_dec",
+]
+_SOURCE_IDS = np.asarray([101, 102, 103, 104], dtype=np.int64)
+_RAS = np.asarray([10.0, 20.0, 30.0, 40.0], dtype=np.float64)
+_DECS = np.asarray([-3.0, -2.0, -1.0, 0.0], dtype=np.float64)
+_LEGACY_IDS = ["cap-a", "cap-b", "cap-c", "cap-d"]
+_ROSTER_SOURCE_ROWS = (3, 0, 1)
+_SOURCE_BANDS = np.asarray(
+ [[b"des-g", b"des-r", b"des-i", b"des-z"]] * 4,
+ dtype="S5",
+)
+
+
+def _write_source(path: Path) -> None:
+ row_values = np.arange(1, 5, dtype=np.float32)
+ with h5py.File(path, "w") as handle:
+ provabgs = handle.create_group("provabgs_desi_edr_sv3")
+ provabgs.create_dataset("object_id", data=_SOURCE_IDS)
+ provabgs.create_dataset("ra", data=_RAS)
+ provabgs.create_dataset("dec", data=_DECS)
+
+ desi = handle.create_group("desi")
+ desi.create_dataset("object_id", data=_SOURCE_IDS)
+ desi.create_dataset("ra", data=_RAS)
+ desi.create_dataset("dec", data=_DECS)
+ # Deliberately present: the vision-only extractor must never copy spectra.
+ desi.create_dataset("spectrum_flux", data=np.arange(28, dtype=np.float32).reshape(4, 7))
+
+ legacy = handle.create_group("legacysurvey")
+ legacy.create_dataset("object_id", data=np.asarray(_LEGACY_IDS, dtype="S16"))
+ legacy.create_dataset("ra", data=_RAS)
+ legacy.create_dataset("dec", data=_DECS)
+ legacy.create_dataset(
+ "image_array",
+ data=np.broadcast_to(row_values[:, None, None, None], (4, 4, 160, 160)).copy(),
+ )
+ legacy.create_dataset("image_band", data=_SOURCE_BANDS)
+ legacy.create_dataset(
+ "image_ivar",
+ data=np.broadcast_to((row_values + 10)[:, None, None, None], (4, 4, 160, 160)).copy(),
+ )
+ legacy.create_dataset("image_mask", data=np.zeros((4, 160, 160), dtype=np.bool_))
+ legacy.create_dataset(
+ "image_scale",
+ data=np.broadcast_to(row_values[:, None], (4, 4)).copy(),
+ )
+ legacy.create_dataset("object_mask", data=np.zeros((4, 160, 160), dtype=np.uint8))
+
+
+def _write_roster(path: Path, *, fields: list[str] | None = None) -> None:
+ rows = []
+ for source_row in _ROSTER_SOURCE_ROWS:
+ rows.append(
+ {
+ "paired_targetid": str(_SOURCE_IDS[source_row]),
+ "caption_object_id": _LEGACY_IDS[source_row],
+ "paired_ra": str(_RAS[source_row]),
+ "paired_dec": str(_DECS[source_row]),
+ "caption_ra": str(_RAS[source_row]),
+ "caption_dec": str(_DECS[source_row]),
+ }
+ )
+ fieldnames = fields or _ROSTER_FIELDS
+ with gzip.open(path, "wt", encoding="utf-8", newline="") as handle:
+ writer = csv.DictWriter(handle, fieldnames=fieldnames, extrasaction="ignore")
+ writer.writeheader()
+ writer.writerows(rows)
+
+
+def _build_case(tmp_path: Path) -> RawImageConfig:
+ source = tmp_path / "source.h5"
+ roster = tmp_path / "roster.csv.gz"
+ _write_source(source)
+ _write_roster(roster)
+ return RawImageConfig(
+ roster_path=roster,
+ output_dir=tmp_path / "output",
+ source=source,
+ expected_count=3,
+ expected_roster_sha256=sha256_file(roster),
+ expected_source_length=source.stat().st_size,
+ expected_source_etag=None,
+ caption_legacy_coordinate_tolerance_arcsec=0.025,
+ )
+
+
+def _manifest(config: RawImageConfig) -> dict[str, Any]:
+ return json.loads((config.output_dir / "manifest.json").read_text(encoding="utf-8"))
+
+
+class _ProbeResponse:
+ def __init__(self, *, status: int = 206) -> None:
+ self.status = status
+ self.headers = {
+ "Content-Range": "bytes 0-0/1234",
+ "Content-Length": "1",
+ "ETag": '"test-etag"',
+ }
+
+ def __enter__(self) -> _ProbeResponse:
+ return self
+
+ def __exit__(self, *args: object) -> None:
+ return None
+
+ def getcode(self) -> int:
+ return self.status
+
+ def read(self, size: int) -> bytes:
+ assert size == 2
+ return b"x"
+
+ def geturl(self) -> str:
+ return "https://example.test/final.h5"
+
+
+def test_extracts_all_mmap_outputs_in_roster_order_with_manifest_contract(
+ tmp_path: Path,
+) -> None:
+ config = _build_case(tmp_path)
+
+ summary = extract_raw_images(config)
+
+ assert summary["status"] == "prepared"
+ assert summary["candidate_count"] == 3
+ assert summary["contiguous_run_count"] == 2
+ assert summary["spectra_included"] is False
+
+ manifest = _manifest(config)
+ assert manifest["roster"]["sha256"] == sha256_file(config.roster_path)
+ assert manifest["selection"] == {
+ "acceptance_status": "pending_review",
+ "candidate_count": 3,
+ "contiguous_run_count": 2,
+ "http_cache_type": None,
+ "raw_row_max": 3,
+ "raw_row_min": 0,
+ "raw_rows_in_roster_order": [3, 0, 1],
+ "roster_order_preserved": True,
+ "split": None,
+ }
+ assert manifest["identity_tolerances_arcsec"]["caption_to_legacysurvey"] == 0.025
+ assert manifest["spectra_included"] is False
+ assert [candidate["paired_targetid"] for candidate in manifest["candidates"]] == [
+ "104",
+ "101",
+ "102",
+ ]
+ assert [candidate["caption_object_id"] for candidate in manifest["candidates"]] == [
+ "cap-d",
+ "cap-a",
+ "cap-b",
+ ]
+ assert all(
+ candidate["acceptance_status"] == "pending_review"
+ and candidate["split"] is None
+ and candidate["coordinate_crosscheck_arcsec"]["caption_to_legacysurvey"] == 0.0
+ for candidate in manifest["candidates"]
+ )
+
+ expected_outputs = {
+ "image_array",
+ "image_band",
+ "aion_image_band",
+ "image_ivar",
+ "image_mask",
+ "image_scale",
+ "object_mask",
+ }
+ assert set(manifest["outputs"]) == expected_outputs
+ assert not any("spectrum" in item["path"] for item in manifest["outputs"].values())
+ assert len({item["path"] for item in manifest["outputs"].values()}) == len(expected_outputs)
+ for item in manifest["outputs"].values():
+ output_path = config.output_dir / item["path"]
+ array = np.load(output_path, mmap_mode="r")
+ assert isinstance(array, np.memmap)
+ assert item["mmap_safe"] is True
+ assert item["sha256"] == sha256_file(output_path)
+ assert item["size_bytes"] == output_path.stat().st_size
+
+ image_array = np.load(config.output_dir / "image_array.npy", mmap_mode="r")
+ assert image_array[:, 0, 0, 0].tolist() == [4.0, 1.0, 2.0]
+ source_bands = np.load(config.output_dir / "image_band.npy", mmap_mode="r")
+ aion_bands = np.load(config.output_dir / "aion_image_band.npy", mmap_mode="r")
+ assert source_bands.tolist() == _SOURCE_BANDS[[3, 0, 1]].tolist()
+ assert aion_bands.tolist() == [
+ [b"DES-G", b"DES-R", b"DES-I", b"DES-Z"],
+ [b"DES-G", b"DES-R", b"DES-I", b"DES-Z"],
+ [b"DES-G", b"DES-R", b"DES-I", b"DES-Z"],
+ ]
+
+
+@pytest.mark.parametrize(
+ "missing_field",
+ ["caption_object_id", "caption_ra", "caption_dec"],
+)
+def test_requires_caption_identity_fields(tmp_path: Path, missing_field: str) -> None:
+ config = _build_case(tmp_path)
+ _write_roster(config.roster_path, fields=[f for f in _ROSTER_FIELDS if f != missing_field])
+ config = replace(config, expected_roster_sha256=sha256_file(config.roster_path))
+
+ with pytest.raises(RawImageValidationError, match="missing required columns"):
+ extract_raw_images(config)
+
+ assert not config.output_dir.exists()
+
+
+@pytest.mark.parametrize(
+ ("mutation", "message"),
+ [
+ ("identity", "does not match caption_object_id"),
+ ("coordinate", "caption/Legacy Survey coordinate separation"),
+ ],
+)
+def test_rejects_caption_legacy_mismatch_without_partial_output(
+ tmp_path: Path, mutation: str, message: str
+) -> None:
+ config = _build_case(tmp_path)
+ if mutation == "identity":
+ with h5py.File(config.source, "r+") as handle:
+ handle["legacysurvey/object_id"][3] = b"wrong-id"
+ config = replace(config, expected_source_length=Path(config.source).stat().st_size)
+ else:
+ with gzip.open(config.roster_path, "rt", encoding="utf-8", newline="") as handle:
+ rows = list(csv.DictReader(handle))
+ rows[0]["caption_ra"] = str(_RAS[3] + 0.01)
+ with gzip.open(config.roster_path, "wt", encoding="utf-8", newline="") as handle:
+ writer = csv.DictWriter(handle, fieldnames=_ROSTER_FIELDS)
+ writer.writeheader()
+ writer.writerows(rows)
+ config = replace(config, expected_roster_sha256=sha256_file(config.roster_path))
+
+ with pytest.raises(RawImageValidationError, match=message):
+ extract_raw_images(config)
+
+ assert not config.output_dir.exists()
+ assert not list(tmp_path.glob(".output.tmp-*"))
+
+
+def test_rejects_swapped_band_order_and_cleans_staging(tmp_path: Path) -> None:
+ config = _build_case(tmp_path)
+ with h5py.File(config.source, "r+") as handle:
+ handle["legacysurvey/image_band"][3] = np.asarray(
+ [b"des-r", b"des-g", b"des-i", b"des-z"], dtype="S5"
+ )
+ config = replace(config, expected_source_length=Path(config.source).stat().st_size)
+
+ with pytest.raises(RawImageValidationError, match="strict des-g/des-r/des-i/des-z order"):
+ extract_raw_images(config)
+
+ assert not config.output_dir.exists()
+ assert not list(tmp_path.glob(".output.tmp-*"))
+
+
+def test_http_probe_requires_one_byte_pinned_range(monkeypatch: pytest.MonkeyPatch) -> None:
+ calls: list[tuple[Any, float]] = []
+
+ def fake_urlopen(request: Any, timeout: float) -> _ProbeResponse:
+ calls.append((request, timeout))
+ return _ProbeResponse()
+
+ monkeypatch.setattr(raw_images.urllib.request, "urlopen", fake_urlopen)
+
+ info = raw_images._probe_http_source(
+ "https://example.test/source.h5",
+ expected_length=1234,
+ expected_etag='"test-etag"',
+ timeout_seconds=7.5,
+ )
+
+ request, timeout = calls.pop()
+ headers = dict(request.header_items())
+ assert request.get_method() == "GET"
+ assert headers["Range"] == "bytes=0-0"
+ assert headers["If-match"] == '"test-etag"'
+ assert timeout == 7.5
+ assert info.kind == "http"
+ assert info.location == "https://example.test/final.h5"
+ assert info.length_bytes == 1234
+ assert info.range_verified is True
+ assert info.status_code == 206
+ assert info.content_range == "bytes 0-0/1234"
+
+
+def test_http_probe_retries_transport_failures_with_exponential_backoff(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ outcomes: list[OSError | _ProbeResponse] = [
+ raw_images.urllib.error.URLError("temporary DNS failure"),
+ TimeoutError("TLS handshake timed out"),
+ _ProbeResponse(),
+ ]
+ sleeps: list[float] = []
+ calls = 0
+
+ def fake_urlopen(request: Any, timeout: float) -> _ProbeResponse:
+ nonlocal calls
+ calls += 1
+ outcome = outcomes.pop(0)
+ if isinstance(outcome, OSError):
+ raise outcome
+ return outcome
+
+ monkeypatch.setattr(raw_images.urllib.request, "urlopen", fake_urlopen)
+ monkeypatch.setattr(raw_images.time, "sleep", sleeps.append)
+
+ info = raw_images._probe_http_source(
+ "https://example.test/source.h5",
+ expected_length=1234,
+ expected_etag='"test-etag"',
+ timeout_seconds=7.5,
+ attempts=3,
+ backoff_seconds=0.25,
+ )
+
+ assert info.range_verified is True
+ assert calls == 3
+ assert sleeps == [0.25, 0.5]
+
+
+def test_http_probe_exhausts_bounded_transport_retries(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ sleeps: list[float] = []
+ calls = 0
+
+ def fake_urlopen(request: Any, timeout: float) -> _ProbeResponse:
+ nonlocal calls
+ calls += 1
+ raise TimeoutError("TLS handshake timed out")
+
+ monkeypatch.setattr(raw_images.urllib.request, "urlopen", fake_urlopen)
+ monkeypatch.setattr(raw_images.time, "sleep", sleeps.append)
+
+ with pytest.raises(RawImageValidationError, match="failed .* after 3 attempts"):
+ raw_images._probe_http_source(
+ "https://example.test/source.h5",
+ expected_length=1234,
+ expected_etag='"test-etag"',
+ timeout_seconds=7.5,
+ attempts=3,
+ backoff_seconds=0.125,
+ )
+
+ assert calls == 3
+ assert sleeps == [0.125, 0.25]
+
+
+def test_http_probe_does_not_retry_invalid_response(monkeypatch: pytest.MonkeyPatch) -> None:
+ calls = 0
+ sleeps: list[float] = []
+
+ def fake_urlopen(request: Any, timeout: float) -> _ProbeResponse:
+ nonlocal calls
+ calls += 1
+ return _ProbeResponse(status=200)
+
+ monkeypatch.setattr(raw_images.urllib.request, "urlopen", fake_urlopen)
+ monkeypatch.setattr(raw_images.time, "sleep", sleeps.append)
+
+ with pytest.raises(RawImageValidationError, match="requires status 206"):
+ raw_images._probe_http_source(
+ "https://example.test/source.h5",
+ expected_length=1234,
+ expected_etag='"test-etag"',
+ timeout_seconds=7.5,
+ attempts=3,
+ backoff_seconds=0.25,
+ )
+
+ assert calls == 1
+ assert sleeps == []
+
+
+def test_http_probe_does_not_retry_http_error(monkeypatch: pytest.MonkeyPatch) -> None:
+ calls = 0
+ sleeps: list[float] = []
+
+ def fake_urlopen(request: Any, timeout: float) -> _ProbeResponse:
+ nonlocal calls
+ calls += 1
+ raise raw_images.urllib.error.HTTPError(
+ request.full_url, 503, "Service Unavailable", hdrs=None, fp=None
+ )
+
+ monkeypatch.setattr(raw_images.urllib.request, "urlopen", fake_urlopen)
+ monkeypatch.setattr(raw_images.time, "sleep", sleeps.append)
+
+ with pytest.raises(RawImageValidationError, match="returned status 503"):
+ raw_images._probe_http_source(
+ "https://example.test/source.h5",
+ expected_length=1234,
+ expected_etag='"test-etag"',
+ timeout_seconds=7.5,
+ attempts=3,
+ backoff_seconds=0.25,
+ )
+
+ assert calls == 1
+ assert sleeps == []
+
+
+def test_existing_output_requires_explicit_overwrite(tmp_path: Path) -> None:
+ config = _build_case(tmp_path)
+ extract_raw_images(config)
+ stale = config.output_dir / "stale.txt"
+ stale.write_text("old", encoding="utf-8")
+
+ with pytest.raises(FileExistsError, match="output already exists"):
+ extract_raw_images(config)
+ assert stale.is_file()
+
+ summary = extract_raw_images(replace(config, overwrite=True))
+
+ assert summary["status"] == "prepared"
+ assert not stale.exists()
+ assert (config.output_dir / "manifest.json").is_file()
diff --git a/tests/test_astroprism_raw_spectra.py b/tests/test_astroprism_raw_spectra.py
new file mode 100644
index 0000000..e20bf5e
--- /dev/null
+++ b/tests/test_astroprism_raw_spectra.py
@@ -0,0 +1,282 @@
+from __future__ import annotations
+
+import argparse
+import csv
+import gzip
+import json
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any
+
+import h5py
+import numpy as np
+import pytest
+import tools.extract_astroprism_raw_spectra as spectrum_cli
+from src.data.astroprism_raw_images import sha256_file
+from src.data.astroprism_raw_spectra import (
+ ASTROPRISM_SPECTRUM_LENGTH,
+ RawSpectrumConfig,
+ RawSpectrumValidationError,
+ extract_raw_spectra,
+)
+
+pytestmark = [pytest.mark.unit, pytest.mark.multimodal]
+
+_ROSTER_FIELDS = [
+ "paired_targetid",
+ "caption_object_id",
+ "paired_ra",
+ "paired_dec",
+ "caption_ra",
+ "caption_dec",
+]
+_SOURCE_IDS = np.asarray([101, 102, 103, 104], dtype=np.int64)
+_RAS = np.asarray([10.0, 20.0, 30.0, 40.0], dtype=np.float64)
+_DECS = np.asarray([-3.0, -2.0, -1.0, 0.0], dtype=np.float64)
+_LEGACY_IDS = ["cap-a", "cap-b", "cap-c", "cap-d"]
+_ROSTER_SOURCE_ROWS = (3, 0, 1)
+
+
+def _write_source(path: Path) -> None:
+ source_count = len(_SOURCE_IDS)
+ row_values = np.arange(1, source_count + 1, dtype=np.float32)
+ pixels = np.arange(ASTROPRISM_SPECTRUM_LENGTH, dtype=np.float32)
+ flux = np.stack([100.0 * index + pixels / 100.0 for index in range(source_count)])
+ ivar = np.broadcast_to(row_values[:, None], flux.shape).copy()
+ mask = np.zeros(flux.shape, dtype=np.bool_)
+ mask[:, ::257] = True
+ wavelength = np.broadcast_to(3500.0 + pixels, flux.shape).copy()
+
+ with h5py.File(path, "w") as handle:
+ provabgs = handle.create_group("provabgs_desi_edr_sv3")
+ provabgs.create_dataset("object_id", data=_SOURCE_IDS)
+ provabgs.create_dataset("ra", data=_RAS)
+ provabgs.create_dataset("dec", data=_DECS)
+
+ desi = handle.create_group("desi")
+ desi.create_dataset("object_id", data=_SOURCE_IDS)
+ desi.create_dataset("ra", data=_RAS)
+ desi.create_dataset("dec", data=_DECS)
+ desi.create_dataset("spectrum_flux", data=flux)
+ desi.create_dataset("spectrum_ivar", data=ivar)
+ desi.create_dataset("spectrum_mask", data=mask)
+ desi.create_dataset("spectrum_lambda", data=wavelength)
+
+ # The shared source validator checks the complete official schema. These
+ # image fields must be present, but the spectrum extractor must not emit them.
+ legacy = handle.create_group("legacysurvey")
+ legacy.create_dataset("object_id", data=np.asarray(_LEGACY_IDS, dtype="S16"))
+ legacy.create_dataset("ra", data=_RAS)
+ legacy.create_dataset("dec", data=_DECS)
+ legacy.create_dataset(
+ "image_array",
+ data=np.broadcast_to(
+ row_values[:, None, None, None], (source_count, 4, 160, 160)
+ ).copy(),
+ )
+ legacy.create_dataset(
+ "image_band",
+ data=np.asarray([[b"des-g", b"des-r", b"des-i", b"des-z"]] * source_count),
+ )
+ legacy.create_dataset(
+ "image_ivar",
+ data=np.broadcast_to(
+ (row_values + 10)[:, None, None, None],
+ (source_count, 4, 160, 160),
+ ).copy(),
+ )
+ legacy.create_dataset(
+ "image_mask", data=np.zeros((source_count, 160, 160), dtype=np.bool_)
+ )
+ legacy.create_dataset(
+ "image_scale",
+ data=np.broadcast_to(row_values[:, None], (source_count, 4)).copy(),
+ )
+ legacy.create_dataset(
+ "object_mask", data=np.zeros((source_count, 160, 160), dtype=np.uint8)
+ )
+
+
+def _write_roster(path: Path) -> None:
+ rows = [
+ {
+ "paired_targetid": str(_SOURCE_IDS[source_row]),
+ "caption_object_id": _LEGACY_IDS[source_row],
+ "paired_ra": str(_RAS[source_row]),
+ "paired_dec": str(_DECS[source_row]),
+ "caption_ra": str(_RAS[source_row]),
+ "caption_dec": str(_DECS[source_row]),
+ }
+ for source_row in _ROSTER_SOURCE_ROWS
+ ]
+ with gzip.open(path, "wt", encoding="utf-8", newline="") as handle:
+ writer = csv.DictWriter(handle, fieldnames=_ROSTER_FIELDS)
+ writer.writeheader()
+ writer.writerows(rows)
+
+
+def _build_case(tmp_path: Path) -> RawSpectrumConfig:
+ source = tmp_path / "source.h5"
+ roster = tmp_path / "roster.csv.gz"
+ _write_source(source)
+ _write_roster(roster)
+ return RawSpectrumConfig(
+ roster_path=roster,
+ output_dir=tmp_path / "output",
+ source=source,
+ expected_count=len(_ROSTER_SOURCE_ROWS),
+ expected_roster_sha256=sha256_file(roster),
+ expected_source_length=source.stat().st_size,
+ expected_source_etag=None,
+ caption_legacy_coordinate_tolerance_arcsec=0.025,
+ )
+
+
+def test_extracts_mmap_spectra_in_roster_order_with_manifest_contract(
+ tmp_path: Path,
+) -> None:
+ config = _build_case(tmp_path)
+
+ summary = extract_raw_spectra(config)
+
+ assert summary["status"] == "prepared"
+ assert summary["candidate_count"] == 3
+ assert summary["contiguous_run_count"] == 2
+ assert summary["spectrum_length"] == ASTROPRISM_SPECTRUM_LENGTH
+
+ manifest = json.loads(
+ (config.output_dir / "manifest.json").read_text(encoding="utf-8")
+ )
+ assert manifest["artifact_type"] == "astroprism_raw_spectra"
+ assert manifest["roster"]["sha256"] == sha256_file(config.roster_path)
+ assert manifest["selection"] == {
+ "candidate_count": 3,
+ "contiguous_run_count": 2,
+ "http_cache_type": None,
+ "raw_row_max": 3,
+ "raw_row_min": 0,
+ "raw_rows_in_roster_order": [3, 0, 1],
+ "roster_order_preserved": True,
+ }
+ assert manifest["spectrum"] == {
+ "aion_modality": "DESISpectrum",
+ "aion_token_count": 273,
+ "length": ASTROPRISM_SPECTRUM_LENGTH,
+ "mask_true_means_excluded": True,
+ "observed_frame": True,
+ }
+ assert manifest["images_included"] is False
+ assert [record["paired_targetid"] for record in manifest["candidates"]] == [
+ "104",
+ "101",
+ "102",
+ ]
+ assert [record["raw_row"] for record in manifest["candidates"]] == [3, 0, 1]
+
+ expected = {
+ "spectrum_flux": ("spectrum_flux.npy", np.dtype("float32")),
+ "spectrum_ivar": ("spectrum_ivar.npy", np.dtype("float32")),
+ "spectrum_mask": ("spectrum_mask.npy", np.dtype("bool")),
+ "spectrum_wavelength": ("spectrum_wavelength.npy", np.dtype("float32")),
+ }
+ assert set(manifest["outputs"]) == set(expected)
+ assert not any("image" in value["path"] for value in manifest["outputs"].values())
+ for key, (filename, dtype) in expected.items():
+ metadata = manifest["outputs"][key]
+ output_path = config.output_dir / filename
+ array = np.load(output_path, mmap_mode="r", allow_pickle=False)
+ assert isinstance(array, np.memmap)
+ assert array.shape == (3, ASTROPRISM_SPECTRUM_LENGTH)
+ assert array.dtype == dtype
+ assert metadata["shape"] == [3, ASTROPRISM_SPECTRUM_LENGTH]
+ assert np.dtype(metadata["dtype"]) == dtype
+ assert metadata["mmap_safe"] is True
+ assert metadata["size_bytes"] == output_path.stat().st_size
+ assert metadata["sha256"] == sha256_file(output_path)
+
+ flux = np.load(config.output_dir / "spectrum_flux.npy", mmap_mode="r")
+ ivar = np.load(config.output_dir / "spectrum_ivar.npy", mmap_mode="r")
+ assert flux[:, 0].tolist() == [300.0, 0.0, 100.0]
+ assert ivar[:, 0].tolist() == [4.0, 1.0, 2.0]
+
+
+@pytest.mark.parametrize(
+ ("mutate", "message"),
+ [
+ (
+ lambda handle: handle["desi/spectrum_flux"].__setitem__((3, 12), np.nan),
+ "spectrum_flux contains non-finite values",
+ ),
+ (
+ lambda handle: handle["desi/spectrum_ivar"].__setitem__((3, 12), -1.0),
+ "spectrum_ivar contains negative values",
+ ),
+ (
+ lambda handle: handle["desi/spectrum_lambda"].__setitem__(
+ (3, 12), handle["desi/spectrum_lambda"][3, 11]
+ ),
+ "spectrum_lambda must be strictly increasing",
+ ),
+ (
+ lambda handle: handle["desi/spectrum_mask"].__setitem__(3, True),
+ "spectrum_mask marks every wavelength as invalid",
+ ),
+ ],
+)
+def test_rejects_invalid_selected_spectrum_and_cleans_staging(
+ tmp_path: Path,
+ mutate: Callable[[h5py.File], None],
+ message: str,
+) -> None:
+ config = _build_case(tmp_path)
+ with h5py.File(config.source, "r+") as handle:
+ mutate(handle)
+
+ with pytest.raises(RawSpectrumValidationError, match=message):
+ extract_raw_spectra(config)
+
+ assert not config.output_dir.exists()
+ assert not list(tmp_path.glob(".output.tmp-*"))
+
+
+def test_cli_local_source_does_not_apply_official_http_pins(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+ capsys: pytest.CaptureFixture[str],
+) -> None:
+ source = tmp_path / "source.h5"
+ args = argparse.Namespace(
+ roster=tmp_path / "roster.csv.gz",
+ source=str(source),
+ output=tmp_path / "output",
+ expected_source_length=None,
+ expected_source_etag=None,
+ http_block_size=4096,
+ http_timeout_seconds=7.5,
+ http_probe_attempts=2,
+ http_probe_backoff_seconds=0.25,
+ caption_legacy_coordinate_tolerance_arcsec=0.025,
+ overwrite=True,
+ )
+ captured: dict[str, Any] = {}
+
+ def fake_extract(config: RawSpectrumConfig) -> dict[str, Any]:
+ captured["config"] = config
+ return {"status": "prepared", "candidate_count": 3}
+
+ monkeypatch.setattr(spectrum_cli, "parse_args", lambda: args)
+ monkeypatch.setattr(spectrum_cli, "extract_raw_spectra", fake_extract)
+
+ assert spectrum_cli.main() == 0
+
+ config = captured["config"]
+ assert isinstance(config, RawSpectrumConfig)
+ assert config.source == str(source)
+ assert config.expected_source_length is None
+ assert config.expected_source_etag is None
+ assert config.http_block_size == 4096
+ assert config.overwrite is True
+ assert json.loads(capsys.readouterr().out) == {
+ "candidate_count": 3,
+ "status": "prepared",
+ }
diff --git a/tests/test_astroprism_spectrum_facts.py b/tests/test_astroprism_spectrum_facts.py
new file mode 100644
index 0000000..1b9cf64
--- /dev/null
+++ b/tests/test_astroprism_spectrum_facts.py
@@ -0,0 +1,161 @@
+from __future__ import annotations
+
+import numpy as np
+import pytest
+from src.data.astroprism_spectrum_facts import (
+ SpectralLineDefinition,
+ _robust_weighted_linear_fit,
+ measure_line,
+ measure_spectrum_facts,
+)
+
+
+def _line_definition(nature: str) -> SpectralLineDefinition:
+ return SpectralLineDefinition(
+ key=f"synthetic_{nature}",
+ label=f"synthetic {nature}",
+ rest_wavelength_angstrom=5000.0,
+ nature=nature, # type: ignore[arg-type]
+ line_window_rest_angstrom=(4997.5, 5002.5),
+ continuum_windows_rest_angstrom=((4970.0, 4990.0), (5010.0, 5030.0)),
+ )
+
+
+def test_robust_fit_refits_after_last_clipping_update() -> None:
+ x = np.arange(20, dtype=np.float64)
+ y = 2.0 + 3.0 * x
+ y[10] += 100.0
+ weight = np.ones_like(x)
+
+ coefficients, _, _, keep = _robust_weighted_linear_fit(
+ x, y, weight, iterations=1, clip_sigma=2.0
+ )
+
+ assert not keep[10]
+ expected, _, _, _ = np.linalg.lstsq(
+ np.column_stack((np.ones(np.count_nonzero(keep)), x[keep])), y[keep], rcond=None
+ )
+ np.testing.assert_allclose(coefficients, expected, rtol=1e-12, atol=1e-12)
+
+
+@pytest.mark.parametrize(
+ ("nature", "line_delta", "expected_sign"),
+ [("emission", 5.0, 1), ("absorption", -5.0, -1)],
+)
+def test_measure_line_uses_signed_rest_equivalent_width(
+ nature: str, line_delta: float, expected_sign: int
+) -> None:
+ wavelength = np.arange(4950.0, 5051.0)
+ flux = np.full(wavelength.shape, 10.0)
+ in_line = (wavelength >= 4998.0) & (wavelength <= 5002.0)
+ flux[in_line] += line_delta
+ ivar = np.full(wavelength.shape, 100.0)
+ mask = np.zeros(wavelength.shape, dtype=bool)
+
+ measured = measure_line(
+ wavelength,
+ flux,
+ ivar,
+ mask,
+ catalog_redshift=0.0,
+ definition=_line_definition(nature),
+ )
+
+ assert measured["detection_class"] == "detected"
+ assert np.sign(measured["equivalent_width_signed_rest_angstrom"]) == expected_sign
+ assert abs(measured["equivalent_width_signed_rest_angstrom"]) == pytest.approx(2.5, rel=1e-6)
+ assert measured["significance_sigma"] > 5.0
+
+
+def test_measure_line_marks_masked_feature_for_review() -> None:
+ wavelength = np.arange(4950.0, 5051.0)
+ flux = np.full(wavelength.shape, 10.0)
+ in_line = (wavelength >= 4998.0) & (wavelength <= 5002.0)
+ flux[in_line] += 10.0
+ ivar = np.full(wavelength.shape, 100.0)
+ mask = np.zeros(wavelength.shape, dtype=bool)
+ mask[np.flatnonzero(wavelength == 5000.0)] = True
+
+ measured = measure_line(
+ wavelength,
+ flux,
+ ivar,
+ mask,
+ catalog_redshift=0.0,
+ definition=_line_definition("emission"),
+ )
+
+ assert measured["detection_class"] == "detected"
+ assert measured["quality"] == "low_confidence"
+ assert measured["masked_line_pixels"] == 1
+ assert "masked_pixels_in_line_window" in measured["quality_flags"]
+
+
+def test_measure_line_rejects_partially_truncated_continuum_windows() -> None:
+ wavelength = np.arange(4980.0, 5041.0)
+ flux = np.full(wavelength.shape, 10.0)
+ ivar = np.full(wavelength.shape, 100.0)
+ mask = np.zeros(wavelength.shape, dtype=bool)
+
+ measured = measure_line(
+ wavelength,
+ flux,
+ ivar,
+ mask,
+ catalog_redshift=0.0,
+ definition=_line_definition("emission"),
+ )
+
+ assert measured["status"] == "out_of_coverage"
+ assert measured["detection_class"] == "unavailable"
+
+
+def test_measure_spectrum_facts_uses_narrow_fnu_dn4000() -> None:
+ wavelength = np.arange(3600.0, 7201.0)
+ # A flat F_nu spectrum has f_lambda proportional to lambda^-2. Increase
+ # F_nu by 50% only in the red Dn4000 band.
+ flux = 1.0e8 / wavelength**2
+ red_band = (wavelength >= 4000.0) & (wavelength <= 4100.0)
+ flux[red_band] *= 1.5
+ ivar = np.full(wavelength.shape, 1.0e4)
+ mask = np.zeros(wavelength.shape, dtype=bool)
+
+ measured = measure_spectrum_facts(wavelength, flux, ivar, mask, catalog_redshift=0.0)
+
+ assert measured["dn4000"]["status"] == "measured"
+ assert measured["dn4000"]["value"] == pytest.approx(1.5, rel=1e-12)
+ assert measured["dn4000"]["strength_label"] == "moderate"
+
+
+def test_dn4000_rejects_partially_truncated_band() -> None:
+ wavelength = np.arange(3900.0, 7201.0)
+ flux = 1.0e8 / wavelength**2
+ ivar = np.full(wavelength.shape, 1.0e4)
+ mask = np.zeros(wavelength.shape, dtype=bool)
+
+ measured = measure_spectrum_facts(wavelength, flux, ivar, mask, catalog_redshift=0.0)
+
+ assert measured["dn4000"]["status"] == "unavailable"
+ assert np.isnan(measured["dn4000"]["value"])
+
+
+def test_continuum_rejects_partially_truncated_band() -> None:
+ wavelength = np.arange(3600.0, 6301.0)
+ flux = 1.0e8 / wavelength**2
+ ivar = np.full(wavelength.shape, 1.0e4)
+ mask = np.zeros(wavelength.shape, dtype=bool)
+
+ measured = measure_spectrum_facts(wavelength, flux, ivar, mask, catalog_redshift=0.0)
+
+ assert measured["continuum"]["status"] == "unavailable"
+ assert np.isnan(measured["continuum"]["red_to_blue_ratio"])
+
+
+def test_measure_spectrum_facts_rejects_nonmonotonic_wavelength() -> None:
+ wavelength = np.array([4000.0, 4001.0, 4001.0, 4002.0])
+ flux = np.ones(4)
+ ivar = np.ones(4)
+ mask = np.zeros(4, dtype=bool)
+
+ with pytest.raises(ValueError, match="strictly increasing"):
+ measure_spectrum_facts(wavelength, flux, ivar, mask, catalog_redshift=0.1)
diff --git a/tests/test_astroprism_stage1a.py b/tests/test_astroprism_stage1a.py
new file mode 100644
index 0000000..f2736f1
--- /dev/null
+++ b/tests/test_astroprism_stage1a.py
@@ -0,0 +1,268 @@
+from __future__ import annotations
+
+import csv
+import gzip
+import json
+from dataclasses import replace
+from pathlib import Path
+
+import numpy as np
+import pyarrow as pa
+import pyarrow.parquet as pq
+import pytest
+from src.data.astroprism_stage1a import (
+ IMAGE_EMBEDDING_SHAPE,
+ Stage1AConfig,
+ Stage1AValidationError,
+ prepare_stage1a,
+ sha256_file,
+)
+
+_ROSTER_FIELDS = [
+ "paired_targetid",
+ "caption_object_id",
+ "paired_ra",
+ "paired_dec",
+ "caption_ra",
+ "caption_dec",
+ "separation_arcsec",
+ "paired_candidate_count",
+ "caption_candidate_count",
+ "one_to_one_row_candidate",
+ "unique_source_ids",
+ "provisional_unambiguous_candidate",
+ "paired_shard",
+ "paired_row",
+ "caption_shard",
+ "caption_row",
+]
+_PAIRED_REVISION = "paired-test-revision"
+_CAPTION_REVISION = "caption-test-revision"
+
+
+def _fixed_list_array(values: np.ndarray, dtype: pa.DataType) -> pa.Array:
+ rows, outer_size, inner_size = values.shape
+ flat = pa.array(values.reshape(-1), type=dtype)
+ inner = pa.FixedSizeListArray.from_arrays(flat, inner_size)
+ outer = pa.FixedSizeListArray.from_arrays(inner, outer_size)
+ assert len(outer) == rows
+ return outer
+
+
+def _write_metadata(root: Path, shard: str, revision: str) -> None:
+ source = root / shard
+ metadata = root / ".cache" / "huggingface" / "download" / f"{shard}.metadata"
+ metadata.parent.mkdir(parents=True, exist_ok=True)
+ metadata.write_text(f"{revision}\n{sha256_file(source)}\n0\n", encoding="utf-8")
+
+
+def _build_case(
+ tmp_path: Path,
+ *,
+ paired_ids: list[str] | None = None,
+ surveys: list[str] | None = None,
+ caption_ras: list[float] | None = None,
+ embedding_dtype: np.dtype = np.dtype("float16"), # noqa: B008
+ embedding_shape: tuple[int, int] = IMAGE_EMBEDDING_SHAPE,
+ nonfinite: bool = False,
+) -> Stage1AConfig:
+ paired_root = tmp_path / "paired"
+ caption_root = tmp_path / "captions"
+ paired_shard = "data/paired.parquet"
+ caption_shard = "data/captions.parquet"
+ (paired_root / "data").mkdir(parents=True)
+ (caption_root / "data").mkdir(parents=True)
+
+ canonical_paired_ids = ["9007199254740993", "9007199254740995"]
+ canonical_caption_ids = ["obj-a", "obj-b"]
+ ras = [10.0, 30.0]
+ decs = [20.0, -10.0]
+ source_caption_ras = caption_ras or ras
+
+ embeddings = np.empty((2, *embedding_shape), dtype=embedding_dtype)
+ embeddings[0].fill(1.0)
+ embeddings[1].fill(2.0)
+ if nonfinite:
+ embeddings[1, 0, 0] = np.nan
+ arrow_dtype = pa.from_numpy_dtype(embedding_dtype)
+ paired_table = pa.table(
+ {
+ "object_id_provabgs": paired_ids or canonical_paired_ids,
+ "ra_provabgs": pa.array(ras, type=pa.float64()),
+ "dec_provabgs": pa.array(decs, type=pa.float64()),
+ "embedding_image": _fixed_list_array(embeddings, arrow_dtype),
+ # Deliberately present: extraction must not copy this column.
+ "embedding_spectrum": pa.array([[901.0], [902.0]], type=pa.list_(pa.float16())),
+ }
+ )
+ pq.write_table(paired_table, paired_root / paired_shard)
+
+ caption_table = pa.table(
+ {
+ "object_id": canonical_caption_ids,
+ "survey": surveys or ["legacy_dr10", "legacy_dr10"],
+ "ra": pa.array(source_caption_ras, type=pa.float64()),
+ "dec": pa.array(decs, type=pa.float64()),
+ "caption": [" original caption A\n", "original caption B"],
+ "summary": ["summary A", "summary B"],
+ }
+ )
+ pq.write_table(caption_table, caption_root / caption_shard)
+ _write_metadata(paired_root, paired_shard, _PAIRED_REVISION)
+ _write_metadata(caption_root, caption_shard, _CAPTION_REVISION)
+
+ # Roster order intentionally reverses both source tables.
+ roster_rows = []
+ for source_row in (1, 0):
+ roster_rows.append(
+ {
+ "paired_targetid": canonical_paired_ids[source_row],
+ "caption_object_id": canonical_caption_ids[source_row],
+ "paired_ra": str(ras[source_row]),
+ "paired_dec": str(decs[source_row]),
+ "caption_ra": str(ras[source_row]),
+ "caption_dec": str(decs[source_row]),
+ "separation_arcsec": "0.0",
+ "paired_candidate_count": "1",
+ "caption_candidate_count": "1",
+ "one_to_one_row_candidate": "True",
+ "unique_source_ids": "True",
+ "provisional_unambiguous_candidate": "True",
+ "paired_shard": paired_shard,
+ "paired_row": str(source_row),
+ "caption_shard": caption_shard,
+ "caption_row": str(source_row),
+ }
+ )
+ roster = tmp_path / "roster.csv.gz"
+ with gzip.open(roster, "wt", encoding="utf-8", newline="") as handle:
+ writer = csv.DictWriter(handle, fieldnames=_ROSTER_FIELDS)
+ writer.writeheader()
+ writer.writerows(roster_rows)
+
+ return Stage1AConfig(
+ roster_path=roster,
+ paired_root=paired_root,
+ caption_root=caption_root,
+ output_dir=tmp_path / "output",
+ expected_count=2,
+ expected_roster_sha256=sha256_file(roster),
+ paired_revision=_PAIRED_REVISION,
+ caption_revision=_CAPTION_REVISION,
+ batch_size=1,
+ )
+
+
+def _read_jsonl(path: Path) -> list[dict]:
+ return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines()]
+
+
+def test_extracts_image_only_in_deterministic_roster_order(tmp_path: Path):
+ config = _build_case(tmp_path)
+
+ summary = prepare_stage1a(config)
+
+ assert summary["status"] == "prepared"
+ embeddings = np.load(config.output_dir / "image_embeddings.f16.npy", mmap_mode="r")
+ assert embeddings.shape == (2, *IMAGE_EMBEDDING_SHAPE)
+ assert embeddings.dtype == np.float16
+ assert np.all(embeddings[0] == np.float16(2.0))
+ assert np.all(embeddings[1] == np.float16(1.0))
+
+ records = _read_jsonl(config.output_dir / "candidates.jsonl")
+ assert [record["candidate_index"] for record in records] == [0, 1]
+ assert [record["paired_targetid"] for record in records] == [
+ "9007199254740995",
+ "9007199254740993",
+ ]
+ assert records[1]["caption"]["original"] == " original caption A\n"
+ assert all(record["acceptance_status"] == "pending_review" for record in records)
+ assert all(record["source_group_id"] is None for record in records)
+ assert all(record["split"] is None for record in records)
+ assert "embedding_spectrum" not in json.dumps(records)
+
+ manifest = json.loads((config.output_dir / "manifest.json").read_text())
+ assert manifest["artifact_type"] == "astroprism_aion_feature_parity_reference"
+ assert manifest["role"] == "parity_reference_only"
+ assert manifest["spectra_included"] is False
+ assert manifest["outputs"]["image_embeddings"]["sha256"] == sha256_file(
+ config.output_dir / "image_embeddings.f16.npy"
+ )
+ assert manifest["sources"]["paired"][0]["sha256"] == sha256_file(
+ config.paired_root / "data/paired.parquet"
+ )
+
+
+def test_validate_only_writes_nothing(tmp_path: Path):
+ config = replace(_build_case(tmp_path), validate_only=True)
+
+ summary = prepare_stage1a(config)
+
+ assert summary["status"] == "validated"
+ assert summary["spectra_included"] is False
+ assert not config.output_dir.exists()
+
+
+@pytest.mark.parametrize(
+ ("case_kwargs", "message"),
+ [
+ ({"paired_ids": ["9007199254740993", "wrong-id"]}, "paired ID mismatch"),
+ ({"surveys": ["legacy_dr10", "hsc"]}, "survey must be"),
+ ({"caption_ras": [10.0, 30.01]}, "caption_ra mismatch"),
+ ],
+)
+def test_rejects_identity_survey_and_coordinate_mismatches(
+ tmp_path: Path, case_kwargs: dict, message: str
+):
+ config = _build_case(tmp_path, **case_kwargs)
+
+ with pytest.raises(Stage1AValidationError, match=message):
+ prepare_stage1a(config)
+
+ assert not config.output_dir.exists()
+ assert not list(tmp_path.glob(".output.tmp-*"))
+
+
+@pytest.mark.parametrize(
+ "case_kwargs",
+ [
+ {"embedding_dtype": np.dtype("float32")},
+ {"embedding_shape": (575, 768)},
+ ],
+)
+def test_rejects_wrong_embedding_schema(tmp_path: Path, case_kwargs: dict):
+ config = _build_case(tmp_path, **case_kwargs)
+
+ with pytest.raises(Stage1AValidationError, match="embedding_image must be"):
+ prepare_stage1a(config)
+
+
+def test_rejects_nonfinite_embedding(tmp_path: Path):
+ config = _build_case(tmp_path, nonfinite=True)
+
+ with pytest.raises(Stage1AValidationError, match="non-finite"):
+ prepare_stage1a(config)
+
+
+def test_source_hash_is_verified(tmp_path: Path):
+ config = _build_case(tmp_path)
+ with (config.paired_root / "data/paired.parquet").open("ab") as handle:
+ handle.write(b"tampered")
+
+ with pytest.raises(Stage1AValidationError, match="source SHA-256 mismatch"):
+ prepare_stage1a(config)
+
+
+def test_existing_output_requires_overwrite(tmp_path: Path):
+ config = _build_case(tmp_path)
+ prepare_stage1a(config)
+ stale = config.output_dir / "stale"
+ stale.write_text("old", encoding="utf-8")
+
+ with pytest.raises(FileExistsError, match="output already exists"):
+ prepare_stage1a(config)
+ assert stale.exists()
+
+ prepare_stage1a(replace(config, overwrite=True))
+ assert not stale.exists()
+ assert (config.output_dir / "manifest.json").is_file()
diff --git a/tests/test_astroprism_stage1a_dataset.py b/tests/test_astroprism_stage1a_dataset.py
new file mode 100644
index 0000000..d1868c2
--- /dev/null
+++ b/tests/test_astroprism_stage1a_dataset.py
@@ -0,0 +1,387 @@
+from __future__ import annotations
+
+import json
+from dataclasses import replace
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import pytest
+import torch
+from src.data.astroprism_stage1a_collate import AstroPRISMStage1ACollator
+from src.data.astroprism_stage1a_dataset import (
+ ASTROPRISM_IMAGE_TOKENS,
+ IGNORE_INDEX,
+ AstroPRISMStage1ADataset,
+ AstroPRISMStage1ADatasetConfig,
+ AstroPRISMStage1ADatasetError,
+)
+
+pytestmark = [pytest.mark.unit, pytest.mark.multimodal]
+
+
+class _FakeQwenTokenizer:
+ pad_token_id = 0
+
+ def __init__(self) -> None:
+ self.calls: list[dict[str, Any]] = []
+
+ @staticmethod
+ def _content_ids(content: str) -> list[int]:
+ return [100 + (sum(word.encode("utf-8")) % 500) for word in content.split()]
+
+ def apply_chat_template(
+ self,
+ messages: list[dict[str, str]],
+ *,
+ tokenize: bool,
+ add_generation_prompt: bool,
+ enable_thinking: bool,
+ ) -> list[int]:
+ self.calls.append(
+ {
+ "messages": messages,
+ "tokenize": tokenize,
+ "add_generation_prompt": add_generation_prompt,
+ "enable_thinking": enable_thinking,
+ }
+ )
+ assert tokenize is True
+ role_ids = {"system": 10, "user": 11, "assistant": 12}
+ ids = [1]
+ for message in messages:
+ ids.append(role_ids[message["role"]])
+ ids.extend(self._content_ids(message["content"]))
+ ids.append(2)
+ if add_generation_prompt:
+ ids.append(role_ids["assistant"])
+ return ids
+
+
+def _candidate(
+ index: int,
+ *,
+ status: str = "accepted",
+ caption_status: str = "released_reused",
+ split: str | None = "train",
+ paired_targetid: str | None = None,
+ caption: str | None = None,
+ revised: str | None = None,
+) -> dict[str, Any]:
+ paired_id = paired_targetid or f"paired-{index}"
+ caption_payload = {
+ "original": caption or f"Galaxy {index} has a smooth central bulge.",
+ "summary_original": f"Summary {index}.",
+ "status": caption_status,
+ }
+ if revised is not None:
+ caption_payload["revised"] = revised
+ return {
+ "schema_version": 1,
+ "candidate_index": index,
+ "paired_targetid": paired_id,
+ "caption_object_id": f"caption-{index}",
+ "source_group_id": f"source-{index}",
+ "acceptance_status": status,
+ "split": split,
+ "caption": caption_payload,
+ "caption_source": {"repo_id": "captions", "row": index},
+ "paired_source": {"repo_id": "paired", "row": index},
+ }
+
+
+def _write_case(
+ tmp_path: Path,
+ records: list[dict[str, Any]] | None = None,
+) -> AstroPRISMStage1ADatasetConfig:
+ records = records or [
+ _candidate(0, caption="Compact galaxy."),
+ _candidate(1, caption="Extended galaxy with a smooth central bulge and diffuse disk."),
+ ]
+ raw_root = tmp_path / "raw"
+ candidate_root = tmp_path / "candidates"
+ raw_root.mkdir()
+ candidate_root.mkdir()
+
+ images = np.empty((len(records), 4, 160, 160), dtype=np.float32)
+ for index in range(len(records)):
+ images[index].fill(index + 0.25)
+ np.save(raw_root / "image_array.npy", images, allow_pickle=False)
+
+ raw_candidates = []
+ for index, record in enumerate(records):
+ raw_candidates.append(
+ {
+ "candidate_index": index,
+ "paired_targetid": record["paired_targetid"],
+ "caption_object_id": record["caption_object_id"],
+ "acceptance_status": "pending_review",
+ "split": None,
+ "raw_row": 1000 + index,
+ "bands": {
+ "source": ["des-g", "des-r", "des-i", "des-z"],
+ "aion": ["DES-G", "DES-R", "DES-I", "DES-Z"],
+ },
+ "coordinates_deg": {"provabgs": {"ra": 10.0 + index, "dec": -2.0}},
+ }
+ )
+ raw_manifest = {
+ "artifact_type": "astroprism_raw_images",
+ "schema_version": 1,
+ "roster": {"sha256": "same-roster"},
+ "source": {"location": "fixture.hdf5"},
+ "outputs": {
+ "image_array": {
+ "path": "image_array.npy",
+ "dtype": "float32",
+ "shape": list(images.shape),
+ }
+ },
+ "candidates": raw_candidates,
+ }
+ (raw_root / "manifest.json").write_text(json.dumps(raw_manifest), encoding="utf-8")
+
+ candidate_path = candidate_root / "candidates.jsonl"
+ candidate_path.write_text(
+ "".join(json.dumps(record) + "\n" for record in records), encoding="utf-8"
+ )
+ candidate_manifest = {
+ "artifact_type": "astroprism_stage1a_candidates",
+ "schema_version": 1,
+ "pins": {"roster_sha256": "same-roster"},
+ "outputs": {
+ "candidates": {
+ "path": "candidates.jsonl",
+ "row_count": len(records),
+ }
+ },
+ }
+ (candidate_root / "manifest.json").write_text(json.dumps(candidate_manifest), encoding="utf-8")
+
+ return AstroPRISMStage1ADatasetConfig(
+ raw_manifest_path=raw_root / "manifest.json",
+ candidates_manifest_path=candidate_root / "manifest.json",
+ split="train",
+ max_merged_seq_length=1024,
+ )
+
+
+def test_emits_raw_image_qwen_chat_assistant_labels_and_provenance(tmp_path: Path) -> None:
+ tokenizer = _FakeQwenTokenizer()
+ dataset = AstroPRISMStage1ADataset(_write_case(tmp_path), tokenizer)
+
+ sample = dataset[0]
+
+ assert len(dataset) == 2
+ assert dataset.active_modalities == {"image"}
+ assert sample["image"].shape == (4, 160, 160)
+ assert sample["image"].dtype == torch.float32
+ assert torch.all(sample["image"] == torch.tensor(0.25))
+ assert sample["text"].dtype == torch.long
+ assert sample["labels"].shape == sample["text"].shape
+ assert torch.equal(sample["text_attention_mask"], torch.ones_like(sample["text"]))
+ first_target = int(torch.nonzero(sample["labels"] != IGNORE_INDEX)[0].item())
+ assert torch.all(sample["labels"][:first_target] == IGNORE_INDEX)
+ assert torch.equal(sample["labels"][first_target:], sample["text"][first_target:])
+ assert sample["image_token_count"] == ASTROPRISM_IMAGE_TOKENS
+ assert sample["merged_length"] == ASTROPRISM_IMAGE_TOKENS + sample["text"].numel()
+ assert sample["provenance"]["paired_targetid"] == "paired-0"
+ assert sample["provenance"]["raw_row"] == 1000
+ assert sample["provenance"]["bands"]["aion"] == [
+ "DES-G",
+ "DES-R",
+ "DES-I",
+ "DES-Z",
+ ]
+ assert sample["provenance"]["caption_field"] == "original"
+ assert sample["provenance"]["caption_field_config"] == "auto"
+ assert sample["provenance"]["caption_status"] == "released_reused"
+ assert len(tokenizer.calls) == 2
+ assert all(call["enable_thinking"] is False for call in tokenizer.calls)
+ assert tokenizer.calls[0]["add_generation_prompt"] is True
+ assert tokenizer.calls[1]["add_generation_prompt"] is False
+
+
+def test_auto_selects_revised_text_for_released_revised_caption(tmp_path: Path) -> None:
+ tokenizer = _FakeQwenTokenizer()
+ record = _candidate(
+ 0,
+ caption_status="released_revised",
+ caption="Original text not approved for this crop.",
+ revised="Reviewed text supported by the exact crop.",
+ )
+ dataset = AstroPRISMStage1ADataset(_write_case(tmp_path, [record]), tokenizer)
+
+ sample = dataset[0]
+
+ assert tokenizer.calls[-1]["messages"][-1]["content"] == record["caption"]["revised"]
+ assert sample["provenance"]["caption_field"] == "revised"
+ assert sample["provenance"]["caption_status"] == "released_revised"
+
+
+@pytest.mark.parametrize(
+ ("config_override", "message"),
+ [
+ ({}, "no non-empty caption.revised"),
+ ({"caption_field": "original"}, "requires caption.revised"),
+ (
+ {"caption_field": "original", "allow_unsafe_candidates": True},
+ "requires caption.revised",
+ ),
+ ({"allow_unsafe_candidates": True}, "no non-empty caption.revised"),
+ ],
+)
+def test_released_revised_never_silently_uses_original(
+ tmp_path: Path,
+ config_override: dict[str, Any],
+ message: str,
+) -> None:
+ record = _candidate(
+ 0,
+ caption_status="released_revised",
+ caption="Original text is not crop-reviewed.",
+ )
+ config = replace(_write_case(tmp_path, [record]), **config_override)
+
+ with pytest.raises(AstroPRISMStage1ADatasetError, match=message):
+ AstroPRISMStage1ADataset(config, _FakeQwenTokenizer())
+
+
+@pytest.mark.parametrize(
+ ("record", "message"),
+ [
+ (_candidate(0, status="pending_review", split="train"), "acceptance_status"),
+ (_candidate(0, status="accepted", split=None), "split=null"),
+ (_candidate(0, caption_status="unreviewed", split="train"), "caption.status"),
+ (_candidate(0, caption_status="rejected", split="train"), "caption.status"),
+ ],
+)
+def test_refuses_unreviewed_or_unsplit_candidates_by_default(
+ tmp_path: Path, record: dict[str, Any], message: str
+) -> None:
+ config = replace(_write_case(tmp_path, [record]), split=None)
+
+ with pytest.raises(AstroPRISMStage1ADatasetError, match=message):
+ AstroPRISMStage1ADataset(config, _FakeQwenTokenizer())
+
+
+def test_explicit_unsafe_override_loads_pending_unsplit_fixture(tmp_path: Path) -> None:
+ config = replace(
+ _write_case(tmp_path, [_candidate(0, status="pending_review", split=None)]),
+ split=None,
+ allow_unsafe_candidates=True,
+ )
+
+ tokenizer = _FakeQwenTokenizer()
+ dataset = AstroPRISMStage1ADataset(config, tokenizer)
+ sample = dataset[0]
+
+ assert len(dataset) == 1
+ assert sample["provenance"]["acceptance_status"] == "pending_review"
+ assert sample["provenance"]["split"] is None
+ assert sample["provenance"]["caption_field"] == "original"
+ assert tokenizer.calls[-1]["messages"][-1]["content"] == (
+ "Galaxy 0 has a smooth central bulge."
+ )
+
+
+def test_split_view_ignores_explicitly_excluded_rows(tmp_path: Path) -> None:
+ records = [
+ _candidate(
+ 0,
+ split="train",
+ caption_status="released_revised",
+ revised="Reviewed revision.",
+ ),
+ _candidate(1, status="rejected", split=None, caption_status="rejected"),
+ ]
+ dataset = AstroPRISMStage1ADataset(_write_case(tmp_path, records), _FakeQwenTokenizer())
+
+ assert len(dataset) == 1
+ assert dataset[0]["provenance"]["candidate_index"] == 0
+
+
+def test_rejects_source_identity_across_splits(tmp_path: Path) -> None:
+ first = _candidate(0, split="train", paired_targetid="same-object")
+ second = _candidate(1, split="validation", paired_targetid="same-object")
+ config = _write_case(tmp_path, [first, second])
+
+ with pytest.raises(AstroPRISMStage1ADatasetError, match="source split leakage"):
+ AstroPRISMStage1ADataset(config, _FakeQwenTokenizer())
+
+
+def test_rejects_raw_caption_identity_mismatch(tmp_path: Path) -> None:
+ config = _write_case(tmp_path)
+ raw_manifest = json.loads(config.raw_manifest_path.read_text(encoding="utf-8"))
+ raw_manifest["candidates"][1]["caption_object_id"] = "wrong-object"
+ config.raw_manifest_path.write_text(json.dumps(raw_manifest), encoding="utf-8")
+
+ with pytest.raises(AstroPRISMStage1ADatasetError, match="caption_object_id mismatch"):
+ AstroPRISMStage1ADataset(config, _FakeQwenTokenizer())
+
+
+def test_rejects_unverified_raw_band_order(tmp_path: Path) -> None:
+ config = _write_case(tmp_path)
+ raw_manifest = json.loads(config.raw_manifest_path.read_text(encoding="utf-8"))
+ raw_manifest["candidates"][0]["bands"]["source"] = [
+ "des-r",
+ "des-g",
+ "des-i",
+ "des-z",
+ ]
+ config.raw_manifest_path.write_text(json.dumps(raw_manifest), encoding="utf-8")
+
+ with pytest.raises(AstroPRISMStage1ADatasetError, match="unsupported band order"):
+ AstroPRISMStage1ADataset(config, _FakeQwenTokenizer())
+
+
+def test_merged_length_guard_raises_instead_of_truncating(tmp_path: Path) -> None:
+ tokenizer = _FakeQwenTokenizer()
+ config = replace(_write_case(tmp_path), max_merged_seq_length=577)
+ dataset = AstroPRISMStage1ADataset(config, tokenizer)
+
+ with pytest.raises(AstroPRISMStage1ADatasetError, match="never truncates"):
+ dataset[0]
+
+ assert all("truncation" not in call for call in tokenizer.calls)
+
+
+def test_collator_pads_text_mask_and_labels_in_lockstep(tmp_path: Path) -> None:
+ tokenizer = _FakeQwenTokenizer()
+ config = _write_case(tmp_path)
+ dataset = AstroPRISMStage1ADataset(config, tokenizer)
+ samples = [dataset[0], dataset[1]]
+ original_lengths = [sample["text"].numel() for sample in samples]
+ collator = AstroPRISMStage1ACollator(
+ tokenizer, max_merged_seq_length=config.max_merged_seq_length
+ )
+
+ batch = collator(samples)
+
+ longest = max(original_lengths)
+ shortest_index = original_lengths.index(min(original_lengths))
+ shortest = original_lengths[shortest_index]
+ assert batch["image"].shape == (2, 4, 160, 160)
+ assert batch["text"].shape == batch["labels"].shape == (2, longest)
+ assert batch["text_attention_mask"].shape == (2, longest)
+ assert torch.all(batch["text"][shortest_index, shortest:] == tokenizer.pad_token_id)
+ assert torch.all(batch["text_attention_mask"][shortest_index, shortest:] == 0)
+ assert torch.all(batch["labels"][shortest_index, shortest:] == IGNORE_INDEX)
+ assert [sample["text"].numel() for sample in samples] == original_lengths
+ assert batch["image_token_count"] == ASTROPRISM_IMAGE_TOKENS
+ assert batch["provenance"][1]["candidate_index"] == 1
+
+
+def test_collator_refuses_oversized_sample_without_truncating(tmp_path: Path) -> None:
+ tokenizer = _FakeQwenTokenizer()
+ dataset = AstroPRISMStage1ADataset(_write_case(tmp_path), tokenizer)
+ sample = dataset[0]
+ collator = AstroPRISMStage1ACollator(
+ tokenizer,
+ max_merged_seq_length=sample["merged_length"] - 1,
+ )
+
+ with pytest.raises(ValueError, match="refusing to truncate"):
+ collator([sample])
+
+ assert sample["text"].numel() == sample["merged_length"] - ASTROPRISM_IMAGE_TOKENS
diff --git a/tests/test_astroprism_stage1b_spectrum_dataset.py b/tests/test_astroprism_stage1b_spectrum_dataset.py
new file mode 100644
index 0000000..d9ef613
--- /dev/null
+++ b/tests/test_astroprism_stage1b_spectrum_dataset.py
@@ -0,0 +1,505 @@
+from __future__ import annotations
+
+import hashlib
+import json
+from dataclasses import replace
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import pytest
+import torch
+from src.data.astroprism_spectrum_contract import ASTROPRISM_SPECTRUM_LENGTH
+from src.data.astroprism_stage1a_dataset import (
+ IGNORE_INDEX,
+ AstroPRISMStage1ADatasetError,
+)
+from src.data.astroprism_stage1b_spectrum_collate import (
+ AstroPRISMStage1BSpectrumCollator,
+)
+from src.data.astroprism_stage1b_spectrum_dataset import (
+ ASTROPRISM_SPECTRUM_TOKENS,
+ AstroPRISMStage1BSpectrumDataset,
+ AstroPRISMStage1BSpectrumDatasetConfig,
+)
+
+pytestmark = [pytest.mark.unit, pytest.mark.multimodal]
+
+
+class _FakeQwenTokenizer:
+ pad_token_id = 0
+
+ def __init__(self) -> None:
+ self.calls: list[dict[str, Any]] = []
+
+ @staticmethod
+ def _content_ids(content: str) -> list[int]:
+ return [100 + (sum(word.encode("utf-8")) % 500) for word in content.split()]
+
+ def apply_chat_template(
+ self,
+ messages: list[dict[str, str]],
+ *,
+ tokenize: bool,
+ add_generation_prompt: bool,
+ enable_thinking: bool,
+ ) -> list[int]:
+ self.calls.append(
+ {
+ "messages": messages,
+ "tokenize": tokenize,
+ "add_generation_prompt": add_generation_prompt,
+ "enable_thinking": enable_thinking,
+ }
+ )
+ assert tokenize is True
+ role_ids = {"system": 10, "user": 11, "assistant": 12}
+ ids = [1]
+ for message in messages:
+ ids.append(role_ids[message["role"]])
+ ids.extend(self._content_ids(message["content"]))
+ ids.append(2)
+ if add_generation_prompt:
+ ids.append(role_ids["assistant"])
+ return ids
+
+
+def _caption_record(index: int, *, caption: str, split: str) -> dict[str, Any]:
+ is_train = split == "train"
+ return {
+ "schema_version": 1,
+ "artifact_type": "astroprism_spectrum_caption_candidate",
+ "candidate_index": index,
+ "paired_targetid": f"paired-{index}",
+ "caption_object_id": f"caption-{index}",
+ "source_group_id": f"source-{index}",
+ "split": split,
+ "spectrum_caption": {
+ "original": caption,
+ "revised": None,
+ "status": "accepted",
+ },
+ "review": {
+ "status": "accepted",
+ "decision_source": "user_instruction",
+ },
+ "training_policy": {
+ "dataset_ready": True,
+ "gradient_training_allowed": is_train,
+ "eligible_for_training": is_train,
+ "eligible_for_evaluation": not is_train,
+ "role": f"alignment_{'training' if is_train else split}",
+ },
+ }
+
+
+def _sha256(path: Path) -> str:
+ return hashlib.sha256(path.read_bytes()).hexdigest()
+
+
+def _write_case(tmp_path: Path) -> AstroPRISMStage1BSpectrumDatasetConfig:
+ records = [
+ _caption_record(0, caption="Spectrum caption zero.", split="train"),
+ _caption_record(
+ 1,
+ caption="A longer reviewed spectrum caption for candidate one.",
+ split="train",
+ ),
+ _caption_record(2, caption="Validation spectrum caption.", split="validation"),
+ _caption_record(3, caption="Test spectrum caption.", split="test"),
+ ]
+ raw_root = tmp_path / "raw"
+ caption_root = tmp_path / "captions"
+ raw_root.mkdir()
+ caption_root.mkdir()
+
+ pixels = np.arange(ASTROPRISM_SPECTRUM_LENGTH, dtype=np.float32)
+ arrays = {
+ "spectrum_flux": np.stack([pixels / 100.0 + 0.25 + index for index in range(len(records))]),
+ "spectrum_ivar": np.stack(
+ [
+ np.full(
+ ASTROPRISM_SPECTRUM_LENGTH,
+ 2.0 + index,
+ dtype=np.float32,
+ )
+ for index in range(len(records))
+ ]
+ ),
+ "spectrum_mask": np.zeros((len(records), ASTROPRISM_SPECTRUM_LENGTH), dtype=np.bool_),
+ "spectrum_wavelength": np.broadcast_to(
+ 3500.0 + pixels,
+ (len(records), ASTROPRISM_SPECTRUM_LENGTH),
+ ).copy(),
+ }
+ for index in range(len(records)):
+ arrays["spectrum_mask"][index, :: 257 + index] = True
+ outputs: dict[str, dict[str, Any]] = {}
+ for key, array in arrays.items():
+ array_path = raw_root / f"{key}.npy"
+ np.save(array_path, array, allow_pickle=False)
+ outputs[key] = {
+ "path": array_path.name,
+ "dtype": str(array.dtype),
+ "shape": list(array.shape),
+ }
+
+ raw_candidates = [
+ {
+ "candidate_index": index,
+ "paired_targetid": record["paired_targetid"],
+ "caption_object_id": record["caption_object_id"],
+ "acceptance_status": "pending_review",
+ "split": None,
+ "raw_row": 1000 + index,
+ "coordinates_deg": {
+ "desi": {"ra": 10.0 + index, "dec": -2.0},
+ },
+ }
+ for index, record in enumerate(records)
+ ]
+ raw_manifest = {
+ "artifact_type": "astroprism_raw_spectra",
+ "schema_version": 1,
+ "roster": {"sha256": "same-roster", "row_count": len(records)},
+ "source": {"kind": "local", "location": "fixture.hdf5"},
+ "spectrum": {
+ "length": ASTROPRISM_SPECTRUM_LENGTH,
+ "aion_token_count": ASTROPRISM_SPECTRUM_TOKENS,
+ "observed_frame": True,
+ "mask_true_means_excluded": True,
+ "aion_modality": "DESISpectrum",
+ },
+ "outputs": outputs,
+ "candidates": raw_candidates,
+ }
+ raw_manifest_path = raw_root / "manifest.json"
+ raw_manifest_path.write_text(json.dumps(raw_manifest), encoding="utf-8")
+
+ captions_path = caption_root / "captions.jsonl"
+ captions_path.write_text(
+ "".join(json.dumps(record) + "\n" for record in records),
+ encoding="utf-8",
+ )
+ split_counts = {"train": 2, "validation": 1, "test": 1}
+ captions_manifest = {
+ "artifact_type": "astroprism_spectrum_captions_accepted",
+ "schema_version": 1,
+ "role": "reviewed_alignment_dataset",
+ "row_count": len(records),
+ "split_counts": split_counts,
+ "candidate_indices": list(range(len(records))),
+ "roster": {"sha256": "same-roster", "row_count": len(records)},
+ "review": {
+ "status": "accepted",
+ "decision_source": "user_instruction",
+ },
+ "training_policy": {
+ "dataset_ready": True,
+ "gradient_training_splits": ["train"],
+ "evaluation_splits": ["validation", "test"],
+ "roles": {
+ "train": "alignment_training",
+ "validation": "alignment_validation",
+ "test": "alignment_test",
+ },
+ },
+ "splits": {"row_counts": split_counts},
+ "outputs": {
+ "captions.jsonl": {
+ "size_bytes": captions_path.stat().st_size,
+ "sha256": _sha256(captions_path),
+ }
+ },
+ }
+ captions_manifest_path = caption_root / "manifest.json"
+ captions_manifest_path.write_text(json.dumps(captions_manifest), encoding="utf-8")
+
+ return AstroPRISMStage1BSpectrumDatasetConfig(
+ raw_manifest_path=raw_manifest_path,
+ captions_manifest_path=captions_manifest_path,
+ max_merged_seq_length=1024,
+ verify_file_hashes=True,
+ )
+
+
+def _rewrite_caption_records(
+ config: AstroPRISMStage1BSpectrumDatasetConfig,
+ mutation,
+ *,
+ refresh_manifest: bool = True,
+) -> None:
+ captions_path = config.captions_manifest_path.parent / "captions.jsonl"
+ records = [json.loads(line) for line in captions_path.read_text().splitlines()]
+ mutation(records)
+ captions_path.write_text(
+ "".join(json.dumps(record) + "\n" for record in records),
+ encoding="utf-8",
+ )
+ if refresh_manifest:
+ manifest = json.loads(config.captions_manifest_path.read_text(encoding="utf-8"))
+ output = manifest["outputs"]["captions.jsonl"]
+ output["size_bytes"] = captions_path.stat().st_size
+ output["sha256"] = _sha256(captions_path)
+ config.captions_manifest_path.write_text(json.dumps(manifest), encoding="utf-8")
+
+
+def _rewrite_caption_manifest(config, mutation) -> None:
+ path = config.captions_manifest_path
+ manifest = json.loads(path.read_text(encoding="utf-8"))
+ mutation(manifest)
+ path.write_text(json.dumps(manifest), encoding="utf-8")
+
+
+def _set_nested(record: dict[str, Any], keys: tuple[str, ...], value: Any) -> None:
+ current = record
+ for key in keys[:-1]:
+ current = current[key]
+ current[keys[-1]] = value
+
+
+def test_emits_reviewed_spectrum_caption_labels_and_grounded_provenance(
+ tmp_path: Path,
+) -> None:
+ tokenizer = _FakeQwenTokenizer()
+ dataset = AstroPRISMStage1BSpectrumDataset(_write_case(tmp_path), tokenizer)
+
+ sample = dataset[0]
+
+ assert len(dataset) == 2
+ assert dataset.active_modalities == {"spectrum"}
+ assert set(sample["spectrum"]) == {"flux", "ivar", "mask", "wavelength"}
+ assert {key: value.shape for key, value in sample["spectrum"].items()} == {
+ "flux": (ASTROPRISM_SPECTRUM_LENGTH,),
+ "ivar": (ASTROPRISM_SPECTRUM_LENGTH,),
+ "mask": (ASTROPRISM_SPECTRUM_LENGTH,),
+ "wavelength": (ASTROPRISM_SPECTRUM_LENGTH,),
+ }
+ assert {key: value.dtype for key, value in sample["spectrum"].items()} == {
+ "flux": torch.float32,
+ "ivar": torch.float32,
+ "mask": torch.bool,
+ "wavelength": torch.float32,
+ }
+ assert sample["spectrum"]["flux"][0].item() == pytest.approx(0.25)
+ assert sample["spectrum"]["ivar"][0].item() == pytest.approx(2.0)
+ assert sample["spectrum"]["mask"][0].item() is True
+ assert sample["spectrum"]["wavelength"][0].item() == pytest.approx(3500.0)
+
+ assert sample["text"].dtype == torch.long
+ assert sample["labels"].shape == sample["text"].shape
+ assert torch.equal(sample["text_attention_mask"], torch.ones_like(sample["text"]))
+ first_target = int(torch.nonzero(sample["labels"] != IGNORE_INDEX)[0].item())
+ assert torch.all(sample["labels"][:first_target] == IGNORE_INDEX)
+ assert torch.equal(sample["labels"][first_target:], sample["text"][first_target:])
+
+ assert ASTROPRISM_SPECTRUM_TOKENS == 273
+ assert sample["spectrum_token_count"] == 273
+ assert sample["merged_length"] == 273 + sample["text"].numel()
+ assert sample["provenance"]["raw_row"] == 1000
+ assert sample["provenance"]["raw_source"] == {
+ "kind": "local",
+ "location": "fixture.hdf5",
+ }
+ assert sample["provenance"]["roster_sha256"] == "same-roster"
+ assert sample["provenance"]["caption_field"] == "original"
+ assert sample["provenance"]["caption_status"] == "accepted"
+ assert sample["provenance"]["caption_review_status"] == "accepted"
+ assert sample["provenance"]["caption_review_decision_source"] == "user_instruction"
+ assert sample["provenance"]["caption_training_role"] == "alignment_training"
+ assert sample["provenance"]["caption_supervision"] == "reviewed_spectrum_caption"
+ assert sample["provenance"]["spectral_grounding_claimed"] is True
+ assert sample["provenance"]["chat_template"] == "qwen_non_thinking"
+ assert len(tokenizer.calls) == 2
+ assert tokenizer.calls[1]["messages"][-1] == {
+ "role": "assistant",
+ "content": "Spectrum caption zero.",
+ }
+ assert all(call["enable_thinking"] is False for call in tokenizer.calls)
+ assert tokenizer.calls[0]["add_generation_prompt"] is True
+ assert tokenizer.calls[1]["add_generation_prompt"] is False
+
+
+def test_current_accepted_schema_uses_original_not_revised(tmp_path: Path) -> None:
+ config = _write_case(tmp_path)
+ _rewrite_caption_records(
+ config,
+ lambda rows: rows[0]["spectrum_caption"].update(
+ {"revised": "A later, unsupported revision."}
+ ),
+ )
+ tokenizer = _FakeQwenTokenizer()
+
+ AstroPRISMStage1BSpectrumDataset(config, tokenizer)[0]
+
+ assert tokenizer.calls[1]["messages"][-1]["content"] == "Spectrum caption zero."
+
+
+def test_merged_length_guard_counts_273_spectrum_tokens_without_truncation(
+ tmp_path: Path,
+) -> None:
+ config = _write_case(tmp_path)
+ baseline = AstroPRISMStage1BSpectrumDataset(config, _FakeQwenTokenizer())[0]
+ constrained = replace(config, max_merged_seq_length=baseline["merged_length"] - 1)
+ tokenizer = _FakeQwenTokenizer()
+ dataset = AstroPRISMStage1BSpectrumDataset(constrained, tokenizer)
+
+ with pytest.raises(
+ AstroPRISMStage1ADatasetError,
+ match=r"273 spectrum tokens.*refusing silent truncation",
+ ):
+ dataset[0]
+
+ assert all("truncation" not in call for call in tokenizer.calls)
+
+
+def test_collator_stacks_nested_dtypes_and_pads_labels_in_lockstep(
+ tmp_path: Path,
+) -> None:
+ tokenizer = _FakeQwenTokenizer()
+ config = _write_case(tmp_path)
+ dataset = AstroPRISMStage1BSpectrumDataset(config, tokenizer)
+ samples = [dataset[0], dataset[1]]
+ original_lengths = [sample["text"].numel() for sample in samples]
+ collator = AstroPRISMStage1BSpectrumCollator(
+ tokenizer,
+ max_merged_seq_length=config.max_merged_seq_length,
+ )
+
+ batch = collator(samples)
+
+ longest = max(original_lengths)
+ shortest_index = original_lengths.index(min(original_lengths))
+ shortest = original_lengths[shortest_index]
+ assert {key: value.shape for key, value in batch["spectrum"].items()} == {
+ "flux": (2, ASTROPRISM_SPECTRUM_LENGTH),
+ "ivar": (2, ASTROPRISM_SPECTRUM_LENGTH),
+ "mask": (2, ASTROPRISM_SPECTRUM_LENGTH),
+ "wavelength": (2, ASTROPRISM_SPECTRUM_LENGTH),
+ }
+ assert batch["text"].shape == batch["labels"].shape == (2, longest)
+ assert batch["text_attention_mask"].shape == (2, longest)
+ assert torch.all(batch["text"][shortest_index, shortest:] == tokenizer.pad_token_id)
+ assert torch.all(batch["text_attention_mask"][shortest_index, shortest:] == 0)
+ assert torch.all(batch["labels"][shortest_index, shortest:] == IGNORE_INDEX)
+ assert batch["spectrum_token_count"] == 273
+ assert torch.equal(
+ batch["merged_length"],
+ torch.tensor([273 + length for length in original_lengths], dtype=torch.long),
+ )
+ assert batch["provenance"][1]["caption_supervision"] == "reviewed_spectrum_caption"
+
+
+def test_collator_rejects_wrong_nested_dtype(tmp_path: Path) -> None:
+ tokenizer = _FakeQwenTokenizer()
+ config = _write_case(tmp_path)
+ sample = AstroPRISMStage1BSpectrumDataset(config, tokenizer)[0]
+ sample["spectrum"]["mask"] = sample["spectrum"]["mask"].to(torch.float32)
+ collator = AstroPRISMStage1BSpectrumCollator(
+ tokenizer,
+ max_merged_seq_length=config.max_merged_seq_length,
+ )
+
+ with pytest.raises(ValueError, match=r"spectrum\.mask must have dtype torch\.bool"):
+ collator([sample])
+
+
+def test_rejects_non_promoted_caption_manifest(tmp_path: Path) -> None:
+ config = _write_case(tmp_path)
+ _rewrite_caption_manifest(
+ config,
+ lambda manifest: manifest.update({"artifact_type": "astroprism_spectrum_caption_pilot"}),
+ )
+
+ with pytest.raises(AstroPRISMStage1ADatasetError, match="artifact_type"):
+ AstroPRISMStage1BSpectrumDataset(config, _FakeQwenTokenizer())
+
+
+def test_rejects_raw_caption_roster_mismatch(tmp_path: Path) -> None:
+ config = _write_case(tmp_path)
+ _rewrite_caption_manifest(
+ config,
+ lambda manifest: manifest["roster"].update({"sha256": "other-roster"}),
+ )
+
+ with pytest.raises(AstroPRISMStage1ADatasetError, match="roster SHA-256 mismatch"):
+ AstroPRISMStage1BSpectrumDataset(config, _FakeQwenTokenizer())
+
+
+def test_checks_source_split_integrity_across_unselected_rows(tmp_path: Path) -> None:
+ config = _write_case(tmp_path)
+ _rewrite_caption_records(
+ config,
+ lambda rows: rows[2].update({"source_group_id": rows[0]["source_group_id"]}),
+ )
+
+ with pytest.raises(AstroPRISMStage1ADatasetError, match="source split leakage"):
+ AstroPRISMStage1BSpectrumDataset(config, _FakeQwenTokenizer())
+
+
+@pytest.mark.parametrize(
+ "keys,value,message",
+ [
+ (("spectrum_caption", "status"), "pending_review", "spectrum_caption.status"),
+ (("review", "status"), "pending_review", "review.status"),
+ (("review", "decision_source"), "heuristic", "review.decision_source"),
+ (("training_policy", "dataset_ready"), False, "dataset_ready"),
+ (("training_policy", "gradient_training_allowed"), False, "gradient_training_allowed"),
+ (("training_policy", "eligible_for_training"), False, "eligible_for_training"),
+ (("training_policy", "eligible_for_evaluation"), True, "eligible_for_evaluation"),
+ (("training_policy", "role"), "alignment_validation", "training_policy.role"),
+ ],
+)
+def test_train_view_rejects_unsafe_selected_rows(
+ tmp_path: Path,
+ keys: tuple[str, ...],
+ value: Any,
+ message: str,
+) -> None:
+ config = _write_case(tmp_path)
+ _rewrite_caption_records(
+ config,
+ lambda rows: _set_nested(rows[0], keys, value),
+ )
+
+ with pytest.raises(AstroPRISMStage1ADatasetError, match=message):
+ AstroPRISMStage1BSpectrumDataset(config, _FakeQwenTokenizer())
+
+
+def test_train_view_does_not_apply_review_gate_to_unselected_rows(tmp_path: Path) -> None:
+ config = _write_case(tmp_path)
+ _rewrite_caption_records(
+ config,
+ lambda rows: rows[2]["spectrum_caption"].update({"status": "pending_review"}),
+ )
+
+ dataset = AstroPRISMStage1BSpectrumDataset(config, _FakeQwenTokenizer())
+
+ assert len(dataset) == 2
+
+
+def test_validation_view_requires_evaluation_policy(tmp_path: Path) -> None:
+ config = replace(_write_case(tmp_path), split="validation")
+ dataset = AstroPRISMStage1BSpectrumDataset(config, _FakeQwenTokenizer())
+
+ assert len(dataset) == 1
+ assert dataset[0]["provenance"]["caption_training_role"] == "alignment_validation"
+
+ _rewrite_caption_records(
+ config,
+ lambda rows: rows[2]["training_policy"].update({"eligible_for_evaluation": False}),
+ )
+ with pytest.raises(AstroPRISMStage1ADatasetError, match="eligible_for_evaluation"):
+ AstroPRISMStage1BSpectrumDataset(config, _FakeQwenTokenizer())
+
+
+def test_caption_hash_verification_detects_tampering(tmp_path: Path) -> None:
+ config = _write_case(tmp_path)
+ _rewrite_caption_records(
+ config,
+ lambda rows: rows[0]["spectrum_caption"].update({"original": "Spectrum caption hero."}),
+ refresh_manifest=False,
+ )
+
+ with pytest.raises(AstroPRISMStage1ADatasetError, match="SHA-256 mismatch"):
+ AstroPRISMStage1BSpectrumDataset(config, _FakeQwenTokenizer())
diff --git a/tests/test_astroprism_stage1b_spectrum_train_entrypoint.py b/tests/test_astroprism_stage1b_spectrum_train_entrypoint.py
new file mode 100644
index 0000000..7b68c5b
--- /dev/null
+++ b/tests/test_astroprism_stage1b_spectrum_train_entrypoint.py
@@ -0,0 +1,176 @@
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import src.train as train_module
+import yaml
+from hydra import compose, initialize_config_dir
+from omegaconf import OmegaConf
+
+pytestmark = [pytest.mark.unit, pytest.mark.multimodal]
+
+
+class _Tokenizer:
+ pad_token_id = 0
+
+
+def test_spectrum_reviewed_hydra_presets_compose_to_frozen_safe_defaults():
+ config_dir = str((Path(__file__).parents[1] / "src" / "conf").resolve())
+ with initialize_config_dir(config_dir=config_dir, version_base="1.2"):
+ cfg = compose(
+ config_name="config",
+ overrides=[
+ "model=prism_astroprism_spectrum_qwen3_4b",
+ "training=astroprism_stage1b_spectrum",
+ "+data=astroprism_stage1b_spectrum",
+ ],
+ )
+
+ assert cfg.training.task == "astroprism_stage1b_spectrum"
+ assert cfg.training.max_steps == 100
+ assert cfg.training.freeze_llm is True
+ assert cfg.training.freeze_vit is True
+ assert cfg.training.data_num_workers == 0
+ assert cfg.training.eval_enabled is False
+ assert cfg.data.raw_manifest_path == "_local/astroprism/raw_spectra/manifest.json"
+ assert cfg.data.captions_manifest_path == (
+ "_local/astroprism/stage1b_spectrum_caption_accepted/manifest.json"
+ )
+ assert cfg.data.captions_path is None
+ assert cfg.data.split == "train"
+ assert list(cfg.data.accepted_caption_statuses) == ["accepted"]
+ assert "allow_unsafe_candidates" not in cfg.data
+ assert cfg.data.verify_file_hashes is True
+ assert list(cfg.model.modalities) == ["text", "spectrum"]
+ assert cfg.model.spectrum_encoder_type == "aion"
+ assert cfg.model.d_spectrum == 768
+ assert cfg.model.aion_num_spectrum_tokens == 273
+ assert cfg.model.aion_raw_spectrum_length == 7781
+ assert cfg.model.projector_norm_mode == "text_scaled_layernorm"
+ assert cfg.model.projector_target_norm == pytest.approx(1.04)
+ assert cfg.model.projector_modality_embed_pos == "before_norm"
+
+
+def test_spectrum_reviewed_builder_forwards_paths_safety_prompts_and_length(
+ monkeypatch,
+ tmp_path,
+):
+ captured = {}
+
+ class _Dataset:
+ def __init__(self, config, tokenizer):
+ captured["dataset_config"] = config
+ captured["dataset_tokenizer"] = tokenizer
+
+ class _Collator:
+ def __init__(self, tokenizer, *, max_merged_seq_length):
+ captured["collator_tokenizer"] = tokenizer
+ captured["collator_max"] = max_merged_seq_length
+
+ monkeypatch.setattr(
+ train_module,
+ "AstroPRISMStage1BSpectrumDataset",
+ _Dataset,
+ )
+ monkeypatch.setattr(
+ train_module,
+ "AstroPRISMStage1BSpectrumCollator",
+ _Collator,
+ )
+
+ tokenizer = _Tokenizer()
+ raw_manifest = tmp_path / "raw-spectra" / "manifest.json"
+ captions_manifest = tmp_path / "accepted-spectra" / "manifest.json"
+ captions = tmp_path / "accepted-spectra" / "captions.jsonl"
+ data_cfg = OmegaConf.create(
+ {
+ "raw_manifest_path": str(raw_manifest),
+ "captions_manifest_path": str(captions_manifest),
+ "captions_path": str(captions),
+ "split": "train",
+ "accepted_caption_statuses": ["accepted"],
+ "verify_file_hashes": True,
+ "max_merged_seq_length": 2048,
+ }
+ )
+
+ dataset, collator = train_module._build_astroprism_stage1b_spectrum_data(
+ data_cfg,
+ tokenizer,
+ )
+
+ assert isinstance(dataset, _Dataset)
+ assert isinstance(collator, _Collator)
+ config = captured["dataset_config"]
+ assert config.raw_manifest_path == raw_manifest
+ assert config.captions_manifest_path == captions_manifest
+ assert config.captions_path == captions
+ assert config.split == "train"
+ assert config.system_prompt == train_module.DEFAULT_SPECTRUM_SYSTEM_PROMPT
+ assert config.user_prompt == train_module.DEFAULT_SPECTRUM_USER_PROMPT
+ assert config.accepted_caption_statuses == ("accepted",)
+ assert config.verify_file_hashes is True
+ assert captured["dataset_tokenizer"] is tokenizer
+ assert captured["collator_tokenizer"] is tokenizer
+ assert captured["collator_max"] == 2048
+
+
+def test_spectrum_reviewed_builder_requires_its_data_group():
+ with pytest.raises(RuntimeError, match="data=astroprism_stage1b_spectrum"):
+ train_module._build_astroprism_stage1b_spectrum_data(None, _Tokenizer())
+
+
+def test_spectrum_reviewed_model_contract_accepts_only_raw_aion_spectrum_text():
+ config = SimpleNamespace(
+ spectrum_encoder_type="aion",
+ modalities=["text", "spectrum"],
+ )
+ train_module._validate_astroprism_stage1b_spectrum_model_contract(
+ "astroprism_stage1b_spectrum",
+ config,
+ )
+
+
+@pytest.mark.parametrize(
+ "spectrum_encoder_type,modalities,message",
+ [
+ ("generic", ["text", "spectrum"], "spectrum_encoder_type=aion"),
+ ("aion", ["text"], "requires exactly"),
+ ("aion", ["text", "image", "spectrum"], "requires exactly"),
+ ],
+)
+def test_spectrum_reviewed_model_contract_fails_closed(
+ spectrum_encoder_type,
+ modalities,
+ message,
+):
+ config = SimpleNamespace(
+ spectrum_encoder_type=spectrum_encoder_type,
+ modalities=modalities,
+ )
+ with pytest.raises(RuntimeError, match=message):
+ train_module._validate_astroprism_stage1b_spectrum_model_contract(
+ "astroprism_stage1b_spectrum",
+ config,
+ )
+
+
+def test_spectrum_reviewed_experiment_is_a_100_step_frozen_accelerate_design():
+ path = Path("experiments/astroprism_stage1b_spectrum.yaml")
+ [experiment] = yaml.safe_load(path.read_text(encoding="utf-8"))["experiments"]
+ overrides = experiment["common_overrides"]
+
+ assert experiment["id"] == "ASTROPRISM-STAGE1B-SPECTRUM"
+ assert overrides["training.task"] == "astroprism_stage1b_spectrum"
+ assert overrides["training.max_steps"] == 100
+ assert overrides["training.freeze_llm"] is True
+ assert overrides["training.freeze_vit"] is True
+ assert overrides["model.freeze_backbone"] is True
+ assert overrides["model.freeze_encoders"] is True
+ assert overrides["model.projector_norm_mode"] == "text_scaled_layernorm"
+ assert overrides["model.projector_target_norm"] == pytest.approx(1.04)
+ assert overrides["model.projector_modality_embed_pos"] == "before_norm"
+ assert overrides["data.split"] == "train"
+ assert overrides["data.accepted_caption_statuses"] == ["accepted"]
+ assert "data.allow_unsafe_candidates" not in overrides
+ assert "astroprism_stage1b_spectrum" in train_module._ASTROPRISM_MAP_STYLE_TASKS
diff --git a/tests/test_astroprism_train_entrypoint.py b/tests/test_astroprism_train_entrypoint.py
new file mode 100644
index 0000000..0781c48
--- /dev/null
+++ b/tests/test_astroprism_train_entrypoint.py
@@ -0,0 +1,140 @@
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import src.train as train_module
+from hydra import compose, initialize_config_dir
+from omegaconf import OmegaConf
+
+
+class _Tokenizer:
+ pad_token_id = 0
+
+
+def test_astroprism_hydra_presets_compose_to_safe_raw_image_defaults():
+ config_dir = str((Path(__file__).parents[1] / "src" / "conf").resolve())
+ with initialize_config_dir(config_dir=config_dir, version_base="1.2"):
+ cfg = compose(
+ config_name="config",
+ overrides=[
+ "model=prism_astroprism_qwen3_4b",
+ "training=astroprism_stage1a",
+ "+data=astroprism_stage1a",
+ ],
+ )
+
+ assert cfg.training.task == "astroprism_stage1a"
+ assert cfg.training.freeze_llm is True
+ assert cfg.training.freeze_vit is True
+ assert cfg.training.data_num_workers == 0
+ assert cfg.data.raw_manifest_path == ("_local/astroprism/raw_images/manifest.json")
+ assert cfg.data.candidates_manifest_path == ("_local/astroprism/stage1a_accepted/manifest.json")
+ assert cfg.data.split == "train"
+ assert cfg.data.allow_unsafe_candidates is False
+ assert cfg.data.verify_file_hashes is True
+ assert list(cfg.data.accepted_caption_statuses) == [
+ "released_reused",
+ "released_revised",
+ ]
+ assert cfg.data.caption_field == "auto"
+ assert cfg.data.max_merged_seq_length == 2048
+ assert cfg.model.image_encoder_type == "aion"
+
+
+def test_stage1a_builder_forwards_paths_safety_and_length(monkeypatch, tmp_path):
+ captured = {}
+
+ class _Dataset:
+ def __init__(self, config, tokenizer):
+ captured["dataset_config"] = config
+ captured["dataset_tokenizer"] = tokenizer
+
+ class _Collator:
+ def __init__(self, tokenizer, *, max_merged_seq_length):
+ captured["collator_tokenizer"] = tokenizer
+ captured["collator_max"] = max_merged_seq_length
+
+ monkeypatch.setattr(train_module, "AstroPRISMStage1ADataset", _Dataset)
+ monkeypatch.setattr(train_module, "AstroPRISMStage1ACollator", _Collator)
+
+ tokenizer = _Tokenizer()
+ raw_manifest = tmp_path / "raw" / "manifest.json"
+ candidate_manifest = tmp_path / "candidates" / "manifest.json"
+ data_cfg = OmegaConf.create(
+ {
+ "raw_manifest_path": str(raw_manifest),
+ "candidates_manifest_path": str(candidate_manifest),
+ "candidates_path": None,
+ "split": "train",
+ "caption_field": "auto",
+ "allow_unsafe_candidates": False,
+ "accepted_statuses": ["accepted", "approved"],
+ "accepted_caption_statuses": ["released_reused", "released_revised"],
+ "verify_file_hashes": True,
+ "max_merged_seq_length": 2048,
+ }
+ )
+
+ dataset, collator = train_module._build_astroprism_stage1a_data(
+ data_cfg,
+ tokenizer,
+ )
+
+ assert isinstance(dataset, _Dataset)
+ assert isinstance(collator, _Collator)
+ config = captured["dataset_config"]
+ assert config.raw_manifest_path == raw_manifest
+ assert config.candidates_manifest_path == candidate_manifest
+ assert config.candidates_path is None
+ assert config.split == "train"
+ assert config.caption_field == "auto"
+ assert config.allow_unsafe_candidates is False
+ assert config.accepted_statuses == ("accepted", "approved")
+ assert config.accepted_caption_statuses == (
+ "released_reused",
+ "released_revised",
+ )
+ assert config.verify_file_hashes is True
+ assert captured["dataset_tokenizer"] is tokenizer
+ assert captured["collator_tokenizer"] is tokenizer
+ assert captured["collator_max"] == 2048
+
+
+def test_stage1a_builder_requires_data_group():
+ with pytest.raises(RuntimeError, match="data=astroprism_stage1a"):
+ train_module._build_astroprism_stage1a_data(None, _Tokenizer())
+
+
+def test_stage1a_model_contract_accepts_only_raw_aion_image_text():
+ config = SimpleNamespace(
+ image_encoder_type="aion",
+ modalities=["text", "image"],
+ )
+ train_module._validate_astroprism_stage1a_model_contract(
+ "astroprism_stage1a",
+ config,
+ )
+
+
+@pytest.mark.parametrize(
+ "image_encoder_type,modalities,message",
+ [
+ ("siglip", ["text", "image"], "image_encoder_type=aion"),
+ ("aion", ["text"], "requires exactly"),
+ ("aion", ["text", "image", "table"], "requires exactly"),
+ ],
+)
+def test_stage1a_model_contract_fails_closed(
+ image_encoder_type,
+ modalities,
+ message,
+):
+ config = SimpleNamespace(
+ image_encoder_type=image_encoder_type,
+ modalities=modalities,
+ )
+ with pytest.raises(RuntimeError, match=message):
+ train_module._validate_astroprism_stage1a_model_contract(
+ "astroprism_stage1a",
+ config,
+ )
diff --git a/tests/test_compose_astroprism_joint_adapters.py b/tests/test_compose_astroprism_joint_adapters.py
new file mode 100644
index 0000000..3391080
--- /dev/null
+++ b/tests/test_compose_astroprism_joint_adapters.py
@@ -0,0 +1,47 @@
+from pathlib import Path
+
+import pytest
+import torch
+from safetensors.torch import load_file, save_file
+from tools.compose_astroprism_joint_adapters import (
+ EXPECTED_SHAPES,
+ compose_joint_adapters,
+)
+
+
+def _checkpoint(path: Path, prefix: str, *, fill: float) -> Path:
+ path.mkdir()
+ tensors = {
+ key: torch.full(shape, fill, dtype=torch.float16)
+ for key, shape in EXPECTED_SHAPES.items()
+ if key.startswith(prefix)
+ }
+ tensors["backbone.large_frozen_weight"] = torch.zeros(2)
+ save_file(tensors, path / "model.safetensors")
+ return path
+
+
+def test_compose_joint_adapters_selects_only_both_projectors(tmp_path):
+ image = _checkpoint(tmp_path / "image", "projectors.image.", fill=1.0)
+ spectrum = _checkpoint(tmp_path / "spectrum", "projectors.spectrum.", fill=2.0)
+ output = tmp_path / "joint"
+
+ manifest = compose_joint_adapters(image, spectrum, output)
+ result = load_file(output / "model.safetensors")
+
+ assert set(result) == set(EXPECTED_SHAPES)
+ assert torch.all(result["projectors.image.fc1.bias"] == 1)
+ assert torch.all(result["projectors.spectrum.fc1.bias"] == 2)
+ assert manifest["tensor_count"] == len(EXPECTED_SHAPES)
+ assert len(manifest["merged_adapter_sha256"]) == 64
+ assert (output / "training_state.json").is_file()
+
+
+def test_compose_joint_adapters_refuses_overwrite(tmp_path):
+ image = _checkpoint(tmp_path / "image", "projectors.image.", fill=1.0)
+ spectrum = _checkpoint(tmp_path / "spectrum", "projectors.spectrum.", fill=2.0)
+ output = tmp_path / "joint"
+ compose_joint_adapters(image, spectrum, output)
+
+ with pytest.raises(FileExistsError, match="Refusing to overwrite"):
+ compose_joint_adapters(image, spectrum, output)
diff --git a/tests/test_launch_aurora.py b/tests/test_launch_aurora.py
new file mode 100644
index 0000000..5ff8985
--- /dev/null
+++ b/tests/test_launch_aurora.py
@@ -0,0 +1,135 @@
+"""Focused safety tests for the generic Aurora launcher."""
+
+from __future__ import annotations
+
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+from tools import launch_aurora
+
+pytestmark = [pytest.mark.launcher, pytest.mark.aurora]
+
+
+def _write_design(path: Path) -> Path:
+ design = path / "design.yaml"
+ backbone = path / "local-backbone"
+ backbone.mkdir(exist_ok=True)
+ design.write_text(
+ "\n".join(
+ [
+ "experiments:",
+ " - id: TEST-LAUNCHER",
+ ' name: "launcher safety test"',
+ " resources:",
+ " ngpus: 1",
+ ' walltime: "00:10:00"',
+ " common_overrides:",
+ " training.max_steps: 1",
+ f' model.backbone_id: "{backbone}"',
+ ]
+ ),
+ encoding="utf-8",
+ )
+ return design
+
+
+def _launcher_argv(
+ design: Path,
+ packed_env: Path,
+ *,
+ batch: bool,
+ dry_run: bool,
+) -> list[str]:
+ argv = [
+ "launch_aurora.py",
+ "--file",
+ str(design),
+ "--design",
+ "TEST-LAUNCHER",
+ "--id",
+ "TEST-LAUNCHER",
+ "--nodes",
+ "1",
+ "--packed-env",
+ str(packed_env),
+ ]
+ if batch:
+ argv.append("--batch")
+ if dry_run:
+ argv.append("--dry-run")
+ return argv
+
+
+def _generated_script(tmp_path: Path) -> str:
+ scripts = list((tmp_path / "jobs").glob("run_aurora_TEST-LAUNCHER_*"))
+ assert len(scripts) == 1
+ return scripts[0].read_text(encoding="utf-8")
+
+
+def test_packed_env_uses_documented_nested_venv_layout(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ """The tarball holds .venv-deepspeed/, so activate must look inside it."""
+ design = _write_design(tmp_path)
+ packed_env = tmp_path / "deepspeed_env.tar.gz"
+ monkeypatch.chdir(tmp_path)
+ monkeypatch.setattr(
+ sys,
+ "argv",
+ _launcher_argv(design, packed_env, batch=True, dry_run=True),
+ )
+
+ launch_aurora.main()
+
+ script = _generated_script(tmp_path)
+ nested_activate = "$LOCAL_ENV/.venv-deepspeed/bin/activate"
+ assert nested_activate in script
+ assert "source $LOCAL_ENV/bin/activate" not in script
+ assert 'VIRTUAL_ENV=\\"$LOCAL_ENV/.venv-deepspeed\\"' in script
+ assert '[ "$ENV_TARBALL" -nt "$MARKER_FILE" ]' in script
+ assert 'tar -xzf "$ENV_TARBALL" -C "$LOCAL_ENV"' in script
+ # The wait must test the same staleness condition as the unpack: with a
+ # stale marker on disk, an existence-only wait lets other ranks activate
+ # the environment while rank 0 is deleting it.
+ wait_line = next(
+ line for line in script.splitlines() if line.startswith("while [ ! -f")
+ )
+ assert '[ "$ENV_TARBALL" -nt "$MARKER_FILE" ]' in wait_line
+
+
+def test_batch_mode_submits_generated_script_with_qsub(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ """--batch must qsub the PBS script, not execute it on the login node."""
+ design = _write_design(tmp_path)
+ packed_env = tmp_path / "deepspeed_env.tar.gz"
+ calls: list[tuple[list[str], dict[str, object]]] = []
+
+ def fake_run(command, **kwargs):
+ calls.append((command, kwargs))
+ return subprocess.CompletedProcess(
+ command,
+ returncode=0,
+ stdout="12345.aurora-pbs-0001\n",
+ stderr="",
+ )
+
+ monkeypatch.chdir(tmp_path)
+ monkeypatch.setattr(
+ sys,
+ "argv",
+ _launcher_argv(design, packed_env, batch=True, dry_run=False),
+ )
+ monkeypatch.setattr(launch_aurora.subprocess, "run", fake_run)
+
+ launch_aurora.main()
+
+ assert len(calls) == 1
+ command, kwargs = calls[0]
+ assert command[0] == "qsub"
+ assert Path(command[1]).name.startswith("run_aurora_TEST-LAUNCHER_")
+ assert kwargs == {"check": True, "capture_output": True, "text": True}
diff --git a/tests/test_launch_aurora_web_vla.py b/tests/test_launch_aurora_web_vla.py
index 6b35e35..b593d17 100644
--- a/tests/test_launch_aurora_web_vla.py
+++ b/tests/test_launch_aurora_web_vla.py
@@ -415,3 +415,15 @@ def test_web_shared_venv_uses_environment_venv_path(tmp_path):
)
script_text = (tmp_path / generated_line.split(": ", 1)[1].strip()).read_text()
assert 'source "/flare/ModCon/sww/prism-envs/qwen3-siglip-py3.12/bin/activate"' in script_text
+
+
+def test_skip_model_staging_rejects_remote_model_ids(tmp_path):
+ """Without absolute shared paths, skipping staging would leave nothing to load."""
+ result = _run_launcher(
+ tmp_path,
+ "vlm",
+ extra_args=("--skip-model-staging",),
+ )
+ assert result.returncode != 0
+ assert "--skip-model-staging requires absolute, existing shared paths" in result.stderr
+ assert "model.backbone_id" in result.stderr
diff --git a/tests/test_materialize_astroprism_hf_dataset.py b/tests/test_materialize_astroprism_hf_dataset.py
new file mode 100644
index 0000000..af2338b
--- /dev/null
+++ b/tests/test_materialize_astroprism_hf_dataset.py
@@ -0,0 +1,252 @@
+"""The HF package and the training configs use different manifest envelopes.
+
+These cover the bridge between them, with a synthetic snapshot: no network, no
+gated repo, no 113 GB of arrays.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import numpy as np
+import pytest
+from tools import materialize_astroprism_hf_dataset as mat
+
+BANDS_SOURCE = [b"des-g", b"des-r", b"des-i", b"des-z"]
+BANDS_AION = [b"DES-G", b"DES-R", b"DES-I", b"DES-Z"]
+
+
+def _snapshot(tmp_path: Path, rows: int = 3) -> Path:
+ """Build a miniature stand-in for the published dataset tree."""
+ snap = tmp_path / "snapshot"
+ images = snap / "assets" / "raw" / "images"
+ spectra = snap / "assets" / "raw" / "spectra"
+ meta = snap / "metadata"
+ for d in (images, spectra, meta):
+ d.mkdir(parents=True)
+
+ np.save(images / "image_array.npy", np.zeros((rows, 4, 8, 8), dtype=np.float32))
+ np.save(images / "image_band.npy", np.array([BANDS_SOURCE] * rows, dtype="S5"))
+ np.save(images / "aion_image_band.npy", np.array([BANDS_AION] * rows, dtype="S5"))
+ np.save(spectra / "spectrum_flux.npy", np.zeros((rows, 16), dtype=np.float32))
+
+ (images / "MANIFEST.json").write_text(json.dumps({
+ "artifact_type": "astroprism_portable_raw_images",
+ "row_count": rows,
+ "outputs": {"image_array": {"path": "image_array.npy", "shape": [rows, 4, 8, 8]}},
+ }))
+ (spectra / "MANIFEST.json").write_text(json.dumps({
+ "artifact_type": "astroprism_portable_raw_spectra",
+ "row_count": rows,
+ "outputs": {"spectrum_flux": {"path": "spectrum_flux.npy", "shape": [rows, 16]}},
+ }))
+ # Spectrum captions live in their own Parquet config, not in
+ # candidates.jsonl -- the published dataset splits them that way.
+ import pyarrow as pa
+ import pyarrow.parquet as pq
+
+ spec_cfg = snap / "data" / "spectrum_caption"
+ spec_cfg.mkdir(parents=True)
+ for split, idxs in (("train", list(range(rows))), ("validation", []), ("test", [])):
+ pq.write_table(
+ pa.table({
+ "candidate_index": idxs,
+ "paired_targetid": [1000 + i for i in idxs],
+ "caption_object_id": [f"obj-{i}" for i in idxs],
+ "source_group_id": [f"grp-{i}" for i in idxs],
+ "split": [split] * len(idxs),
+ "caption": [f"a spectrum {i}" for i in idxs],
+ "caption_status": ["accepted"] * len(idxs),
+ "review_status": ["accepted"] * len(idxs),
+ "training_allowed": [split == "train"] * len(idxs),
+ "dataset_ready": [True] * len(idxs),
+ }),
+ spec_cfg / f"{split}.parquet",
+ )
+
+ with open(meta / "candidates.jsonl", "w", encoding="utf-8") as fh:
+ for i in range(rows):
+ fh.write(json.dumps({
+ "candidate_index": i,
+ "paired_targetid": 1000 + i,
+ "caption_object_id": f"obj-{i}",
+ "split": "train",
+ "acceptance_status": "accepted",
+ "caption": {"original": "a galaxy", "status": "released_reused"},
+ }) + "\n")
+ return snap
+
+
+def _run(snap: Path, out: Path, *extra: str) -> int:
+ return mat.main(["--snapshot-dir", str(snap), "--no-download",
+ "--output-dir", str(out), *extra])
+
+
+def test_writes_the_four_manifests_the_configs_read(tmp_path: Path) -> None:
+ out = tmp_path / "local"
+ assert _run(_snapshot(tmp_path), out) == 0
+ for rel in (
+ "raw_images/manifest.json",
+ "raw_spectra/manifest.json",
+ "stage1a_accepted/manifest.json",
+ "stage1b_spectrum_caption_accepted/manifest.json",
+ ):
+ assert (out / rel).is_file(), rel
+
+
+def test_artifact_types_match_what_the_loaders_demand(tmp_path: Path) -> None:
+ """The loaders reject any other value outright."""
+ out = tmp_path / "local"
+ _run(_snapshot(tmp_path), out)
+ images = json.loads((out / "raw_images/manifest.json").read_text())
+ spectra = json.loads((out / "raw_spectra/manifest.json").read_text())
+ cands = json.loads((out / "stage1a_accepted/manifest.json").read_text())
+ assert images["artifact_type"] == "astroprism_raw_images"
+ assert spectra["artifact_type"] == "astroprism_raw_spectra"
+ assert cands["artifact_type"] in {
+ "astroprism_stage1a_candidates",
+ "astroprism_aion_feature_parity_reference",
+ }
+
+
+def test_output_paths_are_absolute_so_the_tree_can_live_anywhere(tmp_path: Path) -> None:
+ """Loaders resolve relative paths against the manifest's own directory."""
+ out = tmp_path / "local"
+ _run(_snapshot(tmp_path), out)
+ images = json.loads((out / "raw_images/manifest.json").read_text())
+ path = Path(images["outputs"]["image_array"]["path"])
+ assert path.is_absolute() and path.is_file()
+
+
+def test_candidate_rows_carry_index_and_band_order(tmp_path: Path) -> None:
+ """The loader cross-checks these row for row against the roster."""
+ out = tmp_path / "local"
+ _run(_snapshot(tmp_path, rows=3), out)
+ images = json.loads((out / "raw_images/manifest.json").read_text())
+ rows = images["candidates"]
+ assert [r["candidate_index"] for r in rows] == [0, 1, 2]
+ assert rows[0]["bands"]["source"] == ["des-g", "des-r", "des-i", "des-z"]
+ assert rows[0]["bands"]["aion"] == ["DES-G", "DES-R", "DES-I", "DES-Z"]
+ assert rows[1]["caption_object_id"] == "obj-1"
+
+
+def test_candidate_row_count_and_hash_are_recorded(tmp_path: Path) -> None:
+ out = tmp_path / "local"
+ _run(_snapshot(tmp_path, rows=3), out)
+ cands = json.loads((out / "stage1a_accepted/manifest.json").read_text())
+ entry = cands["outputs"]["candidates"]
+ assert entry["row_count"] == 3
+ assert len(entry["sha256"]) == 64
+
+
+def test_hash_is_skipped_on_request(tmp_path: Path) -> None:
+ out = tmp_path / "local"
+ _run(_snapshot(tmp_path), out, "--no-verify-hashes")
+ cands = json.loads((out / "stage1a_accepted/manifest.json").read_text())
+ assert "sha256" not in cands["outputs"]["candidates"]
+
+
+def test_missing_array_fails_loudly(tmp_path: Path) -> None:
+ """A partial snapshot must not yield a manifest pointing at nothing."""
+ snap = _snapshot(tmp_path)
+ (snap / "assets/raw/images/image_array.npy").unlink()
+ assert _run(snap, tmp_path / "local") == 1
+
+
+def test_band_array_length_mismatch_fails(tmp_path: Path) -> None:
+ snap = _snapshot(tmp_path, rows=3)
+ np.save(snap / "assets/raw/images/image_band.npy",
+ np.array([BANDS_SOURCE] * 2, dtype="S5"))
+ assert _run(snap, tmp_path / "local") == 1
+
+
+def test_no_download_without_a_snapshot_is_rejected(tmp_path: Path) -> None:
+ assert mat.main(["--no-download", "--output-dir", str(tmp_path / "local")]) == 1
+
+
+def test_missing_token_is_reported_rather_than_a_traceback(
+ tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ monkeypatch.delenv("HF_TOKEN", raising=False)
+ assert mat.main(["--output-dir", str(tmp_path / "local")]) == 1
+
+
+def test_spectrum_roster_comes_from_the_parquet_not_candidates_jsonl(
+ tmp_path: Path,
+) -> None:
+ """The published candidates.jsonl carries image captions only."""
+ out = tmp_path / "local"
+ _run(_snapshot(tmp_path, rows=3), out)
+ rows = [
+ json.loads(line)
+ for line in (out / "stage1b_spectrum_caption_accepted/captions.jsonl")
+ .read_text(encoding="utf-8")
+ .splitlines()
+ if line.strip()
+ ]
+ assert [r["candidate_index"] for r in rows] == [0, 1, 2]
+ assert rows[0]["spectrum_caption"]["original"] == "a spectrum 0"
+ assert rows[0]["spectrum_caption"]["status"] == "accepted"
+
+
+def test_split_roles_and_gradient_policy_follow_the_dataset_card(
+ tmp_path: Path,
+) -> None:
+ """Only train may take gradients; validation and test are inspection-only."""
+ out = tmp_path / "local"
+ _run(_snapshot(tmp_path, rows=2), out)
+ rows = [
+ json.loads(line)
+ for line in (out / "stage1b_spectrum_caption_accepted/captions.jsonl")
+ .read_text(encoding="utf-8")
+ .splitlines()
+ if line.strip()
+ ]
+ policy = rows[0]["training_policy"]
+ assert policy["role"] == "alignment_training"
+ assert policy["gradient_training_allowed"] is True
+ assert policy["eligible_for_evaluation"] is False
+
+
+def test_both_manifests_declare_the_same_roster(tmp_path: Path) -> None:
+ """The loader compares their SHA-256 to prove one roster describes both."""
+ out = tmp_path / "local"
+ _run(_snapshot(tmp_path), out)
+ raw = json.loads((out / "raw_spectra/manifest.json").read_text())
+ caps = json.loads((out / "stage1b_spectrum_caption_accepted/manifest.json").read_text())
+ assert raw["roster"]["sha256"] == caps["roster"]["sha256"]
+ assert raw["roster"]["row_count"] == caps["roster"]["row_count"]
+
+
+def test_missing_spectrum_parquet_fails_with_a_usable_message(tmp_path: Path) -> None:
+ snap = _snapshot(tmp_path)
+ (snap / "data/spectrum_caption/train.parquet").unlink()
+ assert _run(snap, tmp_path / "local") == 1
+
+
+def test_training_allowed_disagreeing_with_the_split_is_rejected(
+ tmp_path: Path,
+) -> None:
+ """A row marked trainable outside train must not quietly become trainable."""
+ import pyarrow as pa
+ import pyarrow.parquet as pq
+
+ snap = _snapshot(tmp_path, rows=2)
+ bad = snap / "data/spectrum_caption/test.parquet"
+ pq.write_table(
+ pa.table({
+ "candidate_index": [0],
+ "paired_targetid": [1000],
+ "caption_object_id": ["obj-0"],
+ "source_group_id": ["grp-0"],
+ "split": ["test"],
+ "caption": ["a spectrum"],
+ "caption_status": ["accepted"],
+ "review_status": ["accepted"],
+ "training_allowed": [True], # wrong for a test-split row
+ "dataset_ready": [True],
+ }),
+ bad,
+ )
+ assert _run(snap, tmp_path / "local") == 1
diff --git a/tests/test_prepare_astroprism_hf_dataset.py b/tests/test_prepare_astroprism_hf_dataset.py
new file mode 100644
index 0000000..0a067a4
--- /dev/null
+++ b/tests/test_prepare_astroprism_hf_dataset.py
@@ -0,0 +1,427 @@
+from __future__ import annotations
+
+import gzip
+import io
+import json
+from pathlib import Path
+
+import numpy as np
+import pyarrow as pa
+import pyarrow.parquet as pq
+import pytest
+from datasets import Features, Image
+from PIL import Image as PILImage
+from tools.prepare_astroprism_hf_dataset import (
+ CAPTION_REPO_ID,
+ CAPTION_REVISION,
+ IMAGE_ARRAY_CONTRACT,
+ PAIRED_REPO_ID,
+ PAIRED_REVISION,
+ SPECTRUM_ARRAY_CONTRACT,
+ HFDatasetConfig,
+ PackagingValidationError,
+ _dataset_card,
+ _validate_accepted_rows,
+ prepare_hf_dataset,
+ sha256_file,
+)
+
+pytestmark = [pytest.mark.unit, pytest.mark.multimodal]
+
+
+def _png(size: tuple[int, int], color: tuple[int, int, int]) -> bytes:
+ stream = io.BytesIO()
+ PILImage.new("RGB", size, color).save(stream, format="PNG")
+ return stream.getvalue()
+
+
+def _write_json(path: Path, value: object) -> None:
+ path.parent.mkdir(parents=True, exist_ok=True)
+ path.write_text(json.dumps(value, sort_keys=True), encoding="utf-8")
+
+
+def _write_jsonl(path: Path, rows: list[dict]) -> None:
+ path.parent.mkdir(parents=True, exist_ok=True)
+ path.write_text(
+ "".join(json.dumps(row, sort_keys=True) + "\n" for row in rows),
+ encoding="utf-8",
+ )
+
+
+def _raw_fixture(
+ root: Path,
+ accepted: list[dict],
+ contract: dict[str, tuple[np.dtype, tuple[int, ...]]],
+ *,
+ spectra: bool,
+) -> Path:
+ root.mkdir(parents=True)
+ outputs = {}
+ for offset, (key, (dtype, trailing_shape)) in enumerate(contract.items()):
+ shape = (len(accepted), *trailing_shape)
+ if key == "spectrum_wavelength":
+ array = np.broadcast_to(np.arange(trailing_shape[0], dtype=np.float32), shape).copy()
+ elif dtype.kind == "S":
+ array = np.full(shape, b"band", dtype=dtype)
+ elif dtype == np.dtype("bool"):
+ array = np.zeros(shape, dtype=dtype)
+ else:
+ array = np.full(shape, offset + 1, dtype=dtype)
+ path = root / f"{key}.npy"
+ np.save(path, array, allow_pickle=False)
+ outputs[key] = {
+ "path": path.name,
+ "dtype": str(array.dtype),
+ "shape": list(array.shape),
+ "size_bytes": path.stat().st_size,
+ "sha256": sha256_file(path),
+ "mmap_safe": True,
+ "source_dataset": f"source/{key}",
+ }
+ manifest = {
+ "schema_version": 1,
+ "artifact_type": "raw_spectra" if spectra else "raw_images",
+ "candidates": [
+ {
+ "candidate_index": row["candidate_index"],
+ "paired_targetid": row["paired_targetid"],
+ "caption_object_id": row["caption_object_id"],
+ }
+ for row in accepted
+ ],
+ "outputs": outputs,
+ "source": {
+ "url": "https://example.test/raw.hdf5",
+ "etag": '"fixture"',
+ "length_bytes": 123,
+ "hdf5_row_count": 3,
+ },
+ **(
+ {
+ "spectrum": {
+ "length": 7781,
+ "observed_frame": True,
+ "mask_true_means_excluded": True,
+ }
+ }
+ if spectra
+ else {"bands": {"source_labels_preserved": True}}
+ ),
+ }
+ manifest_path = root / "manifest.json"
+ _write_json(manifest_path, manifest)
+ return manifest_path
+
+
+def _spectrum_bundle(root: Path, accepted: list[dict], indices: list[int]) -> Path:
+ plots = root / "plots"
+ plots.mkdir(parents=True)
+ records = []
+ for index in indices:
+ source = accepted[index]
+ records.append(
+ {
+ "candidate_index": index,
+ "paired_targetid": source["paired_targetid"],
+ "caption_object_id": source["caption_object_id"],
+ "source_group_id": source["source_group_id"],
+ "split": source["split"],
+ "student_input": {
+ "wavelength_frame": "observed_vacuum",
+ "spectrum_length": 7781,
+ },
+ "spectrum_caption": {
+ "original": f"Grounded spectrum caption {index}.",
+ "status": "accepted",
+ },
+ "review": {
+ "status": "accepted",
+ "decision_source": "user_instruction",
+ },
+ "teacher_reference": {
+ "catalog_context": {"redshift": 0.1 + index, "redshift_error": 0.01}
+ },
+ "audit_sidecar": {"claims": [{"text": "grounded", "fact_ids": ["x"]}]},
+ "generation": {"method": "fixture"},
+ "training_policy": {
+ "dataset_ready": True,
+ "gradient_training_allowed": source["split"] == "train",
+ "eligible_for_training": source["split"] == "train",
+ "eligible_for_evaluation": source["split"] != "train",
+ },
+ }
+ )
+ (plots / f"candidate_{index:04d}.png").write_bytes(_png((32, 16), (10 + index, 20, 30)))
+ captions = root / "captions.jsonl"
+ _write_jsonl(captions, records)
+ outputs = {
+ "captions.jsonl": {
+ "sha256": sha256_file(captions),
+ "size_bytes": captions.stat().st_size,
+ }
+ }
+ for path in plots.iterdir():
+ outputs[f"plots/{path.name}"] = {
+ "sha256": sha256_file(path),
+ "size_bytes": path.stat().st_size,
+ }
+ _write_json(root / "manifest.json", {"row_count": len(records), "outputs": outputs})
+ return root
+
+
+def _complete_fixture(tmp_path: Path) -> HFDatasetConfig:
+ caption_root = tmp_path / "caption_source"
+ paired_root = tmp_path / "paired_source"
+ (caption_root / "data").mkdir(parents=True)
+ (paired_root / "data").mkdir(parents=True)
+ caption_shard = caption_root / "data/captions.parquet"
+ paired_shard = paired_root / "data/paired.parquet"
+ captions = [f"Accepted image caption {index}." for index in range(3)]
+ summaries = [f"Summary {index}." for index in range(3)]
+ pq.write_table(
+ pa.Table.from_pylist(
+ [
+ {
+ "object_id": f"caption-{index}",
+ "survey": "legacy_dr10",
+ "ra": float(index),
+ "dec": -float(index),
+ "image": _png((160, 160), (index, 2 * index, 3 * index)),
+ "caption": captions[index],
+ "summary": summaries[index],
+ }
+ for index in range(3)
+ ]
+ ),
+ caption_shard,
+ )
+ pq.write_table(
+ pa.Table.from_pylist(
+ [
+ {
+ "object_id_provabgs": f"100{index}",
+ "ra_provabgs": float(index),
+ "dec_provabgs": -float(index),
+ }
+ for index in range(3)
+ ]
+ ),
+ paired_shard,
+ )
+ caption_hash = sha256_file(caption_shard)
+ paired_hash = sha256_file(paired_shard)
+ splits = ("train", "validation", "test")
+ accepted = []
+ for index, split in enumerate(splits):
+ accepted.append(
+ {
+ "schema_version": 1,
+ "candidate_index": index,
+ "paired_targetid": f"100{index}",
+ "caption_object_id": f"caption-{index}",
+ "coordinates_deg": {
+ "paired_ra": float(index),
+ "paired_dec": -float(index),
+ "caption_ra": float(index),
+ "caption_dec": -float(index),
+ },
+ "separation_arcsec": 0.0,
+ "source_group_id": f"group-{index}",
+ "split": split,
+ "acceptance_status": "accepted",
+ "caption": {
+ "original": captions[index],
+ "summary_original": summaries[index],
+ "status": "released_reused",
+ },
+ "caption_source": {
+ "repo_id": CAPTION_REPO_ID,
+ "revision": CAPTION_REVISION,
+ "shard": "data/captions.parquet",
+ "row": index,
+ "sha256": caption_hash,
+ "survey": "legacy_dr10",
+ },
+ "paired_source": {
+ "repo_id": PAIRED_REPO_ID,
+ "revision": PAIRED_REVISION,
+ "shard": "data/paired.parquet",
+ "row": index,
+ "sha256": paired_hash,
+ },
+ }
+ )
+
+ candidates = tmp_path / "accepted/candidates.jsonl"
+ _write_jsonl(candidates, accepted)
+ roster = tmp_path / "roster.csv.gz"
+ with gzip.open(roster, "wt", encoding="utf-8", newline="") as handle:
+ handle.write("candidate_index\n0\n1\n2\n")
+ roster_hash = sha256_file(roster)
+ accepted_manifest = tmp_path / "accepted/manifest.json"
+ _write_json(
+ accepted_manifest,
+ {
+ "outputs": {
+ "candidates": {
+ "path": "candidates.jsonl",
+ "row_count": 3,
+ "sha256": sha256_file(candidates),
+ }
+ },
+ "pins": {"roster_sha256": roster_hash},
+ "splits": {
+ "row_counts": {"train": 1, "validation": 1, "test": 1},
+ "group_count": 3,
+ "policy": "fixture-whole-group",
+ },
+ },
+ )
+ image_manifest = _raw_fixture(
+ tmp_path / "raw_images", accepted, IMAGE_ARRAY_CONTRACT, spectra=False
+ )
+ spectrum_manifest = _raw_fixture(
+ tmp_path / "raw_spectra", accepted, SPECTRUM_ARRAY_CONTRACT, spectra=True
+ )
+ spectrum_bundle = _spectrum_bundle(tmp_path / "spectrum_accepted", accepted, [0, 1, 2])
+ return HFDatasetConfig(
+ output_dir=tmp_path / "upload",
+ accepted_candidates_path=candidates,
+ accepted_manifest_path=accepted_manifest,
+ raw_images_manifest_path=image_manifest,
+ raw_spectra_manifest_path=spectrum_manifest,
+ caption_root=caption_root,
+ paired_root=paired_root,
+ spectrum_caption_bundle=spectrum_bundle,
+ roster_path=roster,
+ expected_count=3,
+ expected_split_counts={"train": 1, "validation": 1, "test": 1},
+ expected_roster_sha256=roster_hash,
+ link_mode="copy",
+ )
+
+
+def test_stages_portable_arrays_and_two_three_split_configs(tmp_path: Path) -> None:
+ config = _complete_fixture(tmp_path)
+ summary = prepare_hf_dataset(config)
+ output = config.output_dir
+
+ assert summary["configs"] == {
+ "image_caption": {"train": 1, "validation": 1, "test": 1},
+ "spectrum_caption": {"train": 1, "validation": 1, "test": 1},
+ }
+ assert (output / "assets/raw/images/image_array.npy").is_file()
+ assert (output / "assets/raw/spectra/spectrum_flux.npy").is_file()
+ assert not list(output.rglob("*embedding*.npy"))
+ roster = [
+ json.loads(line) for line in (output / "metadata/candidates.jsonl").read_text().splitlines()
+ ]
+ assert [row["candidate_index"] for row in roster] == [0, 1, 2]
+ assert all(row["raw_assets"]["array_index"] == row["candidate_index"] for row in roster)
+ assert all("image_embedding" not in row for row in roster)
+
+ image_features = Features.from_arrow_schema(
+ pq.read_schema(output / "data/image_caption/train.parquet")
+ )
+ spectrum_features = Features.from_arrow_schema(
+ pq.read_schema(output / "data/spectrum_caption/train.parquet")
+ )
+ assert isinstance(image_features["image"], Image)
+ assert image_features["image"].decode is False
+ assert isinstance(spectrum_features["spectrum_plot"], Image)
+ assert spectrum_features["spectrum_plot"].decode is False
+
+ spectrum_rows = {
+ split: pq.read_table(output / f"data/spectrum_caption/{split}.parquet").to_pylist()[0]
+ for split in ("train", "validation", "test")
+ }
+ assert all(row["caption_status"] == "accepted" for row in spectrum_rows.values())
+ assert all(row["review_status"] == "accepted" for row in spectrum_rows.values())
+ assert all(row["reviewed"] is True for row in spectrum_rows.values())
+ assert all(row["dataset_ready"] is True for row in spectrum_rows.values())
+ assert spectrum_rows["train"]["role"] == "alignment_training"
+ assert spectrum_rows["train"]["training_allowed"] is True
+ for split in ("validation", "test"):
+ assert spectrum_rows[split]["role"] == "held_out_evaluation"
+ assert spectrum_rows[split]["training_allowed"] is False
+ assert not (output / "data/spectrum_caption_inspection").exists()
+ assert "spectrum_caption_inspection" not in (output / "README.md").read_text()
+ assert (output / "MANIFEST.json").is_file()
+ assert (output / "SHA256SUMS").is_file()
+
+ for path in output.rglob("*"):
+ if path.is_file() and path.suffix in {".json", ".jsonl", ".md"}:
+ text = path.read_text(encoding="utf-8")
+ assert "/lus/" not in text
+ assert "_local/" not in text
+ assert "BaseMM_PRISM" not in text
+
+
+@pytest.mark.parametrize(
+ ("section", "field", "value", "index", "match"),
+ [
+ ("spectrum_caption", "status", "pending_review", 0, "not accepted"),
+ ("review", "status", "pending_review", 0, "lacks an accepted"),
+ ("training_policy", "dataset_ready", False, 0, "not dataset_ready"),
+ (
+ "training_policy",
+ "gradient_training_allowed",
+ True,
+ 1,
+ "does not match",
+ ),
+ ],
+)
+def test_rejects_invalid_spectrum_promotion(
+ tmp_path: Path,
+ section: str,
+ field: str,
+ value: object,
+ index: int,
+ match: str,
+) -> None:
+ config = _complete_fixture(tmp_path)
+ bundle = config.spectrum_caption_bundle
+ captions_path = bundle / "captions.jsonl"
+ rows = [json.loads(line) for line in captions_path.read_text().splitlines()]
+ rows[index][section][field] = value
+ _write_jsonl(captions_path, rows)
+
+ manifest_path = bundle / "manifest.json"
+ manifest = json.loads(manifest_path.read_text())
+ entry = manifest["outputs"]["captions.jsonl"]
+ entry["sha256"] = sha256_file(captions_path)
+ entry["size_bytes"] = captions_path.stat().st_size
+ _write_json(manifest_path, manifest)
+
+ with pytest.raises(PackagingValidationError, match=match):
+ prepare_hf_dataset(config)
+
+
+def test_rejects_source_group_split_leakage(tmp_path: Path) -> None:
+ config = _complete_fixture(tmp_path)
+ rows = [json.loads(line) for line in config.accepted_candidates_path.read_text().splitlines()]
+ rows[1]["source_group_id"] = rows[0]["source_group_id"]
+
+ with pytest.raises(PackagingValidationError, match="source groups cross"):
+ _validate_accepted_rows(
+ rows,
+ expected_count=3,
+ expected_split_counts={"train": 1, "validation": 1, "test": 1},
+ caption_repo_id=CAPTION_REPO_ID,
+ caption_revision=CAPTION_REVISION,
+ paired_repo_id=PAIRED_REPO_ID,
+ paired_revision=PAIRED_REVISION,
+ )
+
+
+def test_dataset_card_documents_three_split_spectrum_policy() -> None:
+ card = _dataset_card(split_counts={"train": 516, "validation": 65, "test": 65})
+
+ assert "config_name: image_caption" in card
+ assert "config_name: spectrum_caption" in card
+ assert "data/spectrum_caption/validation.parquet" in card
+ assert "data/spectrum_caption/test.parquet" in card
+ assert "config_name: spectrum_caption_inspection" not in card
+ assert "`reviewed=true`" in card
+ assert "`dataset_ready=true`" in card
diff --git a/tests/test_prepare_astroprism_spectrum_caption_pilot.py b/tests/test_prepare_astroprism_spectrum_caption_pilot.py
new file mode 100644
index 0000000..52022f9
--- /dev/null
+++ b/tests/test_prepare_astroprism_spectrum_caption_pilot.py
@@ -0,0 +1,234 @@
+from __future__ import annotations
+
+import json
+
+import pytest
+from tools.prepare_astroprism_spectrum_caption_pilot import (
+ COMPACT_PAYLOAD_SCHEMA_VERSION,
+ _compact_caption_payload,
+ _compact_policy_metadata,
+ _review_markdown,
+ _select_candidate_records,
+)
+
+
+def _line(
+ key: str,
+ *,
+ detection_class: str = "detected",
+ quality: str = "valid",
+ flags: list[str] | None = None,
+ significance: float = 5.678,
+) -> dict[str, object]:
+ return {
+ "key": key,
+ "label": key.replace("_", " "),
+ "nature": "emission",
+ "detection_class": detection_class,
+ "quality": quality,
+ "quality_flags": flags or [],
+ "significance_sigma": significance,
+ "observed_centroid_angstrom": 5000.06,
+ "equivalent_width_signed_rest_angstrom": 2.346,
+ "equivalent_width_uncertainty_angstrom": 0.456,
+ # Deliberately verbose audit-only fields that must not reach the model.
+ "source_reference": "raw_spectra:candidate_index=7",
+ "provenance": "spectral_measurement",
+ "continuum_windows_rest_angstrom": [[4900.0, 4950.0], [5050.0, 5100.0]],
+ }
+
+
+def _measurement(
+ *,
+ snr: float = 4.126,
+ continuum_ratio: float = 1.3334,
+ continuum_error: float = 0.0104,
+ dn4000_value: float = 1.5004,
+ dn4000_error: float = 0.1004,
+ lines: list[dict[str, object]] | None = None,
+) -> dict[str, object]:
+ return {
+ "quality": {"median_pixel_snr_rest_3850_7000": snr},
+ "continuum": {
+ "status": "measured",
+ "label": "red/rising",
+ "red_to_blue_ratio": continuum_ratio,
+ "red_to_blue_ratio_uncertainty": continuum_error,
+ },
+ "dn4000": {
+ "status": "measured",
+ "strength_label": "moderate",
+ "value": dn4000_value,
+ "uncertainty": dn4000_error,
+ },
+ "lines": lines or [],
+ }
+
+
+def test_compact_payload_schema_rounding_and_no_audit_envelope() -> None:
+ payload = _compact_caption_payload(
+ _measurement(lines=[_line("halpha_6565")]),
+ {"redshift": 0.12345678},
+ )
+
+ assert payload == {
+ "schema_version": COMPACT_PAYLOAD_SCHEMA_VERSION,
+ "type": "galaxy",
+ "z_cat": 0.123457,
+ "snr": 4.13,
+ "cont": {"class": "red/rising", "ratio": 1.333, "err": 0.01},
+ "dn4000": {"class": "moderate", "value": 1.5, "err": 0.1, "cautious": True},
+ "lines": [
+ {
+ "id": "spec.line.halpha_6565",
+ "name": "halpha 6565",
+ "nature": "emission",
+ "conf": "detected",
+ "sig": 5.68,
+ "obs_A": 5000.1,
+ "ew_rest_A": 2.35,
+ "ew_err_A": 0.46,
+ }
+ ],
+ }
+ serialized = json.dumps(payload, sort_keys=True)
+ for forbidden in (
+ "allowed_uses",
+ "continuum_windows_rest_angstrom",
+ "evidence_locator",
+ "provenance",
+ "required_student_inputs",
+ "restrictions",
+ "source_reference",
+ ):
+ assert forbidden not in serialized
+
+
+def test_compact_payload_excludes_unusable_and_low_snr_tentative_lines() -> None:
+ measurement = _measurement(
+ snr=1.999,
+ lines=[
+ _line("kept_detected"),
+ _line("suppressed_tentative", detection_class="tentative", significance=4.2),
+ _line("masked_detected", flags=["masked_pixels_in_line_window"]),
+ _line("unavailable", detection_class="unavailable", quality="unavailable"),
+ _line("not_detected", detection_class="not_detected"),
+ ],
+ )
+
+ payload = _compact_caption_payload(measurement, {"redshift": 0.1})
+ assert [line["id"] for line in payload["lines"]] == ["spec.line.kept_detected"]
+
+ policy = _compact_policy_metadata(measurement)
+ reasons = {item["fact_id"]: item["reason"] for item in policy["excluded_lines"]}
+ assert reasons == {
+ "spec.line.suppressed_tentative": "tentative_suppressed_low_snr",
+ "spec.line.masked_detected": "masked_line_window",
+ "spec.line.unavailable": "unavailable",
+ "spec.line.not_detected": "not_detected",
+ }
+
+
+def test_tentative_and_observed_frame_flags_are_compact_and_cautious() -> None:
+ measurement = _measurement(
+ snr=2.0,
+ lines=[
+ _line("tentative", detection_class="tentative", significance=3.5),
+ _line(
+ "flagged",
+ flags=[
+ "desi_arm_overlap_region",
+ "telluric_a_band_neighborhood",
+ "telluric_b_band_neighborhood",
+ "red_sky_residual_sensitive",
+ ],
+ ),
+ ],
+ )
+
+ lines = _compact_caption_payload(measurement, {"redshift": 0.2})["lines"]
+ assert lines[0]["cautious"] is True
+ assert "flags" not in lines[0]
+ assert lines[1]["cautious"] is True
+ assert lines[1]["flags"] == ["arm_overlap", "telluric", "red_sky"]
+
+
+def test_weak_continuum_and_dn4000_significance_are_cautious() -> None:
+ payload = _compact_caption_payload(
+ _measurement(
+ continuum_ratio=1.05,
+ continuum_error=0.02,
+ dn4000_value=1.1,
+ dn4000_error=0.06,
+ ),
+ {"redshift": 0.3},
+ )
+
+ assert payload["cont"]["cautious"] is True
+ assert payload["dn4000"]["cautious"] is True
+
+
+@pytest.mark.parametrize(
+ ("value", "error", "cautious"),
+ [
+ (1.403, 0.145, True),
+ (1.845, 0.299, True),
+ (1.435, 0.170, True),
+ (1.328, 0.071, True),
+ (1.850, 0.049, False),
+ ],
+)
+def test_dn4000_class_boundary_crossing_is_cautious(
+ value: float, error: float, cautious: bool
+) -> None:
+ payload = _compact_caption_payload(
+ _measurement(dn4000_value=value, dn4000_error=error),
+ {"redshift": 0.3},
+ )
+
+ assert payload["dn4000"].get("cautious", False) is cautious
+
+
+def test_candidate_selector_supports_train_limit_and_explicit_order() -> None:
+ rows = [
+ {"candidate_index": 0, "split": "test"},
+ {"candidate_index": 1, "split": "train"},
+ {"candidate_index": 2, "split": "validation"},
+ {"candidate_index": 3, "split": "train"},
+ ]
+
+ selected = _select_candidate_records(rows, requested_indices=None, train_limit=2)
+ assert [row["candidate_index"] for row in selected] == [1, 3]
+ selected = _select_candidate_records(rows, requested_indices=[3, 0], train_limit=None)
+ assert [row["candidate_index"] for row in selected] == [3, 0]
+
+ with pytest.raises(ValueError, match="mutually exclusive"):
+ _select_candidate_records(rows, requested_indices=[1], train_limit=1)
+ with pytest.raises(ValueError, match="must be positive"):
+ _select_candidate_records(rows, requested_indices=None, train_limit=0)
+ with pytest.raises(ValueError, match="missing"):
+ _select_candidate_records(rows, requested_indices=[99], train_limit=None)
+
+
+def test_review_markdown_handles_unavailable_continuum_ratio() -> None:
+ measurement = _measurement()
+ measurement["catalog_redshift"] = 0.4
+ measurement["continuum"] = {
+ "status": "unavailable",
+ "label": "unavailable",
+ "red_to_blue_ratio": None,
+ "red_to_blue_ratio_uncertainty": None,
+ }
+ record = {
+ "candidate_index": 46,
+ "split": "train",
+ "paired_targetid": "123",
+ "training_policy": {"role": "training_candidate_pending_review"},
+ "spectrum_caption": {"original": "A grounded caption."},
+ "teacher_reference": {"measurement_summary": measurement},
+ "audit_sidecar": {"claims": []},
+ }
+
+ review = _review_markdown([record])
+
+ assert "Continuum: `unavailable`; red/blue f_lambda ratio `— ± —`." in review
diff --git a/tests/test_promote_astroprism_spectrum_captions.py b/tests/test_promote_astroprism_spectrum_captions.py
new file mode 100644
index 0000000..41e47c7
--- /dev/null
+++ b/tests/test_promote_astroprism_spectrum_captions.py
@@ -0,0 +1,335 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import pytest
+from tools.promote_astroprism_spectrum_captions import (
+ PromotionConfig,
+ PromotionValidationError,
+ promote_spectrum_captions,
+ sha256_file,
+)
+
+pytestmark = [pytest.mark.unit, pytest.mark.multimodal]
+
+SPLIT_COUNTS = {"train": 2, "validation": 1, "test": 1}
+
+
+def _write_json(path: Path, value: object) -> None:
+ path.parent.mkdir(parents=True, exist_ok=True)
+ path.write_text(
+ json.dumps(value, indent=2, sort_keys=True, ensure_ascii=False) + "\n",
+ encoding="utf-8",
+ )
+
+
+def _write_jsonl(path: Path, rows: list[dict]) -> None:
+ path.parent.mkdir(parents=True, exist_ok=True)
+ path.write_text(
+ "".join(json.dumps(row, sort_keys=True, ensure_ascii=False) + "\n" for row in rows),
+ encoding="utf-8",
+ )
+
+
+def _read_jsonl(path: Path) -> list[dict]:
+ return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines()]
+
+
+def _accepted_fixture(root: Path) -> tuple[Path, Path, list[dict]]:
+ splits = ("train", "validation", "train", "test")
+ rows = [
+ {
+ "candidate_index": index,
+ "paired_targetid": f"target-{index}",
+ "caption_object_id": f"caption-{index}",
+ "source_group_id": f"group-{index}",
+ "split": split,
+ "acceptance_status": "accepted",
+ }
+ for index, split in enumerate(splits)
+ ]
+ candidates = root / "accepted/candidates.jsonl"
+ _write_jsonl(candidates, rows)
+ manifest = root / "accepted/manifest.json"
+ _write_json(
+ manifest,
+ {
+ "artifact_type": "astroprism_stage1a_candidates",
+ "roster": {"row_count": 4, "sha256": "fixture-roster-sha256"},
+ "outputs": {
+ "candidates": {
+ "path": "candidates.jsonl",
+ "row_count": 4,
+ "size_bytes": candidates.stat().st_size,
+ "sha256": sha256_file(candidates),
+ }
+ },
+ "splits": {"row_counts": SPLIT_COUNTS},
+ },
+ )
+ return candidates, manifest, rows
+
+
+def _bundle_fixture(
+ root: Path, accepted: list[dict], indices: list[int], *, heldout: bool
+) -> Path:
+ root.mkdir(parents=True)
+ records: list[dict] = []
+ requests: list[dict] = []
+ audits: list[dict] = []
+ for index in indices:
+ identity = accepted[index]
+ teacher_reference = {
+ "facts": [{"fact_id": "catalog.redshift", "value": 0.1 + index}],
+ "catalog_context": {
+ "redshift": 0.1 + index,
+ "source": {
+ "path": f"/tmp/astroprism-fixture/cache/data/shard-{index}.parquet",
+ "shard": f"data/shard-{index}.parquet",
+ "row": index,
+ },
+ },
+ }
+ model_digest = f"model-{index}"
+ payload_digest = f"payload-{index}"
+ records.append(
+ {
+ "schema_version": 1,
+ "artifact_type": "astroprism_spectrum_caption_candidate",
+ **identity,
+ "student_input": {
+ "modalities": ["spectrum"],
+ "raw_manifest": "/tmp/astroprism-fixture/raw_spectra/manifest.json",
+ "array_index": index,
+ "wavelength_frame": "observed_vacuum",
+ "spectrum_length": 7781,
+ },
+ "teacher_reference": teacher_reference,
+ "spectrum_caption": {
+ "original": f"Grounded spectrum caption {index}.",
+ "revised": None,
+ "status": "pending_review",
+ },
+ "audit_sidecar": {
+ "claims": [
+ {"text": f"caption {index}", "fact_ids": ["catalog.redshift"]}
+ ],
+ "spectral_grounding_status": "candidate_unverified",
+ },
+ "generation": {
+ "method": "fixture-generation",
+ "model_input_sha256": model_digest,
+ "compact_payload_sha256": payload_digest,
+ },
+ "training_policy": {
+ "eligible_for_training": False,
+ "eligible_after_review": not heldout,
+ "role": (
+ "held_out_inspection_only"
+ if heldout
+ else "training_candidate_pending_review"
+ ),
+ },
+ }
+ )
+ requests.append(
+ {
+ "candidate_index": index,
+ "paired_targetid": identity["paired_targetid"],
+ "split": identity["split"],
+ "model_input": {"fixture": index},
+ "model_input_sha256": model_digest,
+ "compact_payload_sha256": payload_digest,
+ }
+ )
+ audits.append(
+ {
+ "candidate_index": index,
+ "paired_targetid": identity["paired_targetid"],
+ "caption_object_id": identity["caption_object_id"],
+ "source_group_id": identity["source_group_id"],
+ "split": identity["split"],
+ "model_input_sha256": model_digest,
+ "compact_payload_sha256": payload_digest,
+ "teacher_reference": teacher_reference,
+ "training_policy": {"eligible_for_training": False},
+ }
+ )
+ plot = root / f"plots/candidate_{index:04d}.png"
+ plot.parent.mkdir(exist_ok=True)
+ plot.write_bytes(b"\x89PNG\r\n\x1a\n" + f"fixture-{index}".encode())
+
+ _write_jsonl(root / "captions.jsonl", records)
+ _write_jsonl(root / "generation_requests.jsonl", requests)
+ _write_jsonl(root / "request_audit.jsonl", audits)
+ (root / "review.md").write_text("Pending review fixture.\n", encoding="utf-8")
+ outputs = {
+ path.relative_to(root).as_posix(): {
+ "size_bytes": path.stat().st_size,
+ "sha256": sha256_file(path),
+ }
+ for path in sorted(item for item in root.rglob("*") if item.is_file())
+ }
+ split_counts = {
+ split: sum(accepted[index]["split"] == split for index in indices)
+ for split in sorted({accepted[index]["split"] for index in indices})
+ }
+ _write_json(
+ root / "manifest.json",
+ {
+ "schema_version": 1,
+ "artifact_type": "astroprism_spectrum_caption_pilot",
+ "role": "review_only_not_training_eligible",
+ "row_count": len(indices),
+ "candidate_indices": indices,
+ "held_out_candidate_indices": indices if heldout else [],
+ "split_counts": split_counts,
+ "implementation": {"fixture": True},
+ "sources": {"paper": {"url": "https://example.test/paper"}},
+ "outputs": outputs,
+ },
+ )
+ return root
+
+
+def _complete_fixture(tmp_path: Path, *, link_mode: str = "copy") -> PromotionConfig:
+ candidates, accepted_manifest, accepted = _accepted_fixture(tmp_path)
+ train = _bundle_fixture(tmp_path / "train", accepted, [0, 2], heldout=False)
+ heldout = _bundle_fixture(tmp_path / "heldout", accepted, [1, 3], heldout=True)
+ return PromotionConfig(
+ output_dir=tmp_path / "promoted",
+ accepted_candidates_path=candidates,
+ accepted_manifest_path=accepted_manifest,
+ train_bundle=train,
+ heldout_bundle=heldout,
+ expected_count=4,
+ expected_split_counts=SPLIT_COUNTS,
+ link_mode=link_mode,
+ )
+
+
+def test_promotes_all_splits_and_preserves_portable_provenance(tmp_path: Path) -> None:
+ config = _complete_fixture(tmp_path)
+ source_hashes = {
+ path: sha256_file(path)
+ for root in (config.train_bundle, config.heldout_bundle)
+ for path in Path(root).rglob("*")
+ if path.is_file()
+ }
+
+ summary = promote_spectrum_captions(config)
+ output = config.output_dir
+ records = _read_jsonl(output / "captions.jsonl")
+
+ assert summary["candidate_count"] == 4
+ assert summary["gradient_training_count"] == 2
+ assert summary["evaluation_count"] == 2
+ assert [row["candidate_index"] for row in records] == [0, 1, 2, 3]
+ assert all(row["spectrum_caption"]["status"] == "accepted" for row in records)
+ assert all(row["review"]["status"] == "accepted" for row in records)
+ assert all(row["review"]["decision_source"] == "user_instruction" for row in records)
+ assert all(row["training_policy"]["dataset_ready"] is True for row in records)
+
+ expected_roles = {
+ "train": "alignment_training",
+ "validation": "alignment_validation",
+ "test": "alignment_test",
+ }
+ for row in records:
+ is_train = row["split"] == "train"
+ policy = row["training_policy"]
+ assert policy["gradient_training_allowed"] is is_train
+ assert policy["eligible_for_training"] is is_train
+ assert policy["eligible_for_evaluation"] is (not is_train)
+ assert policy["role"] == expected_roles[row["split"]]
+ assert row["student_input"]["raw_manifest"] == "../raw_spectra/manifest.json"
+ source = row["teacher_reference"]["catalog_context"]["source"]
+ assert source["path"] == source["shard"]
+
+ assert records[0]["generation"]["method"] == "fixture-generation"
+ assert records[0]["audit_sidecar"]["claims"][0]["fact_ids"] == ["catalog.redshift"]
+ assert _read_jsonl(output / "generation_requests.jsonl")[3]["model_input"] == {
+ "fixture": 3
+ }
+ assert _read_jsonl(output / "request_audit.jsonl")[1]["teacher_reference"][
+ "catalog_context"
+ ]["source"]["path"] == "data/shard-1.parquet"
+
+ serialized = "\n".join(
+ path.read_text(encoding="utf-8")
+ for path in output.rglob("*")
+ if path.is_file() and path.suffix in {".json", ".jsonl", ".md"}
+ )
+ assert "/lus/" not in serialized
+ for path, digest in source_hashes.items():
+ assert sha256_file(path) == digest
+
+ manifest = json.loads((output / "manifest.json").read_text(encoding="utf-8"))
+ assert manifest["artifact_type"] == "astroprism_spectrum_captions_accepted"
+ assert manifest["row_count"] == 4
+ assert manifest["split_counts"] == SPLIT_COUNTS
+ assert manifest["roster"] == {"row_count": 4, "sha256": "fixture-roster-sha256"}
+ assert manifest["review"]["decision_source"] == "user_instruction"
+ assert manifest["outputs"]["captions.jsonl"]["sha256"] == sha256_file(
+ output / "captions.jsonl"
+ )
+ checksums = (output / "SHA256SUMS").read_text(encoding="utf-8")
+ assert " manifest.json\n" in checksums
+ assert " captions.jsonl\n" in checksums
+ assert "SHA256SUMS" not in checksums
+
+
+def test_hardlink_mode_links_validated_plots(tmp_path: Path) -> None:
+ config = _complete_fixture(tmp_path, link_mode="hardlink")
+
+ promote_spectrum_captions(config)
+
+ source = config.train_bundle / "plots/candidate_0000.png"
+ promoted = config.output_dir / "plots/candidate_0000.png"
+ assert source.stat().st_ino == promoted.stat().st_ino
+
+
+def test_rejects_source_group_leakage_before_writing_output(tmp_path: Path) -> None:
+ config = _complete_fixture(tmp_path)
+ rows = _read_jsonl(config.accepted_candidates_path)
+ rows[1]["source_group_id"] = rows[0]["source_group_id"]
+ _write_jsonl(config.accepted_candidates_path, rows)
+ manifest = json.loads(config.accepted_manifest_path.read_text(encoding="utf-8"))
+ manifest["outputs"]["candidates"]["size_bytes"] = config.accepted_candidates_path.stat().st_size
+ manifest["outputs"]["candidates"]["sha256"] = sha256_file(
+ config.accepted_candidates_path
+ )
+ _write_json(config.accepted_manifest_path, manifest)
+
+ with pytest.raises(PromotionValidationError, match="source groups cross splits"):
+ promote_spectrum_captions(config)
+
+ assert not config.output_dir.exists()
+ assert not list(tmp_path.glob(".promoted.tmp-*"))
+
+
+def test_rejects_plot_checksum_mismatch_without_mutating_sources(tmp_path: Path) -> None:
+ config = _complete_fixture(tmp_path)
+ plot = config.heldout_bundle / "plots/candidate_0001.png"
+ plot.write_bytes(plot.read_bytes() + b"tampered")
+
+ with pytest.raises(PromotionValidationError, match="mismatch"):
+ promote_spectrum_captions(config)
+
+ assert not config.output_dir.exists()
+
+
+def test_overwrite_is_atomic_when_revalidation_fails(tmp_path: Path) -> None:
+ config = _complete_fixture(tmp_path)
+ promote_spectrum_captions(config)
+ original_digest = sha256_file(config.output_dir / "captions.jsonl")
+ plot = config.train_bundle / "plots/candidate_0000.png"
+ plot.write_bytes(plot.read_bytes() + b"tampered")
+
+ overwrite = PromotionConfig(**{**config.__dict__, "overwrite": True})
+ with pytest.raises(PromotionValidationError, match="mismatch"):
+ promote_spectrum_captions(overwrite)
+
+ assert sha256_file(config.output_dir / "captions.jsonl") == original_digest
+ assert not list(tmp_path.glob(".promoted.tmp-*"))
diff --git a/tests/test_zone_a_aion_training_boundary.py b/tests/test_zone_a_aion_training_boundary.py
new file mode 100644
index 0000000..840bf00
--- /dev/null
+++ b/tests/test_zone_a_aion_training_boundary.py
@@ -0,0 +1,154 @@
+from types import SimpleNamespace
+
+import pytest
+import src.training.trainer_zone_a as trainer_module
+import torch
+import torch.nn as nn
+from src.config import TrainingConfig
+from src.training.trainer_zone_a import (
+ ZoneATrainer,
+ _configure_encoder_for_training,
+ _labels_for_batch,
+ _refresh_encoder_frozen_cache,
+)
+
+
+class _ProtectedEncoder(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.codec = nn.Linear(2, 2).float()
+ self.shared = nn.Linear(2, 2).float()
+ self.dtype_casts = []
+
+ def to(self, *args, **kwargs):
+ if "dtype" in kwargs:
+ self.dtype_casts.append(kwargs["dtype"])
+ return super().to(*args, **kwargs)
+
+ def set_shared_encoder_trainable(self, trainable):
+ self.codec.float()
+ self.codec.eval()
+ for parameter in self.codec.parameters():
+ parameter.requires_grad_(False)
+ for parameter in self.shared.parameters():
+ parameter.requires_grad_(trainable)
+
+
+class _Model(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.config = SimpleNamespace(modalities=["image"], is_interleaved_qa=False)
+ self.backbone = None
+ self.encoders = nn.ModuleDict({"image": _ProtectedEncoder()})
+ self.projectors = nn.ModuleDict({"image": nn.Linear(2, 2)})
+ self._encoder_frozen_cache = {"image": False}
+ self.refresh_calls = 0
+
+ def refresh_encoder_frozen_cache(self):
+ self.refresh_calls += 1
+ self._encoder_frozen_cache = {
+ name: not any(parameter.requires_grad for parameter in encoder.parameters())
+ for name, encoder in self.encoders.items()
+ }
+ return self._encoder_frozen_cache
+
+
+class _FakeAccelerator:
+ def __init__(self, *args, **kwargs):
+ self.state = SimpleNamespace(deepspeed_plugin=None)
+ self.is_main_process = False
+ self.device = torch.device("cpu")
+ # The scheduler expresses warmup/decay in optimizer steps by scaling
+ # with the world size; single-process stub means 1.
+ self.num_processes = 1
+
+ def print(self, *args, **kwargs):
+ return None
+
+ def prepare(self, *items):
+ return items if len(items) > 1 else items[0]
+
+
+class _Dataset:
+ active_modalities = ["image"]
+
+
+def test_capability_aware_encoder_configuration_preserves_codec_fp32():
+ encoder = _ProtectedEncoder()
+
+ _configure_encoder_for_training(
+ encoder,
+ trainable=True,
+ target_dtype=torch.bfloat16,
+ )
+
+ assert encoder.dtype_casts == []
+ assert encoder.codec.weight.dtype == torch.float32
+ assert all(not parameter.requires_grad for parameter in encoder.codec.parameters())
+ assert all(parameter.requires_grad for parameter in encoder.shared.parameters())
+
+ _configure_encoder_for_training(
+ encoder,
+ trainable=False,
+ target_dtype=torch.bfloat16,
+ )
+ assert all(not parameter.requires_grad for parameter in encoder.parameters())
+
+
+@pytest.mark.parametrize("freeze_vit", [True, False])
+def test_zone_a_optimizer_respects_shared_encoder_boundary(monkeypatch, freeze_vit):
+ monkeypatch.setattr(trainer_module, "Accelerator", _FakeAccelerator)
+ model = _Model()
+ config = TrainingConfig(
+ max_steps=1,
+ warmup_steps=0,
+ wandb_project=None,
+ device="cpu",
+ freeze_vit=freeze_vit,
+ )
+ loader = SimpleNamespace(dataset=_Dataset())
+
+ trainer = ZoneATrainer(model, config, loader)
+
+ assert model.refresh_calls == 1
+ assert model.encoders["image"].codec.weight.dtype == torch.float32
+ assert all(
+ not parameter.requires_grad for parameter in model.encoders["image"].codec.parameters()
+ )
+
+ trainable_names = {
+ name for name, parameter in model.named_parameters() if parameter.requires_grad
+ }
+ expected = {"projectors.image.weight", "projectors.image.bias"}
+ if not freeze_vit:
+ expected |= {"encoders.image.shared.weight", "encoders.image.shared.bias"}
+ assert trainable_names == expected
+
+ vit_groups = [group for group in trainer.optimizer.param_groups if group["name"] == "vit"]
+ if freeze_vit:
+ assert vit_groups == []
+ assert model._encoder_frozen_cache == {"image": True}
+ else:
+ assert len(vit_groups) == 1
+ optimizer_ids = {id(parameter) for parameter in vit_groups[0]["params"]}
+ shared_ids = {id(parameter) for parameter in model.encoders["image"].shared.parameters()}
+ codec_ids = {id(parameter) for parameter in model.encoders["image"].codec.parameters()}
+ assert optimizer_ids == shared_ids
+ assert optimizer_ids.isdisjoint(codec_ids)
+ assert model._encoder_frozen_cache == {"image": False}
+
+
+def test_encoder_cache_fallback_and_assistant_only_labels():
+ encoder = _ProtectedEncoder()
+ model = SimpleNamespace(
+ encoders={"image": encoder},
+ _encoder_frozen_cache={"image": True},
+ )
+ _configure_encoder_for_training(encoder, trainable=True, target_dtype=torch.float32)
+ _refresh_encoder_frozen_cache(model)
+ assert model._encoder_frozen_cache == {"image": False}
+
+ text = torch.tensor([[1, 2, 3]])
+ labels = torch.tensor([[-100, -100, 3]])
+ assert _labels_for_batch({"text": text, "labels": labels}) is labels
+ assert _labels_for_batch({"text": text}) is text
diff --git a/tools/accept_astroprism_stage1a.py b/tools/accept_astroprism_stage1a.py
new file mode 100755
index 0000000..51f6bd0
--- /dev/null
+++ b/tools/accept_astroprism_stage1a.py
@@ -0,0 +1,496 @@
+#!/usr/bin/env python3
+"""Promote reviewed AstroPRISM Stage 1A candidates into a training artifact.
+
+The source candidate artifact remains immutable. This tool copies its 646
+records in roster order, marks unchanged captions as released, groups nearby
+cutouts with the deterministic 'skycc60-v1' policy, and assigns exact
+516/65/65 train/validation/test splits without separating a group.
+"""
+
+from __future__ import annotations
+
+import argparse
+import copy
+import hashlib
+import json
+import math
+import os
+import shutil
+import tempfile
+import uuid
+from collections import Counter, defaultdict
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+
+EXPECTED_CANDIDATES = 646
+OVERLAP_RADIUS_ARCSEC = 60.0
+GROUP_POLICY = "skycc60-v1"
+SPLIT_POLICY = "sha256-greedy-whole-group-v1"
+SPLIT_RANK_SEED = "astroprism-stage1a-split-v1"
+SPLIT_COUNTS = {"train": 516, "validation": 65, "test": 65}
+ACCEPTANCE_BASIS = "user_directive_assume_reviewed_and_accepted"
+
+
+class AcceptanceValidationError(ValueError):
+ """Raised when promotion would violate the accepted-artifact contract."""
+
+
+@dataclass(frozen=True)
+class AcceptanceConfig:
+ """Inputs for deterministic Stage 1A candidate promotion."""
+
+ source_manifest_path: Path
+ output_dir: Path
+ overwrite: bool = False
+ expected_count: int = EXPECTED_CANDIDATES
+ split_counts: dict[str, int] | None = None
+
+
+def sha256_file(path: Path, chunk_size: int = 8 * 1024 * 1024) -> str:
+ digest = hashlib.sha256()
+ with path.open("rb") as handle:
+ while chunk := handle.read(chunk_size):
+ digest.update(chunk)
+ return digest.hexdigest()
+
+
+def _read_json(path: Path, *, kind: str) -> dict[str, Any]:
+ if not path.is_file():
+ raise AcceptanceValidationError(f"{kind} does not exist: {path}")
+ try:
+ value = json.loads(path.read_text(encoding="utf-8"))
+ except (OSError, json.JSONDecodeError) as exc:
+ raise AcceptanceValidationError(f"cannot read {kind} {path}: {exc}") from exc
+ if not isinstance(value, dict):
+ raise AcceptanceValidationError(f"{kind} must contain a JSON object: {path}")
+ return value
+
+
+def _resolve_output(manifest_path: Path, output: dict[str, Any], *, name: str) -> Path:
+ value = output.get("path")
+ if not isinstance(value, str) or not value:
+ raise AcceptanceValidationError(f"source manifest output {name!r} has no path")
+ path = Path(value)
+ return path if path.is_absolute() else manifest_path.parent / path
+
+
+def _read_jsonl(path: Path) -> list[dict[str, Any]]:
+ if not path.is_file():
+ raise AcceptanceValidationError(f"source candidates do not exist: {path}")
+ records: list[dict[str, Any]] = []
+ try:
+ with path.open("r", encoding="utf-8") as handle:
+ for line_number, line in enumerate(handle, start=1):
+ if not line.strip():
+ raise AcceptanceValidationError(
+ f"source candidates contain a blank row at line {line_number}"
+ )
+ value = json.loads(line)
+ if not isinstance(value, dict):
+ raise AcceptanceValidationError(
+ f"source candidate line {line_number} must be an object"
+ )
+ records.append(value)
+ except json.JSONDecodeError as exc:
+ raise AcceptanceValidationError(
+ f"cannot parse source candidates {path} at line {exc.lineno}: {exc.msg}"
+ ) from exc
+ return records
+
+
+def _paired_coordinate(record: dict[str, Any], *, index: int) -> tuple[float, float]:
+ coordinates = record.get("coordinates_deg")
+ if not isinstance(coordinates, dict):
+ raise AcceptanceValidationError(f"candidate {index} has no coordinates_deg object")
+ try:
+ ra = float(coordinates["paired_ra"])
+ dec = float(coordinates["paired_dec"])
+ except (KeyError, TypeError, ValueError) as exc:
+ raise AcceptanceValidationError(
+ f"candidate {index} has invalid paired_ra/paired_dec"
+ ) from exc
+ if not math.isfinite(ra) or not math.isfinite(dec):
+ raise AcceptanceValidationError(f"candidate {index} coordinates must be finite")
+ if not 0.0 <= ra < 360.0 or not -90.0 <= dec <= 90.0:
+ raise AcceptanceValidationError(f"candidate {index} coordinates are out of range")
+ return math.radians(ra), math.radians(dec)
+
+
+def _separation_arcsec(first: tuple[float, float], second: tuple[float, float]) -> float:
+ ra1, dec1 = first
+ ra2, dec2 = second
+ haversine = (
+ math.sin((dec2 - dec1) / 2.0) ** 2
+ + math.cos(dec1) * math.cos(dec2) * math.sin((ra2 - ra1) / 2.0) ** 2
+ )
+ return 2.0 * math.asin(min(1.0, math.sqrt(haversine))) * 180.0 / math.pi * 3600.0
+
+
+def _validate_source_records(
+ records: list[dict[str, Any]], *, expected_count: int
+) -> list[tuple[float, float]]:
+ if len(records) != expected_count:
+ raise AcceptanceValidationError(
+ f"source artifact must contain exactly {expected_count} rows, got {len(records)}"
+ )
+ paired_ids: set[str] = set()
+ caption_ids: set[str] = set()
+ coordinates: list[tuple[float, float]] = []
+ for index, record in enumerate(records):
+ if record.get("candidate_index") != index:
+ raise AcceptanceValidationError(
+ f"candidate row {index} has candidate_index={record.get('candidate_index')!r}; "
+ "roster order must be preserved"
+ )
+ paired_id = str(record.get("paired_targetid", "")).strip()
+ caption_id = str(record.get("caption_object_id", "")).strip()
+ if not paired_id or paired_id in paired_ids:
+ raise AcceptanceValidationError(
+ f"candidate {index} has an empty or duplicate paired_targetid"
+ )
+ if not caption_id or caption_id in caption_ids:
+ raise AcceptanceValidationError(
+ f"candidate {index} has an empty or duplicate caption_object_id"
+ )
+ paired_ids.add(paired_id)
+ caption_ids.add(caption_id)
+
+ caption = record.get("caption")
+ if not isinstance(caption, dict):
+ raise AcceptanceValidationError(f"candidate {index} has no caption object")
+ original = caption.get("original")
+ if not isinstance(original, str) or not original.strip():
+ raise AcceptanceValidationError(f"candidate {index} has no non-empty caption.original")
+ if str(record.get("acceptance_status", "")).strip() != "pending_review":
+ raise AcceptanceValidationError(
+ f"candidate {index} is not pending_review; refusing to re-promote it"
+ )
+ if str(caption.get("status", "")).strip() != "unreviewed":
+ raise AcceptanceValidationError(
+ f"candidate {index} caption is not unreviewed; refusing to re-promote it"
+ )
+ if record.get("source_group_id") is not None or record.get("split") is not None:
+ raise AcceptanceValidationError(
+ f"candidate {index} already has a source group or split"
+ )
+ image_embedding = record.get("image_embedding")
+ if not isinstance(image_embedding, dict) or not image_embedding.get("file"):
+ raise AcceptanceValidationError(
+ f"candidate {index} has no image_embedding.file provenance"
+ )
+ coordinates.append(_paired_coordinate(record, index=index))
+ return coordinates
+
+
+def _source_groups(
+ records: list[dict[str, Any]], coordinates: list[tuple[float, float]]
+) -> list[tuple[str, tuple[int, ...]]]:
+ parents = list(range(len(records)))
+
+ def find(index: int) -> int:
+ while parents[index] != index:
+ parents[index] = parents[parents[index]]
+ index = parents[index]
+ return index
+
+ def union(first: int, second: int) -> None:
+ first_root = find(first)
+ second_root = find(second)
+ if first_root != second_root:
+ parents[max(first_root, second_root)] = min(first_root, second_root)
+
+ for first in range(len(records)):
+ for second in range(first + 1, len(records)):
+ if _separation_arcsec(coordinates[first], coordinates[second]) <= OVERLAP_RADIUS_ARCSEC:
+ union(first, second)
+
+ components: dict[int, list[int]] = defaultdict(list)
+ for index in range(len(records)):
+ components[find(index)].append(index)
+
+ groups: list[tuple[str, tuple[int, ...]]] = []
+ for members in components.values():
+ identities = sorted(f"desi:{records[index]['paired_targetid']}" for index in members)
+ payload = ("\n".join(identities) + "\n").encode("utf-8")
+ group_id = f"{GROUP_POLICY}:{hashlib.sha256(payload).hexdigest()[:24]}"
+ groups.append((group_id, tuple(sorted(members))))
+ return groups
+
+
+def _assign_splits(
+ groups: list[tuple[str, tuple[int, ...]]], split_counts: dict[str, int]
+) -> dict[str, str]:
+ ranked = sorted(
+ groups,
+ key=lambda item: (
+ hashlib.sha256(f"{SPLIT_RANK_SEED}\0{item[0]}".encode()).digest(),
+ item[0],
+ ),
+ )
+ remaining = ranked
+ assignments: dict[str, str] = {}
+ for split in ("test", "validation"):
+ capacity = split_counts[split]
+ deferred: list[tuple[str, tuple[int, ...]]] = []
+ for group_id, members in remaining:
+ if len(members) <= capacity:
+ assignments[group_id] = split
+ capacity -= len(members)
+ else:
+ deferred.append((group_id, members))
+ if capacity:
+ raise AcceptanceValidationError(
+ f"whole-group allocation could not fill {split!r}; {capacity} rows remain"
+ )
+ remaining = deferred
+
+ for group_id, _members in remaining:
+ assignments[group_id] = "train"
+ train_rows = sum(
+ len(members) for group_id, members in groups if assignments[group_id] == "train"
+ )
+ if train_rows != split_counts["train"]:
+ raise AcceptanceValidationError(
+ f"whole-group allocation produced {train_rows} train rows, "
+ f"expected {split_counts['train']}"
+ )
+ return assignments
+
+
+def _mapping_sha256(records: list[dict[str, Any]]) -> str:
+ payload = "".join(
+ f"{index}\t{record['source_group_id']}\t{record['split']}\n"
+ for index, record in enumerate(records)
+ )
+ return hashlib.sha256(payload.encode("utf-8")).hexdigest()
+
+
+def _write_jsonl(path: Path, records: list[dict[str, Any]]) -> None:
+ with path.open("w", encoding="utf-8", newline="\n") as handle:
+ for record in records:
+ handle.write(json.dumps(record, ensure_ascii=False, sort_keys=True, allow_nan=False))
+ handle.write("\n")
+
+
+def _write_json(path: Path, value: dict[str, Any]) -> None:
+ path.write_text(
+ json.dumps(value, indent=2, sort_keys=True, ensure_ascii=False, allow_nan=False) + "\n",
+ encoding="utf-8",
+ )
+
+
+def _commit_staging(staging: Path, output: Path, *, overwrite: bool) -> None:
+ if not output.exists() and not output.is_symlink():
+ os.replace(staging, output)
+ return
+ if not overwrite:
+ raise FileExistsError(f"output already exists: {output}")
+ backup = output.with_name(f".{output.name}.backup-{uuid.uuid4().hex}")
+ os.replace(output, backup)
+ try:
+ os.replace(staging, output)
+ except BaseException:
+ os.replace(backup, output)
+ raise
+ if backup.is_dir() and not backup.is_symlink():
+ shutil.rmtree(backup)
+ else:
+ backup.unlink(missing_ok=True)
+
+
+def accept_stage1a(config: AcceptanceConfig) -> dict[str, Any]:
+ """Write a separate accepted artifact and return its reproducibility summary."""
+
+ split_counts = dict(config.split_counts or SPLIT_COUNTS)
+ if set(split_counts) != {"train", "validation", "test"}:
+ raise ValueError("split_counts must define train, validation, and test")
+ if any(not isinstance(count, int) or count < 0 for count in split_counts.values()):
+ raise ValueError("split counts must be nonnegative integers")
+ if sum(split_counts.values()) != config.expected_count:
+ raise ValueError("split counts must sum to expected_count")
+
+ source_manifest_path = Path(config.source_manifest_path).resolve()
+ output_dir = Path(config.output_dir).resolve()
+ if (output_dir.exists() or output_dir.is_symlink()) and not config.overwrite:
+ raise FileExistsError(f"output already exists: {output_dir}")
+ source_manifest = _read_json(source_manifest_path, kind="source manifest")
+ if source_manifest.get("artifact_type") not in {
+ "astroprism_stage1a_candidates",
+ "astroprism_aion_feature_parity_reference",
+ }:
+ raise AcceptanceValidationError("source manifest has an unsupported artifact_type")
+ outputs = source_manifest.get("outputs")
+ if not isinstance(outputs, dict) or not isinstance(outputs.get("candidates"), dict):
+ raise AcceptanceValidationError("source manifest is missing outputs.candidates")
+ source_candidates_path = _resolve_output(
+ source_manifest_path, outputs["candidates"], name="candidates"
+ ).resolve()
+ source_candidates_sha256 = sha256_file(source_candidates_path)
+ expected_sha256 = outputs["candidates"].get("sha256")
+ if expected_sha256 and source_candidates_sha256 != expected_sha256:
+ raise AcceptanceValidationError("source candidate SHA-256 does not match its manifest")
+ records = _read_jsonl(source_candidates_path)
+ declared_count = outputs["candidates"].get("row_count")
+ if declared_count is not None and declared_count != len(records):
+ raise AcceptanceValidationError("source candidate row count does not match its manifest")
+ coordinates = _validate_source_records(records, expected_count=config.expected_count)
+
+ image_output = outputs.get("image_embeddings")
+ if not isinstance(image_output, dict):
+ raise AcceptanceValidationError("source manifest is missing outputs.image_embeddings")
+ image_embeddings_path = _resolve_output(
+ source_manifest_path, image_output, name="image_embeddings"
+ ).resolve()
+ if not image_embeddings_path.is_file():
+ raise AcceptanceValidationError(
+ f"source parity embeddings do not exist: {image_embeddings_path}"
+ )
+
+ groups = _source_groups(records, coordinates)
+ assignments = _assign_splits(groups, split_counts)
+ accepted = copy.deepcopy(records)
+ parity_locator = Path(os.path.relpath(image_embeddings_path, start=output_dir)).as_posix()
+ for group_id, members in groups:
+ split = assignments[group_id]
+ for index in members:
+ record = accepted[index]
+ record["acceptance_status"] = "accepted"
+ record["caption"]["status"] = "released_reused"
+ record["source_group_id"] = group_id
+ record["split"] = split
+ record["image_embedding"]["file"] = parity_locator
+
+ observed_counts = Counter(record["split"] for record in accepted)
+ if observed_counts != Counter(split_counts):
+ raise AssertionError(f"internal split-count mismatch: {observed_counts}")
+ for before, after in zip(records, accepted, strict=True):
+ if before["caption"]["original"] != after["caption"]["original"]:
+ raise AssertionError("caption.original changed during promotion")
+
+ output_dir.parent.mkdir(parents=True, exist_ok=True)
+ staging = Path(tempfile.mkdtemp(prefix=f".{output_dir.name}.tmp-", dir=output_dir.parent))
+ try:
+ candidates_path = staging / "candidates.jsonl"
+ _write_jsonl(candidates_path, accepted)
+ group_sizes = Counter(len(members) for _group_id, members in groups)
+ group_split_counts = Counter(assignments.values())
+ manifest = {
+ "schema_version": 1,
+ "artifact_type": "astroprism_stage1a_candidates",
+ "role": "accepted_training_candidates",
+ "created_at_utc": datetime.now(timezone.utc).isoformat(),
+ "acceptance": {
+ "basis": ACCEPTANCE_BASIS,
+ "accepted_count": len(accepted),
+ "acceptance_status": "accepted",
+ "caption_policy": "unchanged_original",
+ "caption_status": "released_reused",
+ },
+ "lineage": {
+ "source_manifest": {
+ "path": Path(
+ os.path.relpath(source_manifest_path, start=output_dir)
+ ).as_posix(),
+ "size_bytes": source_manifest_path.stat().st_size,
+ "sha256": sha256_file(source_manifest_path),
+ },
+ "source_candidates": {
+ "path": Path(
+ os.path.relpath(source_candidates_path, start=output_dir)
+ ).as_posix(),
+ "row_count": len(records),
+ "size_bytes": source_candidates_path.stat().st_size,
+ "sha256": source_candidates_sha256,
+ },
+ },
+ "pins": copy.deepcopy(source_manifest.get("pins", {})),
+ "roster": copy.deepcopy(source_manifest.get("roster", {})),
+ "sources": copy.deepcopy(source_manifest.get("sources", {})),
+ "parity_reference": {
+ "path": parity_locator,
+ "dtype": image_output.get("dtype"),
+ "shape": image_output.get("shape"),
+ "size_bytes": image_output.get("size_bytes"),
+ "sha256": image_output.get("sha256"),
+ },
+ "splits": {
+ "policy": SPLIT_POLICY,
+ "rank_seed": SPLIT_RANK_SEED,
+ "allocation_order": ["test", "validation", "train_remainder"],
+ "grouping_policy": GROUP_POLICY,
+ "grouping_coordinates": "coordinates_deg.paired_ra/paired_dec",
+ "overlap_radius_arcsec": OVERLAP_RADIUS_ARCSEC,
+ "row_counts": dict(sorted(observed_counts.items())),
+ "group_counts": dict(sorted(group_split_counts.items())),
+ "group_count": len(groups),
+ "group_size_histogram": {
+ str(size): count for size, count in sorted(group_sizes.items())
+ },
+ "mapping_sha256": _mapping_sha256(accepted),
+ },
+ "outputs": {
+ "candidates": {
+ "path": candidates_path.name,
+ "row_count": len(accepted),
+ "size_bytes": candidates_path.stat().st_size,
+ "sha256": sha256_file(candidates_path),
+ }
+ },
+ }
+ _write_json(staging / "manifest.json", manifest)
+ _commit_staging(staging, output_dir, overwrite=config.overwrite)
+ staging = None
+ finally:
+ if staging is not None and staging.exists():
+ shutil.rmtree(staging)
+
+ return {
+ "status": "accepted",
+ "output_dir": str(output_dir),
+ "candidate_count": len(accepted),
+ "group_count": len(groups),
+ "split_counts": dict(sorted(observed_counts.items())),
+ "mapping_sha256": _mapping_sha256(accepted),
+ "candidates_sha256": manifest["outputs"]["candidates"]["sha256"],
+ }
+
+
+def parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument(
+ "--source-manifest",
+ type=Path,
+ default=REPO_ROOT / "_local/astroprism/stage1a_candidates/manifest.json",
+ )
+ parser.add_argument(
+ "--output",
+ type=Path,
+ default=REPO_ROOT / "_local/astroprism/stage1a_accepted",
+ )
+ parser.add_argument(
+ "--overwrite",
+ action="store_true",
+ help="Atomically replace an existing accepted artifact.",
+ )
+ return parser.parse_args()
+
+
+def main() -> int:
+ args = parse_args()
+ summary = accept_stage1a(
+ AcceptanceConfig(
+ source_manifest_path=args.source_manifest,
+ output_dir=args.output,
+ overwrite=args.overwrite,
+ )
+ )
+ print(json.dumps(summary, indent=2, sort_keys=True))
+ return 0
+
+
+if __name__ == "__main__":
+ raise SystemExit(main())
diff --git a/tools/compose_astroprism_joint_adapters.py b/tools/compose_astroprism_joint_adapters.py
new file mode 100644
index 0000000..d38564e
--- /dev/null
+++ b/tools/compose_astroprism_joint_adapters.py
@@ -0,0 +1,174 @@
+#!/usr/bin/env python3
+"""Compose the trained image and spectrum adapters into one small checkpoint.
+
+The Stage 1A and Stage 1B checkpoints contain complete frozen Qwen/AION state,
+but a joint run only needs the independently trained projector tensors. This
+tool reads those tensors lazily and writes a compact checkpoint accepted by
+``training.resume_weights_only``.
+"""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+from collections.abc import Iterable
+from pathlib import Path
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+EXPECTED_SHAPES = {
+ "projectors.image.fc1.bias": (2560,),
+ "projectors.image.fc1.weight": (2560, 768),
+ "projectors.image.fc2.bias": (2560,),
+ "projectors.image.fc2.weight": (2560, 2560),
+ "projectors.image.modality_embedding": (1, 1, 2560),
+ "projectors.spectrum.fc1.bias": (2560,),
+ "projectors.spectrum.fc1.weight": (2560, 768),
+ "projectors.spectrum.fc2.bias": (2560,),
+ "projectors.spectrum.fc2.weight": (2560, 2560),
+ "projectors.spectrum.final_norm.bias": (2560,),
+ "projectors.spectrum.final_norm.weight": (2560,),
+ "projectors.spectrum.modality_embedding": (1, 1, 2560),
+}
+
+
+def _weights_path(checkpoint: Path) -> Path:
+ path = checkpoint.expanduser().resolve()
+ if path.is_dir():
+ path = path / "model.safetensors"
+ if not path.is_file():
+ raise FileNotFoundError(f"Checkpoint weights not found: {path}")
+ return path
+
+
+def _tensor_digest(items: Iterable[tuple[str, torch.Tensor]]) -> str:
+ digest = hashlib.sha256()
+ for name, tensor in sorted(items):
+ contiguous = tensor.detach().cpu().contiguous()
+ digest.update(name.encode("utf-8"))
+ digest.update(str(contiguous.dtype).encode("ascii"))
+ digest.update(json.dumps(list(contiguous.shape)).encode("ascii"))
+ digest.update(contiguous.view(torch.uint8).numpy().tobytes())
+ return digest.hexdigest()
+
+
+def _read_adapter(path: Path, prefix: str) -> dict[str, torch.Tensor]:
+ selected: dict[str, torch.Tensor] = {}
+ with safe_open(path, framework="pt", device="cpu") as checkpoint:
+ for key in checkpoint.keys():
+ clean_key = key[7:] if key.startswith("module.") else key
+ if clean_key.startswith(prefix):
+ if clean_key in selected:
+ raise RuntimeError(f"Duplicate adapter tensor after prefix cleanup: {clean_key}")
+ selected[clean_key] = checkpoint.get_tensor(key)
+ return selected
+
+
+def compose_joint_adapters(
+ image_checkpoint: str | Path,
+ spectrum_checkpoint: str | Path,
+ output_dir: str | Path,
+ *,
+ force: bool = False,
+) -> dict:
+ """Write and return provenance for a compact joint-adapter checkpoint."""
+
+ image_path = _weights_path(Path(image_checkpoint))
+ spectrum_path = _weights_path(Path(spectrum_checkpoint))
+ output = Path(output_dir).expanduser().resolve()
+ weights_out = output / "model.safetensors"
+ manifest_out = output / "manifest.json"
+ state_out = output / "training_state.json"
+
+ existing = [path for path in (weights_out, manifest_out, state_out) if path.exists()]
+ if existing and not force:
+ raise FileExistsError(
+ "Refusing to overwrite joint adapter artifact; pass --force for: "
+ + ", ".join(str(path) for path in existing)
+ )
+
+ image_tensors = _read_adapter(image_path, "projectors.image.")
+ spectrum_tensors = _read_adapter(spectrum_path, "projectors.spectrum.")
+ tensors = {**image_tensors, **spectrum_tensors}
+
+ expected_keys = set(EXPECTED_SHAPES)
+ actual_keys = set(tensors)
+ if actual_keys != expected_keys:
+ missing = sorted(expected_keys - actual_keys)
+ unexpected = sorted(actual_keys - expected_keys)
+ raise RuntimeError(
+ f"Adapter tensor contract mismatch; missing={missing}, unexpected={unexpected}"
+ )
+ for key, expected_shape in EXPECTED_SHAPES.items():
+ actual_shape = tuple(tensors[key].shape)
+ if actual_shape != expected_shape:
+ raise RuntimeError(
+ f"Adapter tensor {key} has shape {actual_shape}; expected {expected_shape}"
+ )
+
+ output.mkdir(parents=True, exist_ok=True)
+ image_digest = _tensor_digest(image_tensors.items())
+ spectrum_digest = _tensor_digest(spectrum_tensors.items())
+ save_file(
+ tensors,
+ weights_out,
+ metadata={
+ "format": "pt",
+ "artifact": "astroprism_joint_adapter_init",
+ "image_adapter_sha256": image_digest,
+ "spectrum_adapter_sha256": spectrum_digest,
+ },
+ )
+ merged_digest = _tensor_digest(tensors.items())
+ manifest = {
+ "schema_version": 1,
+ "artifact": "astroprism_joint_adapter_init",
+ "image_checkpoint": str(image_path),
+ "spectrum_checkpoint": str(spectrum_path),
+ "image_adapter_sha256": image_digest,
+ "spectrum_adapter_sha256": spectrum_digest,
+ "merged_adapter_sha256": merged_digest,
+ "tensor_count": len(tensors),
+ "tensor_keys": sorted(tensors),
+ }
+ manifest_out.write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8")
+ state_out.write_text(
+ json.dumps(
+ {
+ "step": 0,
+ "stage": "joint_adapter_initialization",
+ "provenance_manifest": str(manifest_out),
+ },
+ indent=2,
+ )
+ + "\n",
+ encoding="utf-8",
+ )
+ return manifest
+
+
+def _parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument("--image-checkpoint", required=True, type=Path)
+ parser.add_argument("--spectrum-checkpoint", required=True, type=Path)
+ parser.add_argument("--output-dir", required=True, type=Path)
+ parser.add_argument("--force", action="store_true")
+ return parser.parse_args()
+
+
+def main() -> None:
+ args = _parse_args()
+ manifest = compose_joint_adapters(
+ args.image_checkpoint,
+ args.spectrum_checkpoint,
+ args.output_dir,
+ force=args.force,
+ )
+ print(json.dumps(manifest, indent=2))
+
+
+if __name__ == "__main__":
+ main()
diff --git a/tools/extract_astroprism_raw_images.py b/tools/extract_astroprism_raw_images.py
new file mode 100644
index 0000000..0e6fd57
--- /dev/null
+++ b/tools/extract_astroprism_raw_images.py
@@ -0,0 +1,115 @@
+#!/usr/bin/env python3
+"""Select the 646 AstroPRISM raw Legacy Survey image rows."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+from urllib.parse import urlsplit
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(REPO_ROOT) not in sys.path:
+ sys.path.insert(0, str(REPO_ROOT))
+
+from src.data.astroprism_raw_images import ( # noqa: E402
+ RAW_SOURCE_ETAG,
+ RAW_SOURCE_LENGTH,
+ RAW_SOURCE_URL,
+ RawImageConfig,
+ extract_raw_images,
+)
+
+
+def _is_http(source: str) -> bool:
+ return urlsplit(source).scheme.lower() in {"http", "https"}
+
+
+def parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument(
+ "--roster",
+ type=Path,
+ default=REPO_ROOT / "docs/applications/astroprism_data/overlap_candidates_646.csv.gz",
+ )
+ parser.add_argument(
+ "--source",
+ default=RAW_SOURCE_URL,
+ help="Official HTTP HDF5 URL or a local HDF5 path.",
+ )
+ parser.add_argument(
+ "--output",
+ type=Path,
+ default=REPO_ROOT / "_local/astroprism/raw_images",
+ )
+ parser.add_argument(
+ "--expected-source-length",
+ type=int,
+ default=None,
+ help="Required for custom HTTP sources; the official URL uses its pinned length.",
+ )
+ parser.add_argument(
+ "--expected-source-etag",
+ default=None,
+ help="Required for custom HTTP sources; include quotes if present in the ETag.",
+ )
+ parser.add_argument("--http-block-size", type=int, default=256 * 1024)
+ parser.add_argument("--http-timeout-seconds", type=float, default=30.0)
+ parser.add_argument("--http-probe-attempts", type=int, default=3)
+ parser.add_argument("--http-probe-backoff-seconds", type=float, default=1.0)
+ parser.add_argument(
+ "--caption-legacy-coordinate-tolerance-arcsec",
+ type=float,
+ default=1e-4,
+ help="Maximum raw Legacy-to-caption roster coordinate separation.",
+ )
+ parser.add_argument(
+ "--overwrite",
+ action="store_true",
+ help="Atomically replace an existing output after extraction succeeds.",
+ )
+ return parser.parse_args()
+
+
+def main() -> int:
+ args = parse_args()
+ source_is_official = args.source == RAW_SOURCE_URL
+ if _is_http(args.source):
+ expected_length = (
+ RAW_SOURCE_LENGTH
+ if source_is_official and args.expected_source_length is None
+ else args.expected_source_length
+ )
+ expected_etag = (
+ RAW_SOURCE_ETAG
+ if source_is_official and args.expected_source_etag is None
+ else args.expected_source_etag
+ )
+ else:
+ expected_length = args.expected_source_length
+ expected_etag = args.expected_source_etag
+
+ summary = extract_raw_images(
+ RawImageConfig(
+ roster_path=args.roster,
+ output_dir=args.output,
+ source=args.source,
+ overwrite=args.overwrite,
+ expected_source_length=expected_length,
+ expected_source_etag=expected_etag,
+ http_block_size=args.http_block_size,
+ http_timeout_seconds=args.http_timeout_seconds,
+ http_probe_attempts=args.http_probe_attempts,
+ http_probe_backoff_seconds=args.http_probe_backoff_seconds,
+ caption_legacy_coordinate_tolerance_arcsec=(
+ args.caption_legacy_coordinate_tolerance_arcsec
+ ),
+ )
+ )
+ print(json.dumps(summary, indent=2, sort_keys=True))
+ return 0
+
+
+if __name__ == "__main__":
+ raise SystemExit(main())
diff --git a/tools/extract_astroprism_raw_spectra.py b/tools/extract_astroprism_raw_spectra.py
new file mode 100644
index 0000000..c6c3680
--- /dev/null
+++ b/tools/extract_astroprism_raw_spectra.py
@@ -0,0 +1,109 @@
+#!/usr/bin/env python3
+"""Select the 646 AstroPRISM raw observed-frame DESI spectra."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+from urllib.parse import urlsplit
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(REPO_ROOT) not in sys.path:
+ sys.path.insert(0, str(REPO_ROOT))
+
+from src.data.astroprism_raw_spectra import ( # noqa: E402
+ RAW_SOURCE_ETAG,
+ RAW_SOURCE_LENGTH,
+ RAW_SOURCE_URL,
+ RawSpectrumConfig,
+ extract_raw_spectra,
+)
+
+
+def _is_http(source: str) -> bool:
+ return urlsplit(source).scheme.lower() in {"http", "https"}
+
+
+def parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument(
+ "--roster",
+ type=Path,
+ default=REPO_ROOT / "docs/applications/astroprism_data/overlap_candidates_646.csv.gz",
+ )
+ parser.add_argument(
+ "--source",
+ default=RAW_SOURCE_URL,
+ help="Official HTTP HDF5 URL or a local HDF5 path.",
+ )
+ parser.add_argument(
+ "--output",
+ type=Path,
+ default=REPO_ROOT / "_local/astroprism/raw_spectra",
+ )
+ parser.add_argument("--expected-source-length", type=int, default=None)
+ parser.add_argument(
+ "--expected-source-etag",
+ default=None,
+ help="Required for custom HTTP sources; include quotes if present in the ETag.",
+ )
+ parser.add_argument("--http-block-size", type=int, default=256 * 1024)
+ parser.add_argument("--http-timeout-seconds", type=float, default=30.0)
+ parser.add_argument("--http-probe-attempts", type=int, default=3)
+ parser.add_argument("--http-probe-backoff-seconds", type=float, default=1.0)
+ parser.add_argument(
+ "--caption-legacy-coordinate-tolerance-arcsec",
+ type=float,
+ default=1e-4,
+ )
+ parser.add_argument(
+ "--overwrite",
+ action="store_true",
+ help="Atomically replace an existing output after extraction succeeds.",
+ )
+ return parser.parse_args()
+
+
+def main() -> int:
+ args = parse_args()
+ source_is_official = args.source == RAW_SOURCE_URL
+ if _is_http(args.source):
+ expected_length = (
+ RAW_SOURCE_LENGTH
+ if source_is_official and args.expected_source_length is None
+ else args.expected_source_length
+ )
+ expected_etag = (
+ RAW_SOURCE_ETAG
+ if source_is_official and args.expected_source_etag is None
+ else args.expected_source_etag
+ )
+ else:
+ expected_length = args.expected_source_length
+ expected_etag = args.expected_source_etag
+
+ summary = extract_raw_spectra(
+ RawSpectrumConfig(
+ roster_path=args.roster,
+ output_dir=args.output,
+ source=args.source,
+ overwrite=args.overwrite,
+ expected_source_length=expected_length,
+ expected_source_etag=expected_etag,
+ http_block_size=args.http_block_size,
+ http_timeout_seconds=args.http_timeout_seconds,
+ http_probe_attempts=args.http_probe_attempts,
+ http_probe_backoff_seconds=args.http_probe_backoff_seconds,
+ caption_legacy_coordinate_tolerance_arcsec=(
+ args.caption_legacy_coordinate_tolerance_arcsec
+ ),
+ )
+ )
+ print(json.dumps(summary, indent=2, sort_keys=True))
+ return 0
+
+
+if __name__ == "__main__":
+ raise SystemExit(main())
diff --git a/tools/generate_astroprism_joint_examples.py b/tools/generate_astroprism_joint_examples.py
new file mode 100644
index 0000000..073ec52
--- /dev/null
+++ b/tools/generate_astroprism_joint_examples.py
@@ -0,0 +1,328 @@
+#!/usr/bin/env python3
+"""Generate held-out image and spectrum captions from a joint AstroPRISM checkpoint."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import random
+import sys
+import time
+from dataclasses import fields
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from omegaconf import OmegaConf
+from safetensors import safe_open
+from transformers import AutoTokenizer
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(REPO_ROOT) not in sys.path:
+ sys.path.insert(0, str(REPO_ROOT))
+
+from src.config import ModelConfig
+from src.model import UnifiedTransformer
+from src.train import _build_astroprism_joint_data
+from src.training.distributed import load_model_weights_only
+
+
+def _parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser(
+ description=(
+ "Run greedy, prompt-only generation for one held-out image and one "
+ "held-out spectrum using a joint AstroPRISM checkpoint."
+ )
+ )
+ parser.add_argument("--config", type=Path, required=True, help="Saved Hydra config.yaml")
+ parser.add_argument(
+ "--checkpoint",
+ type=Path,
+ required=True,
+ help="Checkpoint directory containing model.safetensors",
+ )
+ parser.add_argument("--output", type=Path, required=True, help="Output JSON path")
+ parser.add_argument("--device", default="xpu:0")
+ parser.add_argument("--split", choices=("validation", "test"), default="test")
+ parser.add_argument("--image-index", type=int, default=43)
+ parser.add_argument("--spectrum-index", type=int, default=48)
+ parser.add_argument("--image-max-new-tokens", type=int, default=320)
+ parser.add_argument("--spectrum-max-new-tokens", type=int, default=128)
+ parser.add_argument("--seed", type=int, default=42)
+ parser.add_argument(
+ "--attn-implementation",
+ choices=("sdpa", "eager"),
+ default=None,
+ help="Optional inference-only override; defaults to the saved run setting.",
+ )
+ return parser.parse_args()
+
+
+def _model_config_from_saved_run(cfg: Any, *, attn_implementation: str | None) -> ModelConfig:
+ raw = OmegaConf.to_container(cfg.model, resolve=True)
+ if not isinstance(raw, dict):
+ raise TypeError("saved model configuration must resolve to a mapping")
+ raw = dict(raw)
+ raw["llm_backbone_id"] = raw.pop("backbone_id")
+ raw["llm_tokenizer_id"] = raw.pop("tokenizer_id")
+ raw["max_merged_seq_length"] = int(cfg.data.max_merged_seq_length)
+ raw["text_dropout"] = 0.0
+ if attn_implementation is not None:
+ raw["attn_implementation"] = attn_implementation
+ allowed = {field.name for field in fields(ModelConfig)}
+ return ModelConfig(**{key: value for key, value in raw.items() if key in allowed})
+
+
+def _prompt_only(sample: dict[str, Any]) -> tuple[torch.Tensor, str]:
+ text = sample.get("text")
+ labels = sample.get("labels")
+ if not isinstance(text, torch.Tensor) or not isinstance(labels, torch.Tensor):
+ raise TypeError("sample text and labels must be tensors")
+ if text.ndim != 1 or labels.shape != text.shape:
+ raise ValueError("sample text and labels must be aligned one-dimensional tensors")
+
+ supervised = torch.nonzero(labels.ne(-100), as_tuple=False).flatten()
+ if supervised.numel() == 0:
+ raise ValueError("sample has no supervised assistant tokens")
+ answer_start = int(supervised[0].item())
+ if answer_start <= 0:
+ raise ValueError("sample has no masked prompt prefix")
+ if not bool(labels[:answer_start].eq(-100).all()):
+ raise ValueError("prompt mask is not a contiguous leading prefix")
+ if not bool(labels[answer_start:].ne(-100).all()):
+ raise ValueError("assistant labels are not a contiguous suffix")
+ return text[:answer_start].clone(), "leading_ignore_index_prefix"
+
+
+def _checkpoint_contract(checkpoint: Path) -> dict[str, Any]:
+ weights = checkpoint / "model.safetensors"
+ if not weights.is_file():
+ raise FileNotFoundError(f"checkpoint weights do not exist: {weights}")
+ with safe_open(weights, framework="pt", device="cpu") as handle:
+ keys = tuple(handle.keys())
+ projector_keys = tuple(key for key in keys if key.startswith("projectors."))
+ aion_keys = tuple(key for key in keys if key.startswith("encoders.aion."))
+ if len(projector_keys) != 12:
+ raise RuntimeError(
+ f"joint checkpoint must contain 12 projector tensors, found {len(projector_keys)}"
+ )
+ if not aion_keys:
+ raise RuntimeError("joint checkpoint contains no shared encoders.aion tensors")
+ return {
+ "weights": str(weights.resolve()),
+ "tensor_count": len(keys),
+ "projector_tensor_count": len(projector_keys),
+ "shared_aion_tensor_count": len(aion_keys),
+ }
+
+
+def _move_inputs(value: Any, device: torch.device) -> Any:
+ if isinstance(value, torch.Tensor):
+ return value.to(device=device)
+ if isinstance(value, dict):
+ return {key: _move_inputs(item, device) for key, item in value.items()}
+ raise TypeError(f"unsupported inference input type: {type(value).__name__}")
+
+
+def _generation_inputs(
+ sample: dict[str, Any],
+ *,
+ modality: str,
+ prompt_ids: torch.Tensor,
+ device: torch.device,
+) -> dict[str, Any]:
+ inputs: dict[str, Any] = {"text": prompt_ids.unsqueeze(0)}
+ if modality == "image":
+ image = sample.get("image")
+ if not isinstance(image, torch.Tensor) or tuple(image.shape) != (4, 160, 160):
+ raise ValueError("image sample must contain float[4,160,160]")
+ if image.dtype != torch.float32:
+ raise TypeError(f"raw image must remain float32, got {image.dtype}")
+ inputs["image"] = image.unsqueeze(0)
+ elif modality == "spectrum":
+ spectrum = sample.get("spectrum")
+ if not isinstance(spectrum, dict):
+ raise TypeError("spectrum sample must contain a field mapping")
+ expected_dtypes = {
+ "flux": torch.float32,
+ "ivar": torch.float32,
+ "mask": torch.bool,
+ "wavelength": torch.float32,
+ }
+ if set(spectrum) != set(expected_dtypes):
+ raise ValueError(f"unexpected spectrum fields: {sorted(spectrum)}")
+ for field, dtype in expected_dtypes.items():
+ tensor = spectrum[field]
+ if not isinstance(tensor, torch.Tensor) or tuple(tensor.shape) != (7781,):
+ raise ValueError(f"spectrum.{field} must have shape [7781]")
+ if tensor.dtype != dtype:
+ raise TypeError(f"spectrum.{field} must remain {dtype}, got {tensor.dtype}")
+ inputs["spectrum"] = {
+ field: tensor.unsqueeze(0) for field, tensor in spectrum.items()
+ }
+ else:
+ raise ValueError(f"unsupported modality {modality!r}")
+ return _move_inputs(inputs, device)
+
+
+def _generate_one(
+ *,
+ model: UnifiedTransformer,
+ tokenizer: Any,
+ sample: dict[str, Any],
+ modality: str,
+ max_new_tokens: int,
+ device: torch.device,
+) -> dict[str, Any]:
+ if max_new_tokens <= 0:
+ raise ValueError("max_new_tokens must be positive")
+ prompt_ids, prompt_mask_contract = _prompt_only(sample)
+ reference_ids = sample["text"][prompt_ids.numel() :]
+ inputs = _generation_inputs(
+ sample,
+ modality=modality,
+ prompt_ids=prompt_ids,
+ device=device,
+ )
+
+ use_autocast = device.type in {"xpu", "cuda"}
+ started = time.perf_counter()
+ with torch.inference_mode(), torch.autocast(
+ device_type=device.type,
+ dtype=torch.bfloat16,
+ enabled=use_autocast,
+ ):
+ generated = model.generate(
+ inputs,
+ max_new_tokens=max_new_tokens,
+ do_sample=False,
+ pad_token_id=tokenizer.pad_token_id,
+ eos_token_id=tokenizer.eos_token_id,
+ use_cache=True,
+ )
+ elapsed = time.perf_counter() - started
+
+ if not isinstance(generated, torch.Tensor) or generated.ndim != 2:
+ raise RuntimeError(
+ f"generation must return a rank-2 token tensor, got {type(generated).__name__}"
+ )
+ generated_ids = generated[0].detach().cpu()
+ eos_id = int(tokenizer.eos_token_id)
+ hit_eos = bool(generated_ids.eq(eos_id).any())
+ generated_text = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
+ if not generated_text:
+ raise RuntimeError(f"{modality} generation decoded to empty text")
+
+ provenance = dict(sample.get("provenance") or {})
+ return {
+ "modality": modality,
+ "candidate_index": int(provenance["candidate_index"]),
+ "paired_targetid": str(provenance["paired_targetid"]),
+ "caption_object_id": str(provenance["caption_object_id"]),
+ "source_group_id": provenance.get("source_group_id"),
+ "split": provenance.get("split"),
+ "prompt_mask_contract": prompt_mask_contract,
+ "reference_supplied_to_model": False,
+ "prompt_token_count": int(prompt_ids.numel()),
+ "prompt_rendered": tokenizer.decode(prompt_ids, skip_special_tokens=False),
+ "reference_text": tokenizer.decode(reference_ids, skip_special_tokens=True).strip(),
+ "generated_text": generated_text,
+ "generated_token_count": int(generated_ids.numel()),
+ "generated_token_ids": generated_ids.tolist(),
+ "max_new_tokens": int(max_new_tokens),
+ "finish_reason": "eos" if hit_eos else "max_new_tokens",
+ "elapsed_seconds": elapsed,
+ }
+
+
+def main() -> None:
+ args = _parse_args()
+ if not args.config.is_file():
+ raise FileNotFoundError(f"saved Hydra config does not exist: {args.config}")
+ if args.image_index < 0 or args.spectrum_index < 0:
+ raise ValueError("held-out dataset indices must be non-negative")
+
+ os.environ.setdefault("HF_HUB_OFFLINE", "1")
+ os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
+ random.seed(args.seed)
+ np.random.seed(args.seed)
+ torch.manual_seed(args.seed)
+
+ cfg = OmegaConf.load(args.config)
+ model_config = _model_config_from_saved_run(
+ cfg,
+ attn_implementation=args.attn_implementation,
+ )
+ checkpoint_contract = _checkpoint_contract(args.checkpoint)
+ tokenizer = AutoTokenizer.from_pretrained(
+ model_config.llm_tokenizer_id,
+ trust_remote_code=True,
+ local_files_only=True,
+ )
+ if tokenizer.pad_token_id is None or tokenizer.eos_token_id is None:
+ raise RuntimeError("tokenizer must define pad_token_id and eos_token_id")
+
+ model = UnifiedTransformer(model_config)
+ checkpoint_step = load_model_weights_only(model, str(args.checkpoint), device="cpu")
+ model.tokenizer = tokenizer
+
+ device = torch.device(args.device)
+ model.to(device=device)
+ model.eval()
+ model.refresh_encoder_frozen_cache()
+ if set(model.encoders) != {"aion"}:
+ raise RuntimeError(f"expected one shared AION owner, found {list(model.encoders)}")
+ if model._encoder_key_by_modality != {"image": "aion", "spectrum": "aion"}:
+ raise RuntimeError(f"unexpected joint route map: {model._encoder_key_by_modality}")
+
+ dataset, _ = _build_astroprism_joint_data(cfg.data, tokenizer, split=args.split)
+ if args.image_index >= len(dataset.image_dataset):
+ raise IndexError("image index is outside the held-out dataset")
+ if args.spectrum_index >= len(dataset.spectrum_dataset):
+ raise IndexError("spectrum index is outside the held-out dataset")
+
+ image_result = _generate_one(
+ model=model,
+ tokenizer=tokenizer,
+ sample=dataset.image_dataset[args.image_index],
+ modality="image",
+ max_new_tokens=args.image_max_new_tokens,
+ device=device,
+ )
+ spectrum_result = _generate_one(
+ model=model,
+ tokenizer=tokenizer,
+ sample=dataset.spectrum_dataset[args.spectrum_index],
+ modality="spectrum",
+ max_new_tokens=args.spectrum_max_new_tokens,
+ device=device,
+ )
+
+ result = {
+ "schema_version": 1,
+ "checkpoint": str(args.checkpoint.resolve()),
+ "checkpoint_step": checkpoint_step,
+ "checkpoint_contract": checkpoint_contract,
+ "saved_config": str(args.config.resolve()),
+ "device": str(device),
+ "dtype": str(next(model.backbone.parameters()).dtype),
+ "attn_implementation": model.backbone.config._attn_implementation,
+ "decoding": {"strategy": "greedy", "do_sample": False, "seed": args.seed},
+ "dataset": {
+ "split": args.split,
+ "image_count": len(dataset.image_dataset),
+ "spectrum_count": len(dataset.spectrum_dataset),
+ },
+ "examples": [image_result, spectrum_result],
+ }
+ args.output.parent.mkdir(parents=True, exist_ok=True)
+ temporary = args.output.with_suffix(args.output.suffix + ".tmp")
+ temporary.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8")
+ temporary.replace(args.output)
+ print(json.dumps(result, indent=2, sort_keys=True), flush=True)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/tools/launch_aurora.py b/tools/launch_aurora.py
index d59c87b..05b72ba 100644
--- a/tools/launch_aurora.py
+++ b/tools/launch_aurora.py
@@ -225,10 +225,11 @@ def main():
"\n"
"# Rank 0 on each node takes care of unpacking\n"
'if [ "$LOCAL_RANK" == "0" ]; then\n'
- ' if [ ! -f "$MARKER_FILE" ]; then\n'
+ ' if [ ! -f "$MARKER_FILE" ] || [ "$ENV_TARBALL" -nt "$MARKER_FILE" ]; then\n'
+ ' rm -rf "$LOCAL_ENV"\n'
' echo "Rank $RANK (Local 0): Unpacking $ENV_TARBALL to $LOCAL_ENV..."\n'
- " mkdir -p $LOCAL_ENV\n"
- " tar -xzf $ENV_TARBALL -C $LOCAL_ENV\n"
+ ' mkdir -p "$LOCAL_ENV"\n'
+ ' tar -xzf "$ENV_TARBALL" -C "$LOCAL_ENV"\n'
" if [ $? -ne 0 ]; then\n"
' echo "ERROR: Failed to unpack $ENV_TARBALL"\n'
" exit 1\n"
@@ -240,12 +241,22 @@ def main():
" fi\n"
"fi\n"
"\n"
- "# Wait for unpack completion\n"
- 'while [ ! -f "$MARKER_FILE" ]; do\n'
+ "# Wait for unpack completion. The marker must also be newer than the\n"
+ "# tarball: a stale marker from a previous run satisfies an existence-only\n"
+ "# test immediately, letting other ranks activate the environment while\n"
+ "# rank 0 is still deleting and re-extracting it.\n"
+ 'while [ ! -f "$MARKER_FILE" ] || [ "$ENV_TARBALL" -nt "$MARKER_FILE" ]; do\n'
" sleep 1\n"
"done\n"
)
- activate_line = "source $LOCAL_ENV/bin/activate"
+ # setup_deepspeed_env.sh builds .venv-deepspeed/ and the documented
+ # tar command preserves that top-level directory. Relocate VIRTUAL_ENV
+ # before sourcing so console scripts and diagnostics refer to /tmp.
+ activate_line = (
+ 'sed -i "s|VIRTUAL_ENV=.*|VIRTUAL_ENV=\\"$LOCAL_ENV/.venv-deepspeed\\"|" '
+ "$LOCAL_ENV/.venv-deepspeed/bin/activate\n"
+ "source $LOCAL_ENV/.venv-deepspeed/bin/activate"
+ )
model_section = (
"# --- Model Copy (Local /tmp) ---\n"
'export SHARED_HF_HOME="' + default_shared_hf_home + '"\n'
@@ -436,6 +447,21 @@ def main():
# print("\n--- SCRIPT CONTENT ---")
# print(run_script_content)
pass
+ elif args.batch:
+ # A PBS script has to be submitted, not executed: running it here would
+ # run the job on the login node and ignore every #PBS directive above.
+ print("Submitting batch job via qsub...")
+ try:
+ result = subprocess.run(
+ ["qsub", script_path],
+ check=True,
+ capture_output=True,
+ text=True,
+ )
+ print(result.stdout.strip())
+ except subprocess.CalledProcessError as e:
+ print(f"qsub failed: {e.stderr}")
+ raise SystemExit(e.returncode) from e
else:
print("Executing...")
try:
diff --git a/tools/launch_aurora_web.py b/tools/launch_aurora_web.py
index a31bf58..a76404e 100644
--- a/tools/launch_aurora_web.py
+++ b/tools/launch_aurora_web.py
@@ -19,6 +19,7 @@
import datetime
import json
import os
+import re
import subprocess
import sys
import textwrap
@@ -61,6 +62,13 @@ def main():
default="deepspeed_env.tar.gz",
help="Path to packed env tarball",
)
+ parser.add_argument(
+ "--skip-model-staging",
+ action="store_true",
+ help="Skip copying shared Hugging Face caches to node-local /tmp. "
+ "Use only when every required model/tokenizer/cache path is an "
+ "absolute, existing shared-filesystem path.",
+ )
parser.add_argument(
"--prism-dir",
default=None,
@@ -685,6 +693,50 @@ def main():
print(f"Model staging: {backbone_id} (local path, no staging needed)")
else:
print(f"Model staging: {backbone_id} -> {hf_model_dir}")
+ if args.skip_model_staging:
+ tokenizer_id = resolve_hydra_override(
+ final_overrides,
+ unknown_args,
+ "model.tokenizer_id",
+ model_group_config.get("tokenizer_id", backbone_id),
+ )
+ required_shared_paths = {
+ "model.backbone_id": backbone_id,
+ "model.tokenizer_id": tokenizer_id,
+ }
+ if image_encoder_id is not None:
+ required_shared_paths["model.image_encoder_id"] = image_encoder_id
+ if image_processor_id is not None:
+ required_shared_paths["model.image_processor_id"] = image_processor_id
+
+ # load_model_group_config uses yaml.safe_load and resolve_hydra_override
+ # only selects values, so Hydra interpolations arrive here as literal
+ # "${oc.env:VAR,default}/..." strings. Expand the oc.env form before
+ # testing the filesystem, or a config that resolves to real staged
+ # assets is rejected for looking relative.
+ def _expand(value: str | None) -> str | None:
+ if not isinstance(value, str):
+ return value
+ def _sub(match: "re.Match[str]") -> str:
+ body = match.group(1)
+ name, _, fallback = body.partition(",")
+ return os.environ.get(name.strip(), fallback.strip())
+ return re.sub(r"\$\{oc\.env:([^}]*)\}", _sub, value)
+
+ invalid_paths = []
+ for key, value in required_shared_paths.items():
+ resolved = _expand(value)
+ if resolved is None or not os.path.isabs(resolved) or not os.path.exists(resolved):
+ invalid_paths.append(f"{key}={value!r}")
+ if invalid_paths:
+ parser.error(
+ "--skip-model-staging requires absolute, existing shared paths; "
+ "invalid: " + ", ".join(invalid_paths)
+ )
+ print(
+ "Model staging disabled: all required model assets resolve to "
+ "absolute shared-filesystem paths."
+ )
if image_encoder_id:
print(f"Image encoder staging: {image_encoder_id} -> {hf_cache_dir(image_encoder_id)}")
if image_processor_id and image_processor_id != image_encoder_id:
@@ -1039,7 +1091,7 @@ def _override_key(override: str) -> str:
fi
while [ ! -f "$MARKER_FILE" ]; do sleep 1; done'''}
-{"# --- Model Cache (shared filesystem) ---" + chr(10) + 'export HF_HOME="' + default_hf_home + '"' + chr(10) + 'echo " HF_HOME set to: $HF_HOME (shared filesystem)"' if args.use_shared_venv else '''# --- Model Copy ---
+{"# --- Model Cache (shared filesystem) ---" + chr(10) + 'export HF_HOME="' + default_hf_home + '"' + chr(10) + 'echo " HF_HOME set to: $HF_HOME (shared filesystem)"' if args.use_shared_venv or args.skip_model_staging else '''# --- Model Copy ---
export SHARED_HF_ROOT="''' + default_hf_home + '''"
export SHARED_HF_HOME="''' + default_shared_hf_home + '''"
export LOCAL_HF_ROOT="/tmp/huggingface"
@@ -1073,7 +1125,12 @@ def _override_key(override: str) -> str:
echo " Linked standalone artifact $ARTIFACT"
fi
done
- if [ ! -f "$MODEL_MARKER" ]; then
+ if [ "''' + ("1" if args.skip_model_staging else "0") + '''" = "1" ]; then
+ # No marker here. It means the node cache is staged, and nothing was
+ # copied; a later staging-enabled launch on the same node would see
+ # it, skip its own copy loop, then fail offline on an empty cache.
+ echo "Rank $RANK (Local 0): Skipping model staging; using absolute shared paths."
+ elif [ ! -f "$MODEL_MARKER" ]; then
echo "Rank $RANK (Local 0): Staging models to /tmp..."
for MODEL in "${MODELS_TO_STAGE[@]}"; do
if [ -d "$SHARED_HF_HOME/$MODEL" ] && [ ! -d "$LOCAL_HF_HOME/$MODEL" ]; then
@@ -1087,7 +1144,9 @@ def _override_key(override: str) -> str:
echo "Rank $RANK (Local 0): Models already staged."
fi
fi
-while [ ! -f "$MODEL_MARKER" ]; do sleep 2; done'''}
+if [ "''' + ("0" if args.skip_model_staging else "1") + '''" = "1" ]; then
+ while [ ! -f "$MODEL_MARKER" ]; do sleep 2; done
+fi'''}
# --- WebDataset Shard Staging (Node-Local) ---
export SHARD_MARKER="$LOCAL_SHARDS_DIR/shards_ready"
@@ -1161,7 +1220,7 @@ def _override_key(override: str) -> str:
export MASTER_ADDR="$MASTER_ADDR"
export MASTER_PORT="$MASTER_PORT"
-{'export HF_HOME="' + default_hf_home + '"' + chr(10) + 'export TRANSFORMERS_CACHE="' + default_shared_hf_home + '"' + chr(10) + 'export HF_HUB_CACHE="' + default_shared_hf_home + '"' + chr(10) + 'export HF_DATASETS_CACHE="' + default_hf_home + '/datasets"' if args.use_shared_venv else '''export HF_HOME="/tmp/huggingface"
+{'export HF_HOME="' + default_hf_home + '"' + chr(10) + 'export TRANSFORMERS_CACHE="' + default_shared_hf_home + '"' + chr(10) + 'export HF_HUB_CACHE="' + default_shared_hf_home + '"' + chr(10) + 'export HF_DATASETS_CACHE="' + default_hf_home + '/datasets"' if args.use_shared_venv or args.skip_model_staging else '''export HF_HOME="/tmp/huggingface"
export TRANSFORMERS_CACHE="/tmp/huggingface/hub"
export HF_HUB_CACHE="/tmp/huggingface/hub"
export HF_DATASETS_CACHE="/tmp/huggingface/datasets"'''}
diff --git a/tools/materialize_astroprism_hf_dataset.py b/tools/materialize_astroprism_hf_dataset.py
new file mode 100644
index 0000000..6dbc85e
--- /dev/null
+++ b/tools/materialize_astroprism_hf_dataset.py
@@ -0,0 +1,405 @@
+#!/usr/bin/env python3
+"""Materialize the published AstroPRISM Hugging Face dataset for training.
+
+The training configs under ``src/conf/data/astroprism_*.yaml`` read a local
+tree of manifests (``artifact_type: astroprism_raw_images`` and friends) whose
+``outputs`` entries point at the ``.npy`` arrays and candidate JSONL beside
+them. The Hugging Face package ships the same payloads under a different
+envelope: its root ``MANIFEST.json`` is an upload index
+(``artifact_type: astroprism_huggingface_dataset``) and ``metadata/
+candidates.jsonl`` is the candidate file itself rather than a manifest naming
+one. Nothing in PRISM resolves the ``hf://`` paths the original run used.
+
+This script bridges the two. It downloads (or reuses) the dataset snapshot and
+writes the manifest envelope the loaders expect, leaving the large arrays in
+the Hugging Face cache and referencing them by absolute path, so nothing is
+copied twice.
+
+ export HF_TOKEN=... # the dataset is gated
+ python tools/materialize_astroprism_hf_dataset.py --output-dir _local/astroprism
+
+Then train with the shipped configs, which read that tree::
+
+ python tools/launch_aurora.py --file experiments/astroprism_stage1a.yaml ...
+"""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import os
+import sys
+from pathlib import Path
+from typing import Any
+
+SNAPSHOT_FILES = [
+ "MANIFEST.json",
+ "metadata/candidates.jsonl",
+ "assets/raw/images/MANIFEST.json",
+ "assets/raw/spectra/MANIFEST.json",
+]
+ARRAY_PATTERNS = ["assets/raw/images/*.npy", "assets/raw/spectra/*.npy"]
+
+REPO_ID = "Sand33p/AstroPRISM"
+
+# src/data/astroprism_stage1b_spectrum_dataset.py::_ROLE_BY_SPLIT
+_ROLE_BY_SPLIT = {
+ "train": "alignment_training",
+ "validation": "alignment_validation",
+ "test": "alignment_test",
+}
+
+
+class MaterializeError(RuntimeError):
+ """Raised when the snapshot does not match what the loaders require."""
+
+
+def _sha256(path: Path) -> str:
+ h = hashlib.sha256()
+ with open(path, "rb") as fh:
+ for chunk in iter(lambda: fh.read(8 * 1024 * 1024), b""):
+ h.update(chunk)
+ return h.hexdigest()
+
+
+def _read_json(path: Path) -> dict[str, Any]:
+ try:
+ return json.loads(path.read_text(encoding="utf-8"))
+ except (OSError, json.JSONDecodeError) as exc:
+ raise MaterializeError(f"cannot read {path}: {exc}") from exc
+
+
+def _absolutize(manifest: dict[str, Any], manifest_path: Path) -> dict[str, Any]:
+ """Rewrite every ``outputs[*].path`` to an absolute path in the snapshot.
+
+ The loaders resolve a relative output path against the manifest's own
+ directory. Writing the manifest elsewhere would break that, so the paths
+ are pinned to where the arrays actually live.
+ """
+ out = json.loads(json.dumps(manifest)) # deep copy
+ outputs = out.get("outputs")
+ if not isinstance(outputs, dict):
+ raise MaterializeError(f"{manifest_path} has no outputs mapping")
+ for name, entry in outputs.items():
+ if not isinstance(entry, dict) or "path" not in entry:
+ continue
+ resolved = (manifest_path.parent / entry["path"]).resolve()
+ if not resolved.is_file():
+ raise MaterializeError(
+ f"{manifest_path.name} names {name} at {resolved}, which does not exist. "
+ "Re-run without --no-download, or pass --snapshot-dir for a complete copy."
+ )
+ entry["path"] = str(resolved)
+ return out
+
+
+def _raw_candidate_rows(snapshot: Path, records: list[dict[str, Any]]) -> list[dict[str, Any]]:
+ """Rebuild the per-candidate list the raw-image loader cross-checks.
+
+ The HF manifest records band order once, as a mapping, plus two parallel
+ ``.npy`` arrays holding the per-row source and AION band labels. The loader
+ wants one object per candidate carrying both, so read the arrays and pair
+ them with the roster rows, preserving index order.
+ """
+ import numpy as np
+
+ images = snapshot / "assets/raw/images"
+ try:
+ source_bands = np.load(images / "image_band.npy", allow_pickle=False)
+ aion_bands = np.load(images / "aion_image_band.npy", allow_pickle=False)
+ except OSError as exc:
+ raise MaterializeError(f"cannot read per-candidate band arrays: {exc}") from exc
+ if len(source_bands) != len(records) or len(aion_bands) != len(records):
+ raise MaterializeError(
+ f"band arrays hold {len(source_bands)}/{len(aion_bands)} rows but the roster "
+ f"has {len(records)}"
+ )
+
+ def _decode(row: Any) -> list[str]:
+ return [b.decode() if isinstance(b, bytes) else str(b) for b in row]
+
+ rows: list[dict[str, Any]] = []
+ for index, record in enumerate(records):
+ rows.append(
+ {
+ "candidate_index": index,
+ "paired_targetid": record.get("paired_targetid"),
+ "caption_object_id": record.get("caption_object_id"),
+ "bands": {
+ "source": _decode(source_bands[index]),
+ "aion": _decode(aion_bands[index]),
+ },
+ }
+ )
+ return rows
+
+
+def _candidate_manifest(candidates: Path, *, verify: bool) -> dict[str, Any]:
+ """Build the manifest wrapper the Stage 1A/1B loaders expect."""
+ rows = sum(1 for line in candidates.open(encoding="utf-8") if line.strip())
+ if rows == 0:
+ raise MaterializeError(f"{candidates} is empty")
+ entry: dict[str, Any] = {"path": str(candidates.resolve()), "row_count": rows}
+ if verify:
+ entry["sha256"] = _sha256(candidates)
+ return {
+ # One of the two artifact types the loaders accept; see
+ # src/data/astroprism_stage1a_dataset.py::_validate_manifest_types.
+ "artifact_type": "astroprism_aion_feature_parity_reference",
+ "schema_version": 1,
+ "source": {"repo_id": REPO_ID, "repo_type": "dataset"},
+ "outputs": {"candidates": entry, "captions.jsonl": entry},
+ }
+
+
+def _spectrum_caption_records(
+ snapshot: Path, image_records: list[dict[str, Any]]
+) -> list[dict[str, Any]]:
+ """Build the Stage 1B roster from the dataset's spectrum_caption Parquet.
+
+ metadata/candidates.jsonl carries only the image captions, so Stage 1B has
+ no roster in the published tree. The spectrum captions live in the
+ ``spectrum_caption`` config instead, one Parquet per split. Join them back
+ onto the image roster by ``candidate_index`` -- which is also the raw-array
+ index -- and restate the per-record review and policy facts the loader
+ checks. Those are not invented: ``caption_status`` and ``training_allowed``
+ come from the Parquet, and the split roles follow the dataset card, which
+ permits gradients on train only and marks validation and test
+ inspection-only.
+ """
+ import pyarrow.parquet as pq
+
+ by_index = {r.get("candidate_index"): r for r in image_records}
+ rows: dict[int, dict[str, Any]] = {}
+ for split in ("train", "validation", "test"):
+ shard = snapshot / "data" / "spectrum_caption" / f"{split}.parquet"
+ if not shard.is_file():
+ raise MaterializeError(
+ f"spectrum captions for split {split!r} are missing at {shard}. "
+ "Re-run without --no-download, or fetch data/spectrum_caption/*.parquet."
+ )
+ table = pq.read_table(
+ shard,
+ columns=["candidate_index", "split", "caption", "caption_status",
+ "review_status", "training_allowed", "dataset_ready"],
+ )
+ for row in table.to_pylist():
+ index = row["candidate_index"]
+ base = by_index.get(index)
+ if base is None:
+ raise MaterializeError(
+ f"spectrum caption {index} has no matching image-roster row"
+ )
+ is_train = row["split"] == "train"
+ if bool(row["training_allowed"]) is not is_train:
+ raise MaterializeError(
+ f"spectrum caption {index} in split {row['split']!r} has "
+ f"training_allowed={row['training_allowed']!r}"
+ )
+ rows[index] = {
+ "candidate_index": index,
+ "paired_targetid": base.get("paired_targetid"),
+ "caption_object_id": base.get("caption_object_id"),
+ "source_group_id": base.get("source_group_id"),
+ "split": row["split"],
+ "coordinates_deg": base.get("coordinates_deg"),
+ "spectrum_caption": {
+ "original": row["caption"],
+ "status": row["caption_status"],
+ },
+ "review": {
+ "status": row["review_status"],
+ "decision_source": "user_instruction",
+ },
+ "training_policy": {
+ "dataset_ready": bool(row["dataset_ready"]),
+ "gradient_training_allowed": is_train,
+ "eligible_for_training": is_train,
+ "eligible_for_evaluation": not is_train,
+ "role": _ROLE_BY_SPLIT[row["split"]],
+ },
+ }
+ # candidate_index is the raw-array index, so roster order must be preserved.
+ return [rows[i] for i in sorted(rows)]
+
+
+def _spectrum_caption_manifest(
+ candidates: Path, records: list[dict[str, Any]], *, verify: bool
+) -> dict[str, Any]:
+ """Stage 1B checks a stricter envelope than Stage 1A around the same rows.
+
+ It requires its own artifact_type plus the review and training-policy
+ blocks that record the roster was accepted and which splits may take
+ gradients. Those facts come from the dataset card (train-only gradient
+ training; validation and test inspection-only), restated here in the shape
+ src/data/astroprism_stage1b_spectrum_dataset.py validates.
+ """
+ from collections import Counter
+
+ rows = len(records)
+ entry: dict[str, Any] = {"path": str(candidates.resolve()), "row_count": rows}
+ if verify:
+ entry["sha256"] = _sha256(candidates)
+ # The loader cross-checks these against the JSONL it reads, so derive them
+ # from the roster rather than restating them.
+ indices = [record.get("candidate_index") for record in records]
+ split_counts = dict(Counter(record.get("split") for record in records))
+ return {
+ "artifact_type": "astroprism_spectrum_captions_accepted",
+ "candidate_indices": indices,
+ "split_counts": split_counts,
+ "splits": {"row_counts": split_counts},
+ "schema_version": 1,
+ "role": "reviewed_alignment_dataset",
+ "row_count": rows,
+ "source": {"repo_id": REPO_ID, "repo_type": "dataset"},
+ "review": {"status": "accepted", "decision_source": "user_instruction"},
+ "training_policy": {
+ "dataset_ready": True,
+ "gradient_training_splits": ["train"],
+ "evaluation_splits": ["validation", "test"],
+ "roles": {
+ "train": "alignment_training",
+ "validation": "alignment_validation",
+ "test": "alignment_test",
+ },
+ },
+ "outputs": {"captions.jsonl": entry, "candidates": entry},
+ }
+
+
+def _fetch_snapshot(revision: str | None, token: str | None) -> Path:
+ from huggingface_hub import snapshot_download
+
+ return Path(
+ snapshot_download(
+ REPO_ID,
+ repo_type="dataset",
+ revision=revision,
+ token=token,
+ allow_patterns=SNAPSHOT_FILES + ARRAY_PATTERNS,
+ )
+ )
+
+
+def main(argv: list[str] | None = None) -> int:
+ ap = argparse.ArgumentParser(description=__doc__.split("\n", 1)[0])
+ ap.add_argument(
+ "--output-dir",
+ default="_local/astroprism",
+ help="where to write the manifest tree (default: %(default)s, what the configs read)",
+ )
+ ap.add_argument("--snapshot-dir", help="use an existing snapshot instead of downloading")
+ ap.add_argument("--revision", help="dataset revision to pin")
+ ap.add_argument("--no-download", action="store_true", help="fail rather than fetch")
+ ap.add_argument(
+ "--no-verify-hashes",
+ action="store_true",
+ help="skip the candidate SHA-256, which the loaders check when present",
+ )
+ args = ap.parse_args(argv)
+
+ if args.snapshot_dir:
+ snap = Path(args.snapshot_dir).expanduser().resolve()
+ if not snap.is_dir():
+ print(f"error: --snapshot-dir {snap} is not a directory", file=sys.stderr)
+ return 1
+ elif args.no_download:
+ print("error: --no-download requires --snapshot-dir", file=sys.stderr)
+ return 1
+ else:
+ token = os.environ.get("HF_TOKEN")
+ if not token:
+ print(
+ "error: the dataset is gated; set HF_TOKEN to a token whose account "
+ "has been granted access at https://huggingface.co/datasets/" + REPO_ID,
+ file=sys.stderr,
+ )
+ return 1
+ snap = _fetch_snapshot(args.revision, token)
+
+ out = Path(args.output_dir).expanduser().resolve()
+ verify = not args.no_verify_hashes
+
+ try:
+ images = _absolutize(_read_json(snap / "assets/raw/images/MANIFEST.json"),
+ snap / "assets/raw/images/MANIFEST.json")
+ spectra = _absolutize(_read_json(snap / "assets/raw/spectra/MANIFEST.json"),
+ snap / "assets/raw/spectra/MANIFEST.json")
+ candidates_file = snap / "metadata/candidates.jsonl"
+ candidates = _candidate_manifest(candidates_file, verify=verify)
+ records = [
+ json.loads(line)
+ for line in candidates_file.read_text(encoding="utf-8").splitlines()
+ if line.strip()
+ ]
+ # The loader cross-checks the raw manifest's candidate list against the
+ # roster row for row, so both must be present and in the same order.
+ images["candidates"] = _raw_candidate_rows(snap, records)
+ spectra["candidates"] = images["candidates"]
+ except MaterializeError as exc:
+ print(f"error: {exc}", file=sys.stderr)
+ return 1
+
+ # The loaders key on these exact artifact_type strings.
+ images["artifact_type"] = "astroprism_raw_images"
+ spectra["artifact_type"] = "astroprism_raw_spectra"
+
+ # Stage 1B needs its own roster: the published candidates.jsonl has image
+ # captions only. Build it from the spectrum_caption Parquet shards.
+ spectrum_dir = out / "stage1b_spectrum_caption_accepted"
+ spectrum_dir.mkdir(parents=True, exist_ok=True)
+ try:
+ spectrum_records = _spectrum_caption_records(snap, records)
+ except MaterializeError as exc:
+ print(f"error: {exc}", file=sys.stderr)
+ return 1
+ spectrum_file = spectrum_dir / "captions.jsonl"
+ spectrum_file.write_text(
+ "".join(json.dumps(r) + "\n" for r in spectrum_records), encoding="utf-8"
+ )
+ spectrum_captions = _spectrum_caption_manifest(
+ spectrum_file, spectrum_records, verify=verify
+ )
+
+ # Both manifests must declare the same roster provenance; the loader
+ # compares their SHA-256 and row count to prove the arrays and the
+ # captions describe one roster. The published package ships the roster
+ # under provenance/, so use it when present and fall back to the
+ # candidate file, which is the roster the arrays were built from.
+ roster_file = snap / "provenance/overlap_candidates_646.csv.gz"
+ if not roster_file.is_file():
+ roster_file = candidates_file
+ roster = {
+ "path": str(roster_file.resolve()),
+ "sha256": _sha256(roster_file),
+ "row_count": len(records),
+ }
+ for payload in (images, spectra, candidates, spectrum_captions):
+ payload["roster"] = roster
+ written = {
+ "raw_images/manifest.json": images,
+ "raw_spectra/manifest.json": spectra,
+ "stage1a_accepted/manifest.json": candidates,
+ "stage1b_spectrum_caption_accepted/manifest.json": spectrum_captions,
+ }
+ for rel, payload in written.items():
+ dest = out / rel
+ dest.parent.mkdir(parents=True, exist_ok=True)
+ dest.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")
+ print(f"wrote {dest}")
+
+ print(f"wrote {spectrum_file}")
+
+ print(f"\nsnapshot: {snap}")
+ print(f"candidates: {candidates['outputs']['candidates']['row_count']} rows")
+ print(
+ "\nThe arrays stay in the Hugging Face cache; these manifests point at them by\n"
+ "absolute path. Move or clear that cache and regenerate this tree."
+ )
+ return 0
+
+
+if __name__ == "__main__":
+ raise SystemExit(main())
diff --git a/tools/prepare_astroprism_hf_dataset.py b/tools/prepare_astroprism_hf_dataset.py
new file mode 100644
index 0000000..f56a6f3
--- /dev/null
+++ b/tools/prepare_astroprism_hf_dataset.py
@@ -0,0 +1,1495 @@
+#!/usr/bin/env python3
+"""Build a portable Hugging Face dataset upload tree for AstroPRISM.
+
+The package deliberately keeps the calibrated NumPy arrays monolithic. Every
+Parquet row retains ``candidate_index``, which is also the first-axis index into
+each raw array. RGB images and spectrum plots embedded in Parquet are review
+previews only; they are not substitutes for the calibrated model inputs.
+
+No source cache, released AION embedding, or local filesystem path is copied.
+"""
+
+from __future__ import annotations
+
+import argparse
+import gzip
+import hashlib
+import io
+import json
+import math
+import os
+import shutil
+import tempfile
+import uuid
+from collections import Counter, defaultdict
+from collections.abc import Iterable
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path, PurePosixPath
+from typing import Any
+
+import numpy as np
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+
+EXPECTED_CANDIDATES = 646
+EXPECTED_SPLIT_COUNTS = {"train": 516, "validation": 65, "test": 65}
+ROSTER_SHA256 = "650935faa58953c9929e7379429bfb332a7486dab54544fedef64ca385abb3bc"
+CAPTION_REPO_ID = "astronolan/galaxy-descriptions"
+CAPTION_REVISION = "6890dc0c8fc793867c07a6ce4fce11e51c167d6e"
+PAIRED_REPO_ID = "EiffL/desi_x_ls_aion"
+PAIRED_REVISION = "57f7c7575260811ea076163bc743787d75b9aa49"
+RAW_SOURCE_URL = "https://users.flatironinstitute.org/~polymathic/data/provabgs_desi_ls.hdf5"
+ASTROSPECLM_PAPER = "https://arxiv.org/abs/2609.07102"
+
+FORBIDDEN_PORTABLE_TEXT = ("/lus/", "_local/", "BaseMM_PRISM")
+TEXT_SUFFIXES = {".json", ".jsonl", ".md", ".txt", ".yaml", ".yml"}
+
+IMAGE_ARRAY_CONTRACT = {
+ "image_array": (np.dtype("float32"), (4, 160, 160)),
+ "image_ivar": (np.dtype("float32"), (4, 160, 160)),
+ "image_mask": (np.dtype("bool"), (160, 160)),
+ "object_mask": (np.dtype("uint8"), (160, 160)),
+ "image_scale": (np.dtype("float32"), (4,)),
+ "image_band": (np.dtype("S5"), (4,)),
+ "aion_image_band": (np.dtype("S5"), (4,)),
+}
+SPECTRUM_ARRAY_CONTRACT = {
+ "spectrum_flux": (np.dtype("float32"), (7781,)),
+ "spectrum_ivar": (np.dtype("float32"), (7781,)),
+ "spectrum_mask": (np.dtype("bool"), (7781,)),
+ "spectrum_wavelength": (np.dtype("float32"), (7781,)),
+}
+
+
+class PackagingValidationError(ValueError):
+ """Raised when an input violates the portable dataset contract."""
+
+
+@dataclass(frozen=True)
+class HFDatasetConfig:
+ """Inputs for one atomic dataset staging run."""
+
+ output_dir: Path
+ accepted_candidates_path: Path = (
+ REPO_ROOT / "_local/astroprism/stage1a_accepted/candidates.jsonl"
+ )
+ accepted_manifest_path: Path = REPO_ROOT / "_local/astroprism/stage1a_accepted/manifest.json"
+ raw_images_manifest_path: Path = REPO_ROOT / "_local/astroprism/raw_images/manifest.json"
+ raw_spectra_manifest_path: Path = REPO_ROOT / "_local/astroprism/raw_spectra/manifest.json"
+ caption_root: Path = REPO_ROOT / "_local/astroprism/hf/datasets/galaxy-descriptions"
+ paired_root: Path = REPO_ROOT / "_local/astroprism/hf/datasets/desi_x_ls_aion"
+ spectrum_caption_bundle: Path = (
+ REPO_ROOT / "_local/astroprism/stage1b_spectrum_caption_accepted"
+ )
+ roster_path: Path = REPO_ROOT / "docs/astroprism/data/overlap_candidates_646.csv.gz"
+ expected_count: int = EXPECTED_CANDIDATES
+ expected_split_counts: dict[str, int] | None = None
+ expected_roster_sha256: str | None = ROSTER_SHA256
+ caption_repo_id: str = CAPTION_REPO_ID
+ caption_revision: str = CAPTION_REVISION
+ paired_repo_id: str = PAIRED_REPO_ID
+ paired_revision: str = PAIRED_REVISION
+ link_mode: str = "auto"
+ overwrite: bool = False
+
+
+@dataclass(frozen=True)
+class _RawAsset:
+ key: str
+ source: Path
+ destination: str
+ dtype: str
+ shape: tuple[int, ...]
+ sha256: str
+ size_bytes: int
+ source_dataset: str | None
+
+
+def sha256_file(path: Path, chunk_size: int = 8 * 1024 * 1024) -> str:
+ digest = hashlib.sha256()
+ with path.open("rb") as handle:
+ while chunk := handle.read(chunk_size):
+ digest.update(chunk)
+ return digest.hexdigest()
+
+
+def sha256_bytes(value: bytes) -> str:
+ return hashlib.sha256(value).hexdigest()
+
+
+def _read_json(path: Path, *, label: str) -> dict[str, Any]:
+ if not path.is_file():
+ raise PackagingValidationError(f"{label} does not exist: {path}")
+ try:
+ value = json.loads(path.read_text(encoding="utf-8"))
+ except (OSError, json.JSONDecodeError) as exc:
+ raise PackagingValidationError(f"cannot read {label} {path}: {exc}") from exc
+ if not isinstance(value, dict):
+ raise PackagingValidationError(f"{label} must contain a JSON object: {path}")
+ return value
+
+
+def _read_jsonl(path: Path, *, label: str) -> list[dict[str, Any]]:
+ if not path.is_file():
+ raise PackagingValidationError(f"{label} does not exist: {path}")
+ result: list[dict[str, Any]] = []
+ try:
+ with path.open("r", encoding="utf-8") as handle:
+ for line_number, line in enumerate(handle, start=1):
+ if not line.strip():
+ raise PackagingValidationError(f"{label} has a blank row at line {line_number}")
+ value = json.loads(line)
+ if not isinstance(value, dict):
+ raise PackagingValidationError(
+ f"{label} line {line_number} must be a JSON object"
+ )
+ result.append(value)
+ except json.JSONDecodeError as exc:
+ raise PackagingValidationError(
+ f"cannot parse {label} {path} at line {exc.lineno}: {exc.msg}"
+ ) from exc
+ return result
+
+
+def _write_json(path: Path, value: Any) -> None:
+ path.parent.mkdir(parents=True, exist_ok=True)
+ path.write_text(
+ json.dumps(value, indent=2, sort_keys=True, ensure_ascii=False, allow_nan=False) + "\n",
+ encoding="utf-8",
+ )
+
+
+def _write_jsonl(path: Path, rows: Iterable[dict[str, Any]]) -> None:
+ path.parent.mkdir(parents=True, exist_ok=True)
+ with path.open("w", encoding="utf-8", newline="\n") as handle:
+ for row in rows:
+ handle.write(
+ json.dumps(row, sort_keys=True, ensure_ascii=False, allow_nan=False) + "\n"
+ )
+
+
+def _safe_child(root: Path, relative: str, *, label: str) -> Path:
+ pure = PurePosixPath(relative)
+ if pure.is_absolute() or not pure.parts or ".." in pure.parts:
+ raise PackagingValidationError(f"unsafe {label} path: {relative!r}")
+ resolved_root = root.resolve()
+ path = (resolved_root / Path(*pure.parts)).resolve()
+ if not path.is_relative_to(resolved_root):
+ raise PackagingValidationError(f"{label} escapes its root: {relative!r}")
+ if not path.is_file():
+ raise PackagingValidationError(f"{label} does not exist: {path}")
+ return path
+
+
+def _require_nonempty_text(value: Any, *, label: str) -> str:
+ if not isinstance(value, str) or not value.strip():
+ raise PackagingValidationError(f"{label} must be a non-empty string")
+ return value
+
+
+def _finite_float(value: Any, *, label: str) -> float:
+ try:
+ result = float(value)
+ except (TypeError, ValueError) as exc:
+ raise PackagingValidationError(f"{label} must be numeric") from exc
+ if not math.isfinite(result):
+ raise PackagingValidationError(f"{label} must be finite")
+ return result
+
+
+def _validate_accepted_rows(
+ rows: list[dict[str, Any]],
+ *,
+ expected_count: int,
+ expected_split_counts: dict[str, int],
+ caption_repo_id: str,
+ caption_revision: str,
+ paired_repo_id: str,
+ paired_revision: str,
+) -> None:
+ if len(rows) != expected_count:
+ raise PackagingValidationError(
+ f"accepted candidates must contain {expected_count} rows, got {len(rows)}"
+ )
+ paired_ids: set[str] = set()
+ caption_ids: set[str] = set()
+ paired_locators: set[tuple[str, int]] = set()
+ caption_locators: set[tuple[str, int]] = set()
+ group_splits: dict[str, set[str]] = defaultdict(set)
+ observed_splits: Counter[str] = Counter()
+
+ for expected_index, row in enumerate(rows):
+ if row.get("candidate_index") != expected_index:
+ raise PackagingValidationError(
+ f"accepted row {expected_index} has candidate_index="
+ f"{row.get('candidate_index')!r}; direct raw-array indexing would be unsafe"
+ )
+ paired_id = _require_nonempty_text(
+ row.get("paired_targetid"), label=f"candidate {expected_index} paired_targetid"
+ )
+ caption_id = _require_nonempty_text(
+ row.get("caption_object_id"), label=f"candidate {expected_index} caption_object_id"
+ )
+ if paired_id in paired_ids or caption_id in caption_ids:
+ raise PackagingValidationError(
+ f"candidate {expected_index} repeats a paired or caption identifier"
+ )
+ paired_ids.add(paired_id)
+ caption_ids.add(caption_id)
+
+ if row.get("acceptance_status") != "accepted":
+ raise PackagingValidationError(f"candidate {expected_index} is not accepted")
+ caption = row.get("caption")
+ if not isinstance(caption, dict):
+ raise PackagingValidationError(f"candidate {expected_index} has no caption object")
+ _require_nonempty_text(
+ caption.get("original"), label=f"candidate {expected_index} caption.original"
+ )
+ _require_nonempty_text(
+ caption.get("summary_original"),
+ label=f"candidate {expected_index} caption.summary_original",
+ )
+ if caption.get("status") != "released_reused":
+ raise PackagingValidationError(
+ f"candidate {expected_index} caption status is not released_reused"
+ )
+
+ split = row.get("split")
+ if split not in expected_split_counts:
+ raise PackagingValidationError(
+ f"candidate {expected_index} has unsupported split {split!r}"
+ )
+ observed_splits[split] += 1
+ group = _require_nonempty_text(
+ row.get("source_group_id"), label=f"candidate {expected_index} source_group_id"
+ )
+ group_splits[group].add(split)
+
+ coordinates = row.get("coordinates_deg")
+ if not isinstance(coordinates, dict):
+ raise PackagingValidationError(
+ f"candidate {expected_index} has no coordinates_deg object"
+ )
+ for field in ("paired_ra", "paired_dec", "caption_ra", "caption_dec"):
+ _finite_float(coordinates.get(field), label=f"candidate {expected_index} {field}")
+ separation = _finite_float(
+ row.get("separation_arcsec"),
+ label=f"candidate {expected_index} separation_arcsec",
+ )
+ if not 0.0 <= separation <= 1.0:
+ raise PackagingValidationError(
+ f"candidate {expected_index} separation_arcsec is outside [0, 1]"
+ )
+
+ for kind, repo_id, revision, locator_set in (
+ ("caption", caption_repo_id, caption_revision, caption_locators),
+ ("paired", paired_repo_id, paired_revision, paired_locators),
+ ):
+ source = row.get(f"{kind}_source")
+ if not isinstance(source, dict):
+ raise PackagingValidationError(
+ f"candidate {expected_index} has no {kind}_source object"
+ )
+ if source.get("repo_id") != repo_id or source.get("revision") != revision:
+ raise PackagingValidationError(
+ f"candidate {expected_index} {kind} source pin does not match the package pin"
+ )
+ shard = _require_nonempty_text(
+ source.get("shard"), label=f"candidate {expected_index} {kind} shard"
+ )
+ try:
+ source_row = int(source["row"])
+ except (KeyError, TypeError, ValueError) as exc:
+ raise PackagingValidationError(
+ f"candidate {expected_index} {kind} row is invalid"
+ ) from exc
+ if source_row < 0 or (shard, source_row) in locator_set:
+ raise PackagingValidationError(
+ f"candidate {expected_index} has an invalid or duplicate {kind} locator"
+ )
+ locator_set.add((shard, source_row))
+ digest = _require_nonempty_text(
+ source.get("sha256"), label=f"candidate {expected_index} {kind} SHA-256"
+ )
+ if len(digest) != 64 or any(char not in "0123456789abcdef" for char in digest):
+ raise PackagingValidationError(
+ f"candidate {expected_index} {kind} SHA-256 is malformed"
+ )
+
+ if dict(observed_splits) != expected_split_counts:
+ raise PackagingValidationError(
+ f"accepted split counts are {dict(observed_splits)}, expected {expected_split_counts}"
+ )
+ leaking_groups = sorted(group for group, splits in group_splits.items() if len(splits) != 1)
+ if leaking_groups:
+ raise PackagingValidationError(
+ f"source groups cross train/held-out splits: {leaking_groups[:5]}"
+ )
+
+
+def _validate_accepted_manifest(
+ path: Path,
+ manifest: dict[str, Any],
+ candidates_path: Path,
+ rows: list[dict[str, Any]],
+ expected_split_counts: dict[str, int],
+ roster_sha256: str,
+) -> None:
+ output = manifest.get("outputs", {}).get("candidates", {})
+ if output.get("sha256") != sha256_file(candidates_path):
+ raise PackagingValidationError("accepted candidate SHA-256 disagrees with its manifest")
+ if output.get("row_count") != len(rows):
+ raise PackagingValidationError("accepted candidate row count disagrees with its manifest")
+ if manifest.get("splits", {}).get("row_counts") != expected_split_counts:
+ raise PackagingValidationError("accepted manifest split counts are not authoritative")
+ if manifest.get("pins", {}).get("roster_sha256") != roster_sha256:
+ raise PackagingValidationError("accepted manifest roster pin mismatch")
+ recorded_groups = manifest.get("splits", {}).get("group_count")
+ actual_groups = len({row["source_group_id"] for row in rows})
+ if recorded_groups != actual_groups:
+ raise PackagingValidationError(
+ f"accepted manifest records {recorded_groups} groups, found {actual_groups}"
+ )
+ if path.parent == candidates_path.parent and output.get("path") != candidates_path.name:
+ raise PackagingValidationError("accepted manifest points to a different candidate file")
+
+
+def _row_groups_for_indices(parquet: Any, indices: list[int]) -> dict[int, list[tuple[int, int]]]:
+ grouped: dict[int, list[tuple[int, int]]] = defaultdict(list)
+ sorted_indices = sorted(indices)
+ pointer = 0
+ offset = 0
+ for group_index in range(parquet.metadata.num_row_groups):
+ count = parquet.metadata.row_group(group_index).num_rows
+ end = offset + count
+ while pointer < len(sorted_indices) and sorted_indices[pointer] < end:
+ source_row = sorted_indices[pointer]
+ if source_row < offset:
+ raise PackagingValidationError("source Parquet row locators are unordered")
+ grouped[group_index].append((source_row, source_row - offset))
+ pointer += 1
+ offset = end
+ if pointer != len(sorted_indices):
+ raise PackagingValidationError(
+ f"source row {sorted_indices[pointer]} is outside the Parquet file"
+ )
+ return grouped
+
+
+def _read_selected_parquet_rows(
+ path: Path, indices: list[int], columns: list[str]
+) -> dict[int, dict[str, Any]]:
+ try:
+ import pyarrow.parquet as pq
+ except ImportError as exc: # pragma: no cover - deployment environment check
+ raise RuntimeError("pyarrow is required to stage AstroPRISM") from exc
+
+ parquet = pq.ParquetFile(path)
+ missing = [name for name in columns if parquet.schema_arrow.get_field_index(name) < 0]
+ if missing:
+ raise PackagingValidationError(f"{path.name} is missing Parquet columns {missing}")
+ selected: dict[int, dict[str, Any]] = {}
+ for group_index, wanted in _row_groups_for_indices(parquet, indices).items():
+ table = parquet.read_row_group(group_index, columns=columns, use_threads=False)
+ for source_row, local_row in wanted:
+ selected[source_row] = table.slice(local_row, 1).to_pylist()[0]
+ return selected
+
+
+def _validate_png(
+ value: Any,
+ *,
+ label: str,
+ expected_size: tuple[int, int] | None = None,
+ expected_mode: str | None = None,
+) -> bytes:
+ if not isinstance(value, bytes) or not value:
+ raise PackagingValidationError(f"{label} is not non-empty PNG bytes")
+ try:
+ from PIL import Image as PILImage
+ except ImportError as exc: # pragma: no cover - deployment environment check
+ raise RuntimeError("Pillow is required to validate embedded previews") from exc
+ try:
+ with PILImage.open(io.BytesIO(value)) as image:
+ image.verify()
+ with PILImage.open(io.BytesIO(value)) as image:
+ if image.format != "PNG":
+ raise PackagingValidationError(f"{label} must be PNG, got {image.format}")
+ if expected_size is not None and image.size != expected_size:
+ raise PackagingValidationError(f"{label} must be {expected_size}, got {image.size}")
+ if expected_mode is not None and image.mode != expected_mode:
+ raise PackagingValidationError(
+ f"{label} must use {expected_mode}, got {image.mode}"
+ )
+ if image.width < 1 or image.height < 1:
+ raise PackagingValidationError(f"{label} has invalid dimensions")
+ except PackagingValidationError:
+ raise
+ except Exception as exc:
+ raise PackagingValidationError(f"{label} is not a valid PNG: {exc}") from exc
+ return value
+
+
+def _validate_source_rows(
+ rows: list[dict[str, Any]],
+ *,
+ caption_root: Path,
+ paired_root: Path,
+) -> dict[int, bytes]:
+ checksum_cache: dict[Path, str] = {}
+ previews: dict[int, bytes] = {}
+
+ def verified_source(root: Path, shard: str, expected_hashes: set[str], label: str) -> Path:
+ if len(expected_hashes) != 1:
+ raise PackagingValidationError(f"conflicting {label} hashes for {shard}")
+ source = _safe_child(root, shard, label=f"{label} shard")
+ observed = checksum_cache.setdefault(source, sha256_file(source))
+ expected = next(iter(expected_hashes))
+ if observed != expected:
+ raise PackagingValidationError(
+ f"{label} source SHA-256 mismatch for {shard}: expected {expected}, got {observed}"
+ )
+ return source
+
+ caption_groups: dict[str, list[dict[str, Any]]] = defaultdict(list)
+ paired_groups: dict[str, list[dict[str, Any]]] = defaultdict(list)
+ for row in rows:
+ caption_groups[row["caption_source"]["shard"]].append(row)
+ paired_groups[row["paired_source"]["shard"]].append(row)
+
+ caption_columns = ["object_id", "survey", "ra", "dec", "image", "caption", "summary"]
+ for shard, group in sorted(caption_groups.items()):
+ source = verified_source(
+ caption_root,
+ shard,
+ {row["caption_source"]["sha256"] for row in group},
+ "caption",
+ )
+ source_rows = _read_selected_parquet_rows(
+ source, [int(row["caption_source"]["row"]) for row in group], caption_columns
+ )
+ for accepted in group:
+ index = int(accepted["candidate_index"])
+ source_row = source_rows[int(accepted["caption_source"]["row"])]
+ if source_row["object_id"] != accepted["caption_object_id"]:
+ raise PackagingValidationError(f"candidate {index} caption source ID mismatch")
+ if source_row["survey"] != "legacy_dr10":
+ raise PackagingValidationError(
+ f"candidate {index} caption survey is not legacy_dr10"
+ )
+ coordinates = accepted["coordinates_deg"]
+ for source_field, accepted_field in (("ra", "caption_ra"), ("dec", "caption_dec")):
+ observed = _finite_float(
+ source_row[source_field], label=f"candidate {index} source {source_field}"
+ )
+ expected = float(coordinates[accepted_field])
+ if abs(observed - expected) > 1e-9:
+ raise PackagingValidationError(
+ f"candidate {index} caption source {source_field} mismatch"
+ )
+ if source_row["caption"] != accepted["caption"]["original"]:
+ raise PackagingValidationError(f"candidate {index} source caption text mismatch")
+ if source_row["summary"] != accepted["caption"]["summary_original"]:
+ raise PackagingValidationError(f"candidate {index} source caption summary mismatch")
+ previews[index] = _validate_png(
+ source_row["image"],
+ label=f"candidate {index} source image",
+ expected_size=(160, 160),
+ expected_mode="RGB",
+ )
+
+ paired_columns = ["object_id_provabgs", "ra_provabgs", "dec_provabgs"]
+ for shard, group in sorted(paired_groups.items()):
+ source = verified_source(
+ paired_root,
+ shard,
+ {row["paired_source"]["sha256"] for row in group},
+ "paired",
+ )
+ source_rows = _read_selected_parquet_rows(
+ source, [int(row["paired_source"]["row"]) for row in group], paired_columns
+ )
+ for accepted in group:
+ index = int(accepted["candidate_index"])
+ source_row = source_rows[int(accepted["paired_source"]["row"])]
+ if str(source_row["object_id_provabgs"]) != accepted["paired_targetid"]:
+ raise PackagingValidationError(f"candidate {index} paired source ID mismatch")
+ coordinates = accepted["coordinates_deg"]
+ for source_field, accepted_field in (
+ ("ra_provabgs", "paired_ra"),
+ ("dec_provabgs", "paired_dec"),
+ ):
+ observed = _finite_float(
+ source_row[source_field], label=f"candidate {index} source {source_field}"
+ )
+ expected = float(coordinates[accepted_field])
+ if abs(observed - expected) > 1e-9:
+ raise PackagingValidationError(
+ f"candidate {index} paired source {source_field} mismatch"
+ )
+
+ if set(previews) != set(range(len(rows))):
+ raise PackagingValidationError("not every accepted row produced an image preview")
+ return previews
+
+
+def _validate_raw_arrays(
+ manifest_path: Path,
+ manifest: dict[str, Any],
+ accepted_rows: list[dict[str, Any]],
+ contract: dict[str, tuple[np.dtype[Any], tuple[int, ...]]],
+ *,
+ destination_root: str,
+) -> list[_RawAsset]:
+ candidates = manifest.get("candidates")
+ if not isinstance(candidates, list) or len(candidates) != len(accepted_rows):
+ raise PackagingValidationError(f"{manifest_path} candidate count mismatch")
+ for index, (raw, accepted) in enumerate(zip(candidates, accepted_rows, strict=True)):
+ for field in ("candidate_index", "paired_targetid", "caption_object_id"):
+ if raw.get(field) != accepted.get(field):
+ raise PackagingValidationError(
+ f"candidate {index} raw/accepted identity mismatch for {field}"
+ )
+
+ outputs = manifest.get("outputs")
+ if not isinstance(outputs, dict) or set(outputs) != set(contract):
+ raise PackagingValidationError(
+ f"{manifest_path} outputs must be exactly {sorted(contract)}"
+ )
+ assets: list[_RawAsset] = []
+ for key, (expected_dtype, trailing_shape) in contract.items():
+ entry = outputs[key]
+ if not isinstance(entry, dict):
+ raise PackagingValidationError(f"raw output {key} has no manifest entry")
+ source = _safe_child(manifest_path.parent, entry.get("path", ""), label=key)
+ expected_shape = (len(accepted_rows), *trailing_shape)
+ try:
+ array = np.load(source, mmap_mode="r", allow_pickle=False)
+ except Exception as exc:
+ raise PackagingValidationError(f"cannot load raw array {source}: {exc}") from exc
+ if array.shape != expected_shape or array.dtype != expected_dtype:
+ raise PackagingValidationError(
+ f"{key} must have shape {expected_shape} and dtype {expected_dtype}, "
+ f"got {array.shape} and {array.dtype}"
+ )
+ if (
+ entry.get("shape") != list(expected_shape)
+ or np.dtype(entry.get("dtype")) != expected_dtype
+ ):
+ raise PackagingValidationError(f"{key} header disagrees with its source manifest")
+ if entry.get("size_bytes") != source.stat().st_size:
+ raise PackagingValidationError(f"{key} size disagrees with its source manifest")
+ digest = sha256_file(source)
+ if entry.get("sha256") != digest:
+ raise PackagingValidationError(f"{key} SHA-256 disagrees with its source manifest")
+
+ if array.dtype.kind == "f":
+ for start in range(0, len(array), 32):
+ chunk = np.asarray(array[start : start + 32])
+ if not np.isfinite(chunk).all():
+ raise PackagingValidationError(f"{key} contains non-finite values")
+ if key.endswith("_ivar") and np.any(chunk < 0):
+ raise PackagingValidationError(f"{key} contains negative inverse variance")
+ if key == "spectrum_wavelength" and np.any(np.diff(array, axis=1) <= 0):
+ raise PackagingValidationError("spectrum_wavelength is not strictly increasing")
+
+ destination = f"{destination_root}/{source.name}"
+ assets.append(
+ _RawAsset(
+ key=key,
+ source=source,
+ destination=destination,
+ dtype=str(array.dtype),
+ shape=tuple(int(value) for value in array.shape),
+ sha256=digest,
+ size_bytes=source.stat().st_size,
+ source_dataset=entry.get("source_dataset"),
+ )
+ )
+ return assets
+
+
+def _copy_or_link(source: Path, destination: Path, *, mode: str) -> str:
+ destination.parent.mkdir(parents=True, exist_ok=True)
+ if mode not in {"auto", "hardlink", "copy"}:
+ raise ValueError("link_mode must be one of: auto, hardlink, copy")
+ if mode in {"auto", "hardlink"}:
+ try:
+ os.link(source, destination)
+ return "hardlink"
+ except OSError:
+ if mode == "hardlink":
+ raise
+ shutil.copy2(source, destination)
+ return "copy"
+
+
+def _portable_raw_manifest(
+ *,
+ kind: str,
+ source_manifest: dict[str, Any],
+ assets: list[_RawAsset],
+ row_count: int,
+) -> dict[str, Any]:
+ source = source_manifest.get("source", {})
+ portable_source = {
+ key: source[key]
+ for key in ("url", "etag", "length_bytes", "hdf5_row_count")
+ if key in source
+ }
+ return {
+ "schema_version": 1,
+ "artifact_type": f"astroprism_portable_raw_{kind}",
+ "row_count": row_count,
+ "array_index_contract": "candidate_index is the first-axis index in every array",
+ "source": portable_source,
+ **({"bands": source_manifest.get("bands")} if kind == "images" else {}),
+ **({"spectrum": source_manifest.get("spectrum")} if kind == "spectra" else {}),
+ "outputs": {
+ asset.key: {
+ "path": Path(asset.destination).name,
+ "dtype": asset.dtype,
+ "shape": list(asset.shape),
+ "size_bytes": asset.size_bytes,
+ "sha256": asset.sha256,
+ "mmap_safe": True,
+ **(
+ {"source_dataset": asset.source_dataset}
+ if asset.source_dataset is not None
+ else {}
+ ),
+ }
+ for asset in assets
+ },
+ }
+
+
+def _validate_bundle_manifest(bundle: Path, records_path: Path, row_count: int) -> dict[str, Any]:
+ manifest = _read_json(bundle / "manifest.json", label="spectrum caption bundle manifest")
+ if manifest.get("row_count") != row_count:
+ raise PackagingValidationError(f"{bundle} manifest row count mismatch")
+ entry = manifest.get("outputs", {}).get("captions.jsonl")
+ if not isinstance(entry, dict) or entry.get("sha256") != sha256_file(records_path):
+ raise PackagingValidationError(f"{bundle} caption checksum mismatch")
+ return manifest
+
+
+def _load_spectrum_bundle(
+ bundle: Path,
+ accepted_rows: list[dict[str, Any]],
+ *,
+ expected_indices: set[int],
+) -> tuple[dict[int, dict[str, Any]], dict[int, bytes]]:
+ records_path = bundle / "captions.jsonl"
+ records = _read_jsonl(records_path, label="spectrum caption candidates")
+ manifest = _validate_bundle_manifest(bundle, records_path, len(records))
+ by_index: dict[int, dict[str, Any]] = {}
+ plots: dict[int, bytes] = {}
+ for record in records:
+ try:
+ index = int(record["candidate_index"])
+ except (KeyError, TypeError, ValueError) as exc:
+ raise PackagingValidationError(
+ "spectrum caption has an invalid candidate_index"
+ ) from exc
+ if index in by_index or index not in expected_indices:
+ raise PackagingValidationError(
+ f"spectrum bundle has unexpected or duplicate candidate {index}"
+ )
+ accepted = accepted_rows[index]
+ if record.get("split") != accepted["split"]:
+ raise PackagingValidationError(f"candidate {index} spectrum split mismatch")
+ for field in ("paired_targetid", "caption_object_id", "source_group_id"):
+ if record.get(field) != accepted.get(field):
+ raise PackagingValidationError(
+ f"candidate {index} spectrum/accepted identity mismatch for {field}"
+ )
+ caption = record.get("spectrum_caption")
+ if not isinstance(caption, dict) or caption.get("status") != "accepted":
+ raise PackagingValidationError(f"candidate {index} spectrum caption is not accepted")
+ _require_nonempty_text(caption.get("original"), label=f"candidate {index} spectrum caption")
+ review = record.get("review")
+ if (
+ not isinstance(review, dict)
+ or review.get("status") != "accepted"
+ or review.get("decision_source") != "user_instruction"
+ ):
+ raise PackagingValidationError(
+ f"candidate {index} spectrum caption lacks an accepted user-directed review"
+ )
+ policy = record.get("training_policy")
+ training_allowed = accepted["split"] == "train"
+ if not isinstance(policy, dict) or policy.get("dataset_ready") is not True:
+ raise PackagingValidationError(
+ f"candidate {index} spectrum caption is not dataset_ready"
+ )
+ if policy.get("gradient_training_allowed") is not training_allowed:
+ raise PackagingValidationError(
+ f"candidate {index} gradient_training_allowed does not match its split"
+ )
+ if policy.get("eligible_for_training") is not training_allowed:
+ raise PackagingValidationError(
+ f"candidate {index} eligible_for_training does not match its split"
+ )
+ if policy.get("eligible_for_evaluation") is not (not training_allowed):
+ raise PackagingValidationError(
+ f"candidate {index} eligible_for_evaluation does not match its split"
+ )
+
+ relative_plot = f"plots/candidate_{index:04d}.png"
+ plot_path = _safe_child(bundle, relative_plot, label="spectrum plot")
+ plot_bytes = _validate_png(plot_path.read_bytes(), label=f"candidate {index} spectrum plot")
+ plot_entry = manifest.get("outputs", {}).get(relative_plot)
+ if not isinstance(plot_entry, dict) or plot_entry.get("sha256") != sha256_bytes(plot_bytes):
+ raise PackagingValidationError(f"candidate {index} spectrum plot checksum mismatch")
+ by_index[index] = record
+ plots[index] = plot_bytes
+
+ if set(by_index) != expected_indices:
+ missing = sorted(expected_indices - set(by_index))
+ raise PackagingValidationError(
+ f"spectrum bundle is incomplete; missing candidate indices {missing[:10]}"
+ )
+ return by_index, plots
+
+
+def _portable_source(source: dict[str, Any]) -> dict[str, Any]:
+ return {
+ key: source[key]
+ for key in ("repo_id", "revision", "shard", "row", "sha256", "survey")
+ if key in source
+ }
+
+
+def _portable_candidate(row: dict[str, Any]) -> dict[str, Any]:
+ index = int(row["candidate_index"])
+ return {
+ "schema_version": 1,
+ "candidate_index": index,
+ "paired_targetid": row["paired_targetid"],
+ "caption_object_id": row["caption_object_id"],
+ "source_group_id": row["source_group_id"],
+ "split": row["split"],
+ "coordinates_deg": row["coordinates_deg"],
+ "separation_arcsec": row["separation_arcsec"],
+ "acceptance_status": row["acceptance_status"],
+ "caption": row["caption"],
+ "caption_source": _portable_source(row["caption_source"]),
+ "paired_source": _portable_source(row["paired_source"]),
+ "raw_assets": {
+ "image_manifest": "assets/raw/images/MANIFEST.json",
+ "spectrum_manifest": "assets/raw/spectra/MANIFEST.json",
+ "array_index": index,
+ },
+ }
+
+
+def _image_parquet_rows(
+ rows: list[dict[str, Any]], previews: dict[int, bytes], split: str
+) -> list[dict[str, Any]]:
+ result = []
+ for row in rows:
+ if row["split"] != split:
+ continue
+ index = int(row["candidate_index"])
+ result.append(
+ {
+ "candidate_index": index,
+ "paired_targetid": row["paired_targetid"],
+ "caption_object_id": row["caption_object_id"],
+ "source_group_id": row["source_group_id"],
+ "split": split,
+ "image": {"bytes": previews[index], "path": None},
+ "caption": row["caption"]["original"],
+ "caption_summary": row["caption"]["summary_original"],
+ "caption_status": row["caption"]["status"],
+ "training_allowed": split == "train",
+ "role": "alignment_training" if split == "train" else "held_out_evaluation",
+ "raw_array_index": index,
+ "raw_image_manifest": "assets/raw/images/MANIFEST.json",
+ "paired_ra_deg": float(row["coordinates_deg"]["paired_ra"]),
+ "paired_dec_deg": float(row["coordinates_deg"]["paired_dec"]),
+ "caption_ra_deg": float(row["coordinates_deg"]["caption_ra"]),
+ "caption_dec_deg": float(row["coordinates_deg"]["caption_dec"]),
+ "separation_arcsec": float(row["separation_arcsec"]),
+ "caption_source_repo": row["caption_source"]["repo_id"],
+ "caption_source_revision": row["caption_source"]["revision"],
+ "caption_source_shard": row["caption_source"]["shard"],
+ "caption_source_row": int(row["caption_source"]["row"]),
+ }
+ )
+ return result
+
+
+def _spectrum_parquet_rows(
+ records: dict[int, dict[str, Any]],
+ plots: dict[int, bytes],
+ *,
+ split: str,
+) -> list[dict[str, Any]]:
+ result = []
+ for index in sorted(records):
+ record = records[index]
+ if record["split"] != split:
+ continue
+ teacher = record.get("teacher_reference", {})
+ catalog = teacher.get("catalog_context", {}) if isinstance(teacher, dict) else {}
+ redshift = catalog.get("redshift")
+ redshift_error = catalog.get("redshift_error")
+ claims = record.get("audit_sidecar", {}).get("claims", [])
+ generation = record.get("generation", {})
+ result.append(
+ {
+ "candidate_index": index,
+ "paired_targetid": record["paired_targetid"],
+ "caption_object_id": record["caption_object_id"],
+ "source_group_id": record["source_group_id"],
+ "split": split,
+ "spectrum_plot": {"bytes": plots[index], "path": None},
+ "caption": record["spectrum_caption"]["original"],
+ "caption_status": "accepted",
+ "review_status": "accepted",
+ "reviewed": True,
+ "dataset_ready": True,
+ "training_allowed": split == "train",
+ "role": "alignment_training" if split == "train" else "held_out_evaluation",
+ "raw_array_index": index,
+ "raw_spectrum_manifest": "assets/raw/spectra/MANIFEST.json",
+ "wavelength_frame": record.get("student_input", {}).get(
+ "wavelength_frame", "observed_vacuum"
+ ),
+ "spectrum_length": int(
+ record.get("student_input", {}).get("spectrum_length", 7781)
+ ),
+ "catalog_redshift": None if redshift is None else float(redshift),
+ "catalog_redshift_error": (
+ None if redshift_error is None else float(redshift_error)
+ ),
+ "claims_json": json.dumps(
+ claims, sort_keys=True, ensure_ascii=False, allow_nan=False
+ ),
+ "generation_json": json.dumps(
+ generation, sort_keys=True, ensure_ascii=False, allow_nan=False
+ ),
+ }
+ )
+ return result
+
+
+def _write_image_parquet(path: Path, rows: list[dict[str, Any]]) -> None:
+ try:
+ from datasets import Dataset, Features, Image, Value, disable_progress_bars
+ except ImportError as exc: # pragma: no cover - deployment environment check
+ raise RuntimeError("datasets is required to write Hugging Face Parquet features") from exc
+
+ disable_progress_bars()
+ features = Features(
+ {
+ "candidate_index": Value("int32"),
+ "paired_targetid": Value("string"),
+ "caption_object_id": Value("string"),
+ "source_group_id": Value("string"),
+ "split": Value("string"),
+ "image": Image(decode=False),
+ "caption": Value("string"),
+ "caption_summary": Value("string"),
+ "caption_status": Value("string"),
+ "training_allowed": Value("bool"),
+ "role": Value("string"),
+ "raw_array_index": Value("int32"),
+ "raw_image_manifest": Value("string"),
+ "paired_ra_deg": Value("float64"),
+ "paired_dec_deg": Value("float64"),
+ "caption_ra_deg": Value("float64"),
+ "caption_dec_deg": Value("float64"),
+ "separation_arcsec": Value("float64"),
+ "caption_source_repo": Value("string"),
+ "caption_source_revision": Value("string"),
+ "caption_source_shard": Value("string"),
+ "caption_source_row": Value("int32"),
+ }
+ )
+ path.parent.mkdir(parents=True, exist_ok=True)
+ # No compression= here. datasets 5.x chooses it per column -- "none" for
+ # columns holding embedded binary, "snappy" elsewhere -- and passing our
+ # own collides with that ("multiple values for keyword argument
+ # 'compression'"). The per-column choice is also the better one.
+ Dataset.from_list(rows, features=features).to_parquet(str(path))
+ _validate_hf_image_feature(path, "image")
+
+
+def _write_spectrum_parquet(path: Path, rows: list[dict[str, Any]]) -> None:
+ try:
+ from datasets import Dataset, Features, Image, Value, disable_progress_bars
+ except ImportError as exc: # pragma: no cover - deployment environment check
+ raise RuntimeError("datasets is required to write Hugging Face Parquet features") from exc
+
+ disable_progress_bars()
+ features = Features(
+ {
+ "candidate_index": Value("int32"),
+ "paired_targetid": Value("string"),
+ "caption_object_id": Value("string"),
+ "source_group_id": Value("string"),
+ "split": Value("string"),
+ "spectrum_plot": Image(decode=False),
+ "caption": Value("string"),
+ "caption_status": Value("string"),
+ "review_status": Value("string"),
+ "reviewed": Value("bool"),
+ "dataset_ready": Value("bool"),
+ "training_allowed": Value("bool"),
+ "role": Value("string"),
+ "raw_array_index": Value("int32"),
+ "raw_spectrum_manifest": Value("string"),
+ "wavelength_frame": Value("string"),
+ "spectrum_length": Value("int32"),
+ "catalog_redshift": Value("float64"),
+ "catalog_redshift_error": Value("float64"),
+ "claims_json": Value("string"),
+ "generation_json": Value("string"),
+ }
+ )
+ path.parent.mkdir(parents=True, exist_ok=True)
+ Dataset.from_list(rows, features=features).to_parquet(str(path))
+ _validate_hf_image_feature(path, "spectrum_plot")
+
+
+def _validate_hf_image_feature(path: Path, field: str) -> None:
+ try:
+ import pyarrow.parquet as pq
+ from datasets import Features, Image
+ except ImportError as exc: # pragma: no cover - deployment environment check
+ raise RuntimeError(
+ "datasets and pyarrow are required to validate Parquet features"
+ ) from exc
+ features = Features.from_arrow_schema(pq.read_schema(path))
+ feature = features.get(field)
+ if not isinstance(feature, Image) or feature.decode is not False:
+ raise PackagingValidationError(
+ f"{path} did not preserve {field} as datasets.Image(decode=False)"
+ )
+
+
+def _dataset_card(
+ *,
+ split_counts: dict[str, int],
+) -> str:
+ configs = [
+ "- config_name: image_caption\n"
+ " data_files:\n"
+ " - split: train\n"
+ " path: data/image_caption/train.parquet\n"
+ " - split: validation\n"
+ " path: data/image_caption/validation.parquet\n"
+ " - split: test\n"
+ " path: data/image_caption/test.parquet",
+ "- config_name: spectrum_caption\n"
+ " data_files:\n"
+ " - split: train\n"
+ " path: data/spectrum_caption/train.parquet\n"
+ " - split: validation\n"
+ " path: data/spectrum_caption/validation.parquet\n"
+ " - split: test\n"
+ " path: data/spectrum_caption/test.parquet",
+ ]
+ return f"""---
+license: other
+language:
+- en
+task_categories:
+- image-to-text
+tags:
+- astronomy
+- multimodal
+- spectra
+- private-research-dataset
+pretty_name: AstroPRISM image and spectrum caption pairs
+configs:
+{chr(10).join(configs)}
+---
+
+# AstroPRISM
+
+Private research staging data for aligning PRISM with astronomical images and spectra. The
+package contains {sum(split_counts.values())} object-level candidates with a fixed,
+source-group-safe split ({split_counts["train"]} train / {split_counts["validation"]}
+validation / {split_counts["test"]} test).
+
+## Configs and training status
+
+| Config | Splits | Rows | Training status |
+|---|---|---:|---|
+| `image_caption` | train, validation, test | {sum(split_counts.values())} | Only train rows are allowed; captions are accepted weak supervision. |
+| `spectrum_caption` | train, validation, test | {sum(split_counts.values())} | All captions are accepted, reviewed, and dataset-ready; only train rows allow gradient training. |
+
+Every spectrum-caption row has `caption_status=accepted`, `review_status=accepted`,
+`reviewed=true`, and `dataset_ready=true`. Dataset readiness means the pair is approved for its
+assigned split; it does not grant gradient-training permission. Only train rows have
+`training_allowed=true`. Validation and test rows have `role=held_out_evaluation` and
+`training_allowed=false` so they remain clean evaluation supervision.
+
+The embedded `image` and `spectrum_plot` columns are inspection previews. Model inputs are the
+calibrated arrays in `assets/raw/images/` and `assets/raw/spectra/`. Released AION embeddings,
+download caches, and encoder weights are intentionally excluded.
+
+## Direct raw-array indexing
+
+`candidate_index` and `raw_array_index` are identical. They index axis 0 of every monolithic
+NumPy array. This remains true in all Parquet splits; indices are deliberately not renumbered.
+
+```python
+from huggingface_hub import hf_hub_download
+import numpy as np
+
+path = hf_hub_download(
+ repo_id="Sand33p/AstroPRISM",
+ repo_type="dataset",
+ filename="assets/raw/images/image_array.npy",
+)
+images = np.load(path, mmap_mode="r")
+image = images[row["candidate_index"]] # float32 [4, 160, 160]
+```
+
+Load pair metadata with an authenticated Hugging Face session:
+
+```python
+from datasets import load_dataset
+
+images = load_dataset("Sand33p/AstroPRISM", "image_caption")
+spectra = load_dataset("Sand33p/AstroPRISM", "spectrum_caption")
+```
+
+See `metadata/candidates.jsonl` for the authoritative cross-modal identity and split mapping,
+`provenance/SOURCES.json` for immutable upstream pins, `LICENSES.md` for source-specific terms,
+and `SHA256SUMS` for integrity verification.
+
+## Important limitations
+
+- Image descriptions are synthetic weak supervision and can contain unsupported interpretation.
+- Spectrum captions were generated from deterministic measured facts and accepted by user instruction.
+- Spectrum flux units were not retained explicitly in the extracted local arrays.
+- Object-level source groups must stay in their assigned split across every derived task.
+- Raw scientific data retain their original survey terms and required acknowledgements.
+"""
+
+
+def _licenses() -> str:
+ return """# Licensing and acknowledgements
+
+This compilation does not apply one blanket license to all payloads. Upstream terms continue to
+govern each component.
+
+- Caption text and RGB review previews come from `astronolan/galaxy-descriptions` at revision
+ `6890dc0c8fc793867c07a6ce4fce11e51c167d6e`. Its dataset card releases descriptions and
+ embeddings under MIT while stating that underlying Legacy Surveys data retain the original
+ data-use policy and required acknowledgements:
+ https://huggingface.co/datasets/astronolan/galaxy-descriptions
+- Cross-match metadata and source locators come from `EiffL/desi_x_ls_aion` at revision
+ `57f7c7575260811ea076163bc743787d75b9aa49`, whose dataset card is CC-BY-4.0:
+ https://huggingface.co/datasets/EiffL/desi_x_ls_aion
+- Calibrated image and spectrum arrays were selected from the official AION tutorial HDF5 source.
+ They derive from Legacy Surveys, DESI, PROVABGS, and Multimodal Universe data; retain all
+ original data-use policies, citations, and acknowledgements.
+- Generated spectrum captions and plots are research artifacts. They do not change the terms of
+ the underlying observations or catalog values.
+
+Users are responsible for reviewing current upstream policies before redistribution or
+publication. Cite the Legacy Surveys, DESI/PROVABGS, Multimodal Universe, AION-1, AION-Search,
+and AstroSpecLM sources as applicable.
+"""
+
+
+def _sources_manifest(
+ accepted_manifest: dict[str, Any],
+ raw_images_manifest: dict[str, Any],
+ raw_spectra_manifest: dict[str, Any],
+ roster_sha256: str,
+) -> dict[str, Any]:
+ return {
+ "schema_version": 1,
+ "roster": {
+ "path": "provenance/overlap_candidates_646.csv.gz",
+ "sha256": roster_sha256,
+ },
+ "caption_source": {
+ "repo_id": CAPTION_REPO_ID,
+ "revision": CAPTION_REVISION,
+ "config": "default",
+ "split": "default",
+ "survey": "legacy_dr10",
+ },
+ "paired_source": {
+ "repo_id": PAIRED_REPO_ID,
+ "revision": PAIRED_REVISION,
+ },
+ "raw_observation_source": {
+ "url": raw_images_manifest.get("source", {}).get("url", RAW_SOURCE_URL),
+ "etag": raw_images_manifest.get("source", {}).get("etag"),
+ "length_bytes": raw_images_manifest.get("source", {}).get("length_bytes"),
+ "image_manifest_source_sha256": None,
+ "spectrum_manifest_source_sha256": None,
+ },
+ "split_policy": accepted_manifest.get("splits", {}),
+ "spectrum_caption_method": {
+ "reference": ASTROSPECLM_PAPER,
+ "status": "synthetic_pending_human_review",
+ },
+ "raw_source_consistency": (
+ raw_images_manifest.get("source", {}).get("etag")
+ == raw_spectra_manifest.get("source", {}).get("etag")
+ and raw_images_manifest.get("source", {}).get("length_bytes")
+ == raw_spectra_manifest.get("source", {}).get("length_bytes")
+ ),
+ }
+
+
+def _all_file_entries(root: Path, *, exclude: set[str] | None = None) -> dict[str, dict[str, Any]]:
+ excluded = exclude or set()
+ result: dict[str, dict[str, Any]] = {}
+ for path in sorted(item for item in root.rglob("*") if item.is_file()):
+ relative = path.relative_to(root).as_posix()
+ if relative in excluded:
+ continue
+ result[relative] = {
+ "sha256": sha256_file(path),
+ "size_bytes": path.stat().st_size,
+ }
+ return result
+
+
+def _write_integrity_files(
+ root: Path,
+ *,
+ config_counts: dict[str, dict[str, int]],
+ link_actions: dict[str, str],
+) -> None:
+ payload = _all_file_entries(root, exclude={"MANIFEST.json", "SHA256SUMS"})
+ manifest = {
+ "schema_version": 1,
+ "artifact_type": "astroprism_huggingface_dataset",
+ "created_at_utc": datetime.now(timezone.utc).isoformat(),
+ "visibility_intent": "private",
+ "array_index_contract": "candidate_index == raw_array_index == first-axis array index",
+ "configs": config_counts,
+ "raw_asset_staging_method": link_actions,
+ "excluded": [
+ "released AION image embeddings",
+ "released AION spectrum embeddings",
+ "Hugging Face download caches",
+ "model weights",
+ ],
+ "files": payload,
+ "integrity_note": "MANIFEST.json and SHA256SUMS are excluded from this self-listing.",
+ }
+ _write_json(root / "MANIFEST.json", manifest)
+ checksummed = _all_file_entries(root, exclude={"SHA256SUMS"})
+ lines = [f"{entry['sha256']} {name}" for name, entry in sorted(checksummed.items())]
+ (root / "SHA256SUMS").write_text("\n".join(lines) + "\n", encoding="utf-8")
+
+
+def _validate_portability(root: Path) -> None:
+ forbidden_bytes = [item.encode("utf-8") for item in FORBIDDEN_PORTABLE_TEXT]
+ for path in sorted(item for item in root.rglob("*") if item.is_file()):
+ relative = path.relative_to(root).as_posix()
+ if any(part == ".cache" for part in Path(relative).parts):
+ raise PackagingValidationError(f"cache path leaked into package: {relative}")
+ if "embedding" in path.name.lower() and path.suffix == ".npy":
+ raise PackagingValidationError(f"released embedding leaked into package: {relative}")
+ if path.suffix.lower() in TEXT_SUFFIXES or path.suffix.lower() == ".parquet":
+ content = path.read_bytes()
+ for forbidden, encoded in zip(FORBIDDEN_PORTABLE_TEXT, forbidden_bytes, strict=True):
+ if encoded in content:
+ raise PackagingValidationError(
+ f"non-portable text {forbidden!r} leaked into {relative}"
+ )
+
+
+def _path_exists(path: Path) -> bool:
+ return path.exists() or path.is_symlink()
+
+
+def _commit_staging(staging: Path, output: Path, *, overwrite: bool) -> None:
+ if not _path_exists(output):
+ os.replace(staging, output)
+ return
+ if not overwrite:
+ raise FileExistsError(f"output already exists: {output}")
+ backup = output.with_name(f".{output.name}.backup-{uuid.uuid4().hex}")
+ os.replace(output, backup)
+ try:
+ os.replace(staging, output)
+ except BaseException:
+ os.replace(backup, output)
+ raise
+ if backup.is_dir() and not backup.is_symlink():
+ shutil.rmtree(backup)
+ else:
+ backup.unlink(missing_ok=True)
+
+
+def prepare_hf_dataset(config: HFDatasetConfig) -> dict[str, Any]:
+ """Validate sources and atomically stage a portable Hugging Face dataset tree."""
+
+ expected_splits = config.expected_split_counts or dict(EXPECTED_SPLIT_COUNTS)
+ if sum(expected_splits.values()) != config.expected_count:
+ raise ValueError("expected split counts must sum to expected_count")
+ if config.link_mode not in {"auto", "hardlink", "copy"}:
+ raise ValueError("link_mode must be one of: auto, hardlink, copy")
+
+ output = Path(config.output_dir).resolve()
+ if _path_exists(output) and not config.overwrite:
+ raise FileExistsError(f"output already exists: {output}")
+
+ roster = Path(config.roster_path)
+ if not roster.is_file():
+ raise PackagingValidationError(f"roster does not exist: {roster}")
+ roster_sha256 = sha256_file(roster)
+ if config.expected_roster_sha256 is not None and roster_sha256 != config.expected_roster_sha256:
+ raise PackagingValidationError(
+ f"roster SHA-256 mismatch: expected {config.expected_roster_sha256}, got {roster_sha256}"
+ )
+ with gzip.open(roster, "rt", encoding="utf-8", newline="") as handle:
+ roster_rows = sum(1 for _line in handle) - 1
+ if roster_rows != config.expected_count:
+ raise PackagingValidationError(
+ f"roster must contain {config.expected_count} data rows, got {roster_rows}"
+ )
+
+ accepted_rows = _read_jsonl(Path(config.accepted_candidates_path), label="accepted candidates")
+ _validate_accepted_rows(
+ accepted_rows,
+ expected_count=config.expected_count,
+ expected_split_counts=expected_splits,
+ caption_repo_id=config.caption_repo_id,
+ caption_revision=config.caption_revision,
+ paired_repo_id=config.paired_repo_id,
+ paired_revision=config.paired_revision,
+ )
+ accepted_manifest = _read_json(
+ Path(config.accepted_manifest_path), label="accepted candidate manifest"
+ )
+ _validate_accepted_manifest(
+ Path(config.accepted_manifest_path),
+ accepted_manifest,
+ Path(config.accepted_candidates_path),
+ accepted_rows,
+ expected_splits,
+ roster_sha256,
+ )
+
+ image_previews = _validate_source_rows(
+ accepted_rows,
+ caption_root=Path(config.caption_root),
+ paired_root=Path(config.paired_root),
+ )
+ raw_images_manifest = _read_json(
+ Path(config.raw_images_manifest_path), label="raw image manifest"
+ )
+ raw_spectra_manifest = _read_json(
+ Path(config.raw_spectra_manifest_path), label="raw spectrum manifest"
+ )
+ image_assets = _validate_raw_arrays(
+ Path(config.raw_images_manifest_path),
+ raw_images_manifest,
+ accepted_rows,
+ IMAGE_ARRAY_CONTRACT,
+ destination_root="assets/raw/images",
+ )
+ spectrum_assets = _validate_raw_arrays(
+ Path(config.raw_spectra_manifest_path),
+ raw_spectra_manifest,
+ accepted_rows,
+ SPECTRUM_ARRAY_CONTRACT,
+ destination_root="assets/raw/spectra",
+ )
+
+ spectrum_indices = {int(row["candidate_index"]) for row in accepted_rows}
+ spectrum_records, spectrum_plots = _load_spectrum_bundle(
+ Path(config.spectrum_caption_bundle),
+ accepted_rows,
+ expected_indices=spectrum_indices,
+ )
+ output.parent.mkdir(parents=True, exist_ok=True)
+ staging = Path(tempfile.mkdtemp(prefix=f".{output.name}.tmp-", dir=output.parent))
+ committed = False
+ try:
+ link_actions: dict[str, str] = {}
+ for asset in [*image_assets, *spectrum_assets]:
+ destination = staging / asset.destination
+ action = _copy_or_link(asset.source, destination, mode=config.link_mode)
+ link_actions[asset.destination] = action
+ if sha256_file(destination) != asset.sha256:
+ raise PackagingValidationError(
+ f"staged raw asset checksum mismatch: {asset.destination}"
+ )
+
+ _write_json(
+ staging / "assets/raw/images/MANIFEST.json",
+ _portable_raw_manifest(
+ kind="images",
+ source_manifest=raw_images_manifest,
+ assets=image_assets,
+ row_count=len(accepted_rows),
+ ),
+ )
+ _write_json(
+ staging / "assets/raw/spectra/MANIFEST.json",
+ _portable_raw_manifest(
+ kind="spectra",
+ source_manifest=raw_spectra_manifest,
+ assets=spectrum_assets,
+ row_count=len(accepted_rows),
+ ),
+ )
+ _write_jsonl(
+ staging / "metadata/candidates.jsonl",
+ (_portable_candidate(row) for row in accepted_rows),
+ )
+ _copy_or_link(
+ roster,
+ staging / "provenance/overlap_candidates_646.csv.gz",
+ mode=config.link_mode,
+ )
+ sources = _sources_manifest(
+ accepted_manifest,
+ raw_images_manifest,
+ raw_spectra_manifest,
+ roster_sha256,
+ )
+ sources["raw_observation_source"]["image_manifest_source_sha256"] = sha256_file(
+ Path(config.raw_images_manifest_path)
+ )
+ sources["raw_observation_source"]["spectrum_manifest_source_sha256"] = sha256_file(
+ Path(config.raw_spectra_manifest_path)
+ )
+ _write_json(staging / "provenance/SOURCES.json", sources)
+
+ for split in ("train", "validation", "test"):
+ pair_rows = _image_parquet_rows(accepted_rows, image_previews, split)
+ if len(pair_rows) != expected_splits[split]:
+ raise PackagingValidationError(f"image Parquet row count mismatch for {split}")
+ _write_image_parquet(staging / f"data/image_caption/{split}.parquet", pair_rows)
+
+ spectrum_counts: dict[str, int] = {}
+ for split in ("train", "validation", "test"):
+ pair_rows = _spectrum_parquet_rows(
+ spectrum_records,
+ spectrum_plots,
+ split=split,
+ )
+ if len(pair_rows) != expected_splits[split]:
+ raise PackagingValidationError(f"spectrum Parquet row count mismatch for {split}")
+ spectrum_counts[split] = len(pair_rows)
+ _write_spectrum_parquet(
+ staging / f"data/spectrum_caption/{split}.parquet",
+ pair_rows,
+ )
+
+ (staging / "README.md").write_text(
+ _dataset_card(split_counts=expected_splits),
+ encoding="utf-8",
+ )
+ (staging / "LICENSES.md").write_text(_licenses(), encoding="utf-8")
+ config_counts = {
+ "image_caption": dict(expected_splits),
+ "spectrum_caption": dict(spectrum_counts),
+ }
+ _write_integrity_files(staging, config_counts=config_counts, link_actions=link_actions)
+ _validate_portability(staging)
+ _commit_staging(staging, output, overwrite=config.overwrite)
+ committed = True
+ finally:
+ if not committed and staging.exists():
+ shutil.rmtree(staging)
+
+ return {
+ "status": "prepared",
+ "output": str(output),
+ "configs": config_counts,
+ "candidate_count": len(accepted_rows),
+ "raw_array_index": "candidate_index",
+ "spectrum_caption_count": len(spectrum_records),
+ "spectrum_captions_reviewed": True,
+ "manifest": str(output / "MANIFEST.json"),
+ "sha256sums": str(output / "SHA256SUMS"),
+ }
+
+
+def parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument(
+ "--output",
+ type=Path,
+ default=REPO_ROOT / "_local/astroprism/hf_upload/AstroPRISM",
+ )
+ parser.add_argument(
+ "--accepted-candidates",
+ type=Path,
+ default=REPO_ROOT / "_local/astroprism/stage1a_accepted/candidates.jsonl",
+ )
+ parser.add_argument(
+ "--accepted-manifest",
+ type=Path,
+ default=REPO_ROOT / "_local/astroprism/stage1a_accepted/manifest.json",
+ )
+ parser.add_argument(
+ "--raw-images-manifest",
+ type=Path,
+ default=REPO_ROOT / "_local/astroprism/raw_images/manifest.json",
+ )
+ parser.add_argument(
+ "--raw-spectra-manifest",
+ type=Path,
+ default=REPO_ROOT / "_local/astroprism/raw_spectra/manifest.json",
+ )
+ parser.add_argument(
+ "--caption-root",
+ type=Path,
+ default=REPO_ROOT / "_local/astroprism/hf/datasets/galaxy-descriptions",
+ )
+ parser.add_argument(
+ "--paired-root",
+ type=Path,
+ default=REPO_ROOT / "_local/astroprism/hf/datasets/desi_x_ls_aion",
+ )
+ parser.add_argument(
+ "--spectrum-caption-bundle",
+ type=Path,
+ default=REPO_ROOT / "_local/astroprism/stage1b_spectrum_caption_accepted",
+ help="Complete promoted train/validation/test spectrum-caption bundle.",
+ )
+ parser.add_argument(
+ "--roster",
+ type=Path,
+ default=REPO_ROOT / "docs/astroprism/data/overlap_candidates_646.csv.gz",
+ )
+ parser.add_argument("--link-mode", choices=("auto", "hardlink", "copy"), default="auto")
+ parser.add_argument("--overwrite", action="store_true")
+ return parser.parse_args()
+
+
+def main() -> int:
+ args = parse_args()
+ summary = prepare_hf_dataset(
+ HFDatasetConfig(
+ output_dir=args.output,
+ accepted_candidates_path=args.accepted_candidates,
+ accepted_manifest_path=args.accepted_manifest,
+ raw_images_manifest_path=args.raw_images_manifest,
+ raw_spectra_manifest_path=args.raw_spectra_manifest,
+ caption_root=args.caption_root,
+ paired_root=args.paired_root,
+ spectrum_caption_bundle=args.spectrum_caption_bundle,
+ roster_path=args.roster,
+ link_mode=args.link_mode,
+ overwrite=args.overwrite,
+ )
+ )
+ print(json.dumps(summary, indent=2, sort_keys=True))
+ return 0
+
+
+if __name__ == "__main__":
+ raise SystemExit(main())
diff --git a/tools/prepare_astroprism_spectrum_caption_pilot.py b/tools/prepare_astroprism_spectrum_caption_pilot.py
new file mode 100644
index 0000000..0c755d0
--- /dev/null
+++ b/tools/prepare_astroprism_spectrum_caption_pilot.py
@@ -0,0 +1,1158 @@
+#!/usr/bin/env python3
+"""Build auditable compact requests or a spectrum-caption review bundle.
+
+The workflow follows AstroSpecLM Figure 1 and Appendix A.5: raw DESI spectra
+are first distilled into a structured fact report, and prose is generated only
+from caption-admissible facts. This pilot never modifies the Stage 1A image
+captions or the current Stage 1B proxy-caption configuration.
+"""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import math
+import os
+import shutil
+import sys
+import tempfile
+from collections import defaultdict
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(REPO_ROOT) not in sys.path:
+ sys.path.insert(0, str(REPO_ROOT))
+
+from src.data.astroprism_spectrum_facts import ( # noqa: E402
+ measure_spectrum_facts,
+)
+
+PAPER_URL = "https://arxiv.org/pdf/2609.07102"
+PAPER_VERSION = "2609.07102v1"
+FACT_EXTRACTOR_VERSION = "astroprism-spectrum-facts-v1"
+COMPACT_PAYLOAD_SCHEMA_VERSION = 1
+COMPACT_POLICY_VERSION = "astroprism-spectrum-caption-conservative-v2"
+TENTATIVE_MIN_MEDIAN_PIXEL_SNR = 2.0
+DN4000_CLASS_BOUNDARIES = (1.3, 1.6)
+DN4000_CLASS_BOUNDARY_SIGMA = 1.0
+DEFAULT_CANDIDATE_INDICES = (0, 1, 2, 3, 4)
+
+COMPACT_FLAG_NAMES = {
+ "desi_arm_overlap_region": "arm_overlap",
+ "telluric_a_band_neighborhood": "telluric",
+ "telluric_b_band_neighborhood": "telluric",
+ "red_sky_residual_sensitive": "red_sky",
+}
+
+# Adapted from the constraints in AstroSpecLM Appendix A.5. We preserve the
+# source citation instead of copying the paper's full prompt into the repo.
+ALIGNMENT_CAPTION_PROMPT = """\
+Act as an experienced astronomical spectroscopist. Write one concise, natural,
+objective caption of the supplied galaxy spectrum using only the compact JSON
+facts. Prioritize continuum morphology, a useful Dn4000 value, and two to five
+principal features. Include approximate observed positions or rest equivalent
+widths only when useful.
+
+The z_cat value and galaxy type are external context, not classifications or a
+redshift measured here. A tentative feature is 3-5 sigma and must be hedged; a
+detected feature is at least 5 sigma. Any fact with cautious=true must be
+hedged or omitted. For a line, cautious means it is tentative or lies in an
+arm-overlap, telluric, or red-sky-sensitive region; for continuum or Dn4000 it
+means the reported trend or Dn4000 class is uncertain. Ca H can blend with
+H-epsilon, and features
+named as blends are integrated blends. Do not infer age, dust, star formation, image
+morphology, or physical mechanisms. Do not mention JSON, fields, confidence
+labels, identifiers, or generation. Return only the caption text.
+"""
+
+
+def sha256_file(path: Path, chunk_size: int = 8 * 1024 * 1024) -> str:
+ digest = hashlib.sha256()
+ with path.open("rb") as handle:
+ while chunk := handle.read(chunk_size):
+ digest.update(chunk)
+ return digest.hexdigest()
+
+
+def sha256_text(value: str) -> str:
+ return hashlib.sha256(value.encode("utf-8")).hexdigest()
+
+
+def _json_safe(value: Any) -> Any:
+ if isinstance(value, dict):
+ return {key: _json_safe(item) for key, item in value.items()}
+ if isinstance(value, list | tuple):
+ return [_json_safe(item) for item in value]
+ if isinstance(value, np.generic):
+ value = value.item()
+ if isinstance(value, float) and not math.isfinite(value):
+ return None
+ return value
+
+
+def _finite_float(value: Any) -> float | None:
+ if value is None:
+ return None
+ numeric = float(value)
+ if not math.isfinite(numeric):
+ return None
+ return numeric
+
+
+def _round_for_payload(value: Any, digits: int) -> float | None:
+ numeric = _finite_float(value)
+ if numeric is None:
+ return None
+ rounded = round(numeric, digits)
+ return 0.0 if rounded == 0 else rounded
+
+
+def _line_is_caption_admissible(line: dict[str, Any]) -> bool:
+ return (
+ line.get("detection_class") in {"detected", "tentative"}
+ and line.get("quality") != "unavailable"
+ and "masked_pixels_in_line_window" not in line.get("quality_flags", [])
+ )
+
+
+def _compact_line_exclusion_reason(
+ line: dict[str, Any], median_pixel_snr: float | None
+) -> str | None:
+ if line.get("quality") == "unavailable" or line.get("detection_class") == "unavailable":
+ return "unavailable"
+ if "masked_pixels_in_line_window" in line.get("quality_flags", []):
+ return "masked_line_window"
+ if line.get("detection_class") == "not_detected":
+ return "not_detected"
+ if line.get("detection_class") not in {"detected", "tentative"}:
+ return "not_caption_admissible"
+ if line["detection_class"] == "tentative" and (
+ median_pixel_snr is None
+ or not math.isfinite(median_pixel_snr)
+ or median_pixel_snr < TENTATIVE_MIN_MEDIAN_PIXEL_SNR
+ ):
+ return "tentative_suppressed_low_snr"
+ return None
+
+
+def _compact_flags(line: dict[str, Any]) -> list[str]:
+ result: list[str] = []
+ for flag in line.get("quality_flags", []):
+ compact = COMPACT_FLAG_NAMES.get(flag, flag)
+ if compact not in result:
+ result.append(compact)
+ return result
+
+
+def _dn4000_class_crosses_boundary(value: float, error: float) -> bool:
+ if error <= 0 or not math.isfinite(error):
+ return True
+ lower = value - DN4000_CLASS_BOUNDARY_SIGMA * error
+ upper = value + DN4000_CLASS_BOUNDARY_SIGMA * error
+ return any(lower < boundary <= upper for boundary in DN4000_CLASS_BOUNDARIES)
+
+
+def _compact_caption_payload(
+ measurement: dict[str, Any], catalog: dict[str, Any]
+) -> dict[str, Any]:
+ """Return only information needed by the caption language model."""
+
+ median_pixel_snr = _finite_float(measurement["quality"]["median_pixel_snr_rest_3850_7000"])
+ payload: dict[str, Any] = {
+ "schema_version": COMPACT_PAYLOAD_SCHEMA_VERSION,
+ "type": "galaxy",
+ "z_cat": _round_for_payload(catalog["redshift"], 6),
+ }
+ if median_pixel_snr is not None:
+ payload["snr"] = _round_for_payload(median_pixel_snr, 2)
+
+ continuum = measurement["continuum"]
+ if continuum["status"] == "measured":
+ continuum_ratio = float(continuum["red_to_blue_ratio"])
+ continuum_error = float(continuum["red_to_blue_ratio_uncertainty"])
+ compact_continuum = {
+ "class": continuum["label"],
+ "ratio": _round_for_payload(continuum_ratio, 3),
+ "err": _round_for_payload(continuum_error, 3),
+ }
+ if continuum_error <= 0 or abs(continuum_ratio - 1.0) / continuum_error < 3.0:
+ compact_continuum["cautious"] = True
+ payload["cont"] = compact_continuum
+
+ dn4000 = measurement["dn4000"]
+ if dn4000["status"] == "measured":
+ dn4000_value = float(dn4000["value"])
+ dn4000_error = float(dn4000["uncertainty"])
+ compact_dn4000 = {
+ "class": dn4000["strength_label"],
+ "value": _round_for_payload(dn4000_value, 3),
+ "err": _round_for_payload(dn4000_error, 3),
+ }
+ if (
+ dn4000_error <= 0
+ or (dn4000_value - 1.0) / dn4000_error < 2.0
+ or _dn4000_class_crosses_boundary(dn4000_value, dn4000_error)
+ ):
+ compact_dn4000["cautious"] = True
+ payload["dn4000"] = compact_dn4000
+
+ compact_lines: list[dict[str, Any]] = []
+ for line in measurement["lines"]:
+ if _compact_line_exclusion_reason(line, median_pixel_snr) is not None:
+ continue
+ flags = _compact_flags(line)
+ compact_line = {
+ "id": f"spec.line.{line['key']}",
+ "name": line["label"],
+ "nature": line["nature"],
+ "conf": line["detection_class"],
+ "sig": _round_for_payload(line.get("significance_sigma"), 2),
+ "obs_A": _round_for_payload(line.get("observed_centroid_angstrom"), 1),
+ "ew_rest_A": _round_for_payload(line.get("equivalent_width_signed_rest_angstrom"), 2),
+ "ew_err_A": _round_for_payload(line.get("equivalent_width_uncertainty_angstrom"), 2),
+ }
+ if flags:
+ compact_line["flags"] = flags
+ if line["detection_class"] == "tentative" or flags:
+ compact_line["cautious"] = True
+ compact_lines.append(compact_line)
+ payload["lines"] = compact_lines
+ return payload
+
+
+def _compact_policy_metadata(measurement: dict[str, Any]) -> dict[str, Any]:
+ median_pixel_snr = _finite_float(measurement["quality"]["median_pixel_snr_rest_3850_7000"])
+ included: list[str] = []
+ excluded: list[dict[str, str]] = []
+ for line in measurement["lines"]:
+ fact_id = f"spec.line.{line['key']}"
+ reason = _compact_line_exclusion_reason(line, median_pixel_snr)
+ if reason is None:
+ included.append(fact_id)
+ else:
+ excluded.append({"fact_id": fact_id, "reason": reason})
+ return {
+ "version": COMPACT_POLICY_VERSION,
+ "tentative_min_median_pixel_snr": TENTATIVE_MIN_MEDIAN_PIXEL_SNR,
+ "dn4000_class_boundaries": list(DN4000_CLASS_BOUNDARIES),
+ "dn4000_class_boundary_sigma": DN4000_CLASS_BOUNDARY_SIGMA,
+ "included_line_fact_ids": included,
+ "excluded_lines": excluded,
+ }
+
+
+def sha256_json(value: Any) -> str:
+ canonical = json.dumps(
+ _json_safe(value), sort_keys=True, separators=(",", ":"), allow_nan=False
+ )
+ return sha256_text(canonical)
+
+
+def _load_json(path: Path) -> dict[str, Any]:
+ with path.open("r", encoding="utf-8") as handle:
+ value = json.load(handle)
+ if not isinstance(value, dict):
+ raise ValueError(f"expected a JSON object: {path}")
+ return value
+
+
+def _load_jsonl(path: Path) -> list[dict[str, Any]]:
+ rows: list[dict[str, Any]] = []
+ with path.open("r", encoding="utf-8") as handle:
+ for line_number, line in enumerate(handle, start=1):
+ if not line.strip():
+ continue
+ row = json.loads(line)
+ if not isinstance(row, dict):
+ raise ValueError(f"{path}:{line_number} is not a JSON object")
+ rows.append(row)
+ return rows
+
+
+def _select_candidate_records(
+ accepted_rows: list[dict[str, Any]],
+ *,
+ requested_indices: list[int] | None,
+ train_limit: int | None,
+) -> list[dict[str, Any]]:
+ if train_limit is not None and requested_indices is not None:
+ raise ValueError("--indices and --train-limit are mutually exclusive")
+
+ accepted_by_index = {int(row["candidate_index"]): row for row in accepted_rows}
+ if len(accepted_by_index) != len(accepted_rows):
+ raise ValueError("accepted candidates contain duplicate candidate_index values")
+
+ if train_limit is not None:
+ if train_limit <= 0:
+ raise ValueError("--train-limit must be positive")
+ records = [row for row in accepted_rows if row["split"] == "train"][:train_limit]
+ if len(records) != train_limit:
+ raise ValueError(
+ f"requested {train_limit} train candidates, but only {len(records)} are available"
+ )
+ return records
+
+ indices = list(
+ requested_indices if requested_indices is not None else DEFAULT_CANDIDATE_INDICES
+ )
+ if not indices or len(indices) != len(set(indices)):
+ raise ValueError("indices must be unique and nonempty")
+ if any(index < 0 for index in indices):
+ raise ValueError("indices must be nonnegative")
+ missing = sorted(set(indices) - set(accepted_by_index))
+ if missing:
+ raise ValueError(f"requested candidate indices are missing: {missing}")
+ return [accepted_by_index[index] for index in indices]
+
+
+def _write_json(path: Path, value: Any) -> None:
+ path.write_text(
+ json.dumps(_json_safe(value), indent=2, sort_keys=True, allow_nan=False) + "\n",
+ encoding="utf-8",
+ )
+
+
+def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None:
+ with path.open("w", encoding="utf-8") as handle:
+ for row in rows:
+ handle.write(json.dumps(_json_safe(row), sort_keys=True, allow_nan=False) + "\n")
+
+
+def _load_catalog_rows(
+ records: list[dict[str, Any]], paired_root: Path
+) -> dict[int, dict[str, Any]]:
+ try:
+ import pyarrow.parquet as pq
+ except ImportError as exc: # pragma: no cover - runtime environment check
+ raise RuntimeError("pyarrow is required to read the paired metadata") from exc
+
+ columns = [
+ "object_id_provabgs",
+ "Z_HP_provabgs",
+ "ZERR_provabgs",
+ "TSNR2_BGS_provabgs",
+ ]
+ grouped: dict[str, list[dict[str, Any]]] = defaultdict(list)
+ for record in records:
+ grouped[record["paired_source"]["shard"]].append(record)
+
+ result: dict[int, dict[str, Any]] = {}
+ for shard, selected in grouped.items():
+ path = paired_root / shard
+ if not path.is_file():
+ raise FileNotFoundError(f"paired metadata shard is missing: {path}")
+ expected_sha256s = {record["paired_source"]["sha256"] for record in selected}
+ if len(expected_sha256s) != 1:
+ raise ValueError(f"conflicting accepted checksums for paired shard: {path}")
+ actual_sha256 = sha256_file(path)
+ if actual_sha256 not in expected_sha256s:
+ raise ValueError(f"paired metadata checksum mismatch: {path}")
+
+ table = pq.read_table(path, columns=columns)
+ for record in selected:
+ row_number = int(record["paired_source"]["row"])
+ row = table.slice(row_number, 1).to_pylist()[0]
+ expected_id = record["paired_targetid"]
+ if str(row["object_id_provabgs"]) != expected_id:
+ raise ValueError(
+ f"candidate {record['candidate_index']}: paired metadata ID "
+ f"{row['object_id_provabgs']!r} != {expected_id!r}"
+ )
+ result[int(record["candidate_index"])] = {
+ "redshift": float(row["Z_HP_provabgs"]),
+ "redshift_error": float(row["ZERR_provabgs"]),
+ "tsnr2_bgs": float(row["TSNR2_BGS_provabgs"]),
+ "source": {
+ "path": str(path.resolve()),
+ "shard": shard,
+ "row": row_number,
+ "sha256": actual_sha256,
+ },
+ }
+ return result
+
+
+def _fact(
+ fact_id: str,
+ name: str,
+ value: Any,
+ *,
+ provenance: str,
+ source_reference: str,
+ evidence_kind: str,
+ required_student_inputs: list[str],
+ quality: str,
+ status: str,
+ evidence_locator: dict[str, Any],
+ allowed_uses: list[str],
+ uncertainty: Any = None,
+ unit: str | None = None,
+ restrictions: list[str] | None = None,
+) -> dict[str, Any]:
+ return {
+ "fact_id": fact_id,
+ "name": name,
+ "value": value,
+ "unit": unit,
+ "uncertainty": uncertainty,
+ "provenance": provenance,
+ "source_reference": source_reference,
+ "evidence_kind": evidence_kind,
+ "required_student_inputs": required_student_inputs,
+ "quality": quality,
+ "status": status,
+ "evidence_locator": evidence_locator,
+ "allowed_uses": allowed_uses,
+ "restrictions": restrictions or [],
+ }
+
+
+def _structured_facts(
+ measurement: dict[str, Any],
+ catalog: dict[str, Any],
+ *,
+ candidate_index: int,
+) -> list[dict[str, Any]]:
+ source_reference = f"raw_spectra:candidate_index={candidate_index}"
+ facts = [
+ _fact(
+ "catalog.object_type",
+ "object_type",
+ "galaxy",
+ provenance="catalog_label",
+ source_reference="EiffL/desi_x_ls_aion dataset sample definition",
+ evidence_kind="external_label",
+ required_student_inputs=["spectrum"],
+ quality="low_confidence",
+ status="supplied",
+ evidence_locator={},
+ allowed_uses=["caption", "alignment_supervision"],
+ restrictions=[
+ "Sample-level galaxy label; no row-level DESI class field was available.",
+ "Must not be described as a direct spectral classification.",
+ ],
+ ),
+ _fact(
+ "catalog.redshift",
+ "redshift",
+ catalog["redshift"],
+ provenance="catalog_label",
+ source_reference=(
+ f"{catalog['source']['shard']}:{catalog['source']['row']}:Z_HP_provabgs"
+ ),
+ evidence_kind="external_label",
+ required_student_inputs=["spectrum"],
+ quality="valid",
+ status="supplied",
+ evidence_locator={},
+ allowed_uses=["caption", "alignment_supervision"],
+ uncertainty=catalog["redshift_error"],
+ restrictions=[
+ "External PROVABGS redshift, not independently recovered here.",
+ "All rest-frame measurement windows are conditioned on this value.",
+ ],
+ ),
+ _fact(
+ "spec.coverage",
+ "wavelength_coverage",
+ measurement["coverage"],
+ provenance="spectral_measurement",
+ source_reference=source_reference,
+ evidence_kind="direct",
+ required_student_inputs=["spectrum"],
+ quality="valid",
+ status="measured",
+ evidence_locator={
+ "observed_wavelength_angstrom": [
+ measurement["coverage"]["observed_min_angstrom"],
+ measurement["coverage"]["observed_max_angstrom"],
+ ]
+ },
+ allowed_uses=["grounded_explanation"],
+ ),
+ _fact(
+ "spec.quality",
+ "median_pixel_snr",
+ measurement["quality"]["median_pixel_snr_rest_3850_7000"],
+ provenance="spectral_measurement",
+ source_reference=source_reference,
+ evidence_kind="derived",
+ required_student_inputs=["spectrum"],
+ quality="valid",
+ status="measured",
+ evidence_locator={"rest_wavelength_angstrom": [3850.0, 7000.0]},
+ allowed_uses=["caption", "grounded_explanation"],
+ restrictions=[measurement["quality"]["definition"]],
+ ),
+ _fact(
+ "spec.continuum",
+ "continuum_morphology",
+ measurement["continuum"],
+ provenance="spectral_measurement",
+ source_reference=source_reference,
+ evidence_kind="derived",
+ required_student_inputs=["spectrum"],
+ quality=(
+ "valid" if measurement["continuum"]["status"] == "measured" else "unavailable"
+ ),
+ status=measurement["continuum"]["status"],
+ evidence_locator={
+ "blue_rest_wavelength_angstrom": measurement["continuum"][
+ "blue_window_rest_angstrom"
+ ],
+ "red_rest_wavelength_angstrom": measurement["continuum"][
+ "red_window_rest_angstrom"
+ ],
+ },
+ allowed_uses=(
+ ["caption", "grounded_explanation"]
+ if measurement["continuum"]["status"] == "measured"
+ else []
+ ),
+ uncertainty=measurement["continuum"]["red_to_blue_ratio_uncertainty"],
+ restrictions=[
+ "Observed f_lambda morphology only; do not infer age, dust, or star formation."
+ ],
+ ),
+ _fact(
+ "spec.dn4000",
+ "dn4000",
+ measurement["dn4000"],
+ provenance="spectral_measurement",
+ source_reference=source_reference,
+ evidence_kind="derived",
+ required_student_inputs=["spectrum"],
+ quality=("valid" if measurement["dn4000"]["status"] == "measured" else "unavailable"),
+ status=measurement["dn4000"]["status"],
+ evidence_locator={
+ "blue_rest_wavelength_angstrom": [3850.0, 3950.0],
+ "red_rest_wavelength_angstrom": [4000.0, 4100.0],
+ },
+ allowed_uses=(
+ ["caption", "grounded_explanation"]
+ if measurement["dn4000"]["status"] == "measured"
+ else []
+ ),
+ uncertainty=measurement["dn4000"]["uncertainty"],
+ restrictions=[
+ "Do not infer stellar age from Dn4000 alone.",
+ "Rest-frame bands are conditioned on the external catalog redshift.",
+ ],
+ ),
+ ]
+
+ for line in measurement["lines"]:
+ observed_interval = [
+ line["line_window_rest_angstrom"][0] * (1.0 + catalog["redshift"]),
+ line["line_window_rest_angstrom"][1] * (1.0 + catalog["redshift"]),
+ ]
+ caption_allowed = _line_is_caption_admissible(line)
+ restrictions = [
+ "Line lookup and rest-frame window are conditioned on the external catalog redshift."
+ ]
+ if line.get("caveat"):
+ restrictions.append(line["caveat"])
+ restrictions.extend(f"Review flag: {flag}." for flag in line.get("quality_flags", []))
+ facts.append(
+ _fact(
+ f"spec.line.{line['key']}",
+ line["label"],
+ line,
+ provenance="spectral_measurement",
+ source_reference=source_reference,
+ evidence_kind="derived",
+ required_student_inputs=["spectrum"],
+ quality=line["quality"],
+ status=line["status"],
+ evidence_locator={
+ "observed_wavelength_angstrom": observed_interval,
+ "rest_wavelength_angstrom": line["line_window_rest_angstrom"],
+ },
+ allowed_uses=(
+ ["caption", "grounded_explanation"]
+ if caption_allowed
+ else ["grounded_explanation"]
+ ),
+ uncertainty=line.get("integrated_strength_uncertainty_source_flux_angstrom"),
+ restrictions=restrictions,
+ )
+ )
+ return facts
+
+
+def _validate_draft(
+ draft: dict[str, Any],
+ facts: list[dict[str, Any]],
+ *,
+ candidate_index: int,
+) -> None:
+ caption = draft.get("student_caption")
+ claims = draft.get("claims")
+ if not isinstance(caption, str) or not caption.strip():
+ raise ValueError(f"candidate {candidate_index}: draft caption is empty")
+ if not isinstance(claims, list) or not claims:
+ raise ValueError(f"candidate {candidate_index}: draft claims are missing")
+ facts_by_id = {fact["fact_id"]: fact for fact in facts}
+ for claim in claims:
+ text = claim.get("text")
+ if not isinstance(text, str) or text not in caption:
+ raise ValueError(
+ f"candidate {candidate_index}: claim text is not present in caption: {text!r}"
+ )
+ fact_ids = claim.get("fact_ids")
+ if not isinstance(fact_ids, list) or not fact_ids:
+ raise ValueError(f"candidate {candidate_index}: claim has no fact IDs")
+ for fact_id in fact_ids:
+ if fact_id not in facts_by_id:
+ raise ValueError(f"candidate {candidate_index}: unknown fact ID {fact_id!r}")
+ if "caption" not in facts_by_id[fact_id]["allowed_uses"]:
+ raise ValueError(
+ f"candidate {candidate_index}: fact {fact_id!r} is not caption-admissible"
+ )
+
+
+def _line_by_key(measurement: dict[str, Any]) -> dict[str, dict[str, Any]]:
+ return {line["key"]: line for line in measurement["lines"]}
+
+
+def _plot_record(
+ path: Path,
+ wavelength: np.ndarray,
+ flux: np.ndarray,
+ ivar: np.ndarray,
+ mask: np.ndarray,
+ measurement: dict[str, Any],
+ *,
+ candidate_index: int,
+ target_id: str,
+) -> None:
+ import matplotlib
+
+ matplotlib.use("Agg")
+ import matplotlib.pyplot as plt
+
+ valid = (~mask) & (ivar > 0)
+ clean_flux = np.asarray(flux, dtype=np.float64).copy()
+ clean_flux[~valid] = np.nan
+ interpolated = np.interp(wavelength, wavelength[valid], clean_flux[valid])
+ kernel = np.ones(7, dtype=np.float64) / 7.0
+ smoothed = np.convolve(interpolated, kernel, mode="same")
+ redshift = measurement["catalog_redshift"]
+ line_map = _line_by_key(measurement)
+
+ figure, axes = plt.subplots(2, 2, figsize=(14, 8), constrained_layout=True)
+ full_axis = axes[0, 0]
+ full_axis.plot(wavelength, clean_flux, color="0.72", linewidth=0.45)
+ full_axis.plot(wavelength, smoothed, color="black", linewidth=0.9)
+ lower, upper = np.nanpercentile(clean_flux, [0.5, 99.5])
+ margin = 0.12 * (upper - lower)
+ full_axis.set_ylim(lower - margin, upper + margin)
+ full_axis.set_xlim(wavelength[0], wavelength[-1])
+ full_axis.set_title(f"Candidate {candidate_index} | TARGETID {target_id} | z={redshift:.5f}")
+ full_axis.set_xlabel("Observed vacuum wavelength (Angstrom)")
+ full_axis.set_ylabel("Flux (source array units)")
+
+ zooms = [
+ ("[O II] / Ca H+K", 3650.0, 4050.0, {"oii_3727_3729", "ca_k_3934", "ca_h_3969"}),
+ ("H-beta / [O III]", 4770.0, 5050.0, {"hbeta_4863", "oiii_4960", "oiii_5008"}),
+ (
+ "H-alpha / [N II] / [S II]",
+ 6480.0,
+ 6790.0,
+ {"nii_6550", "halpha_6565", "nii_6585", "sii_6718", "sii_6733"},
+ ),
+ ]
+ for axis, (title, rest_lo, rest_hi, keys) in zip(
+ (axes[0, 1], axes[1, 0], axes[1, 1]), zooms, strict=False
+ ):
+ observed_lo = rest_lo * (1.0 + redshift)
+ observed_hi = rest_hi * (1.0 + redshift)
+ selected = (wavelength >= observed_lo) & (wavelength <= observed_hi)
+ axis.plot(wavelength[selected], clean_flux[selected], color="0.65", linewidth=0.6)
+ axis.plot(wavelength[selected], smoothed[selected], color="black", linewidth=1.0)
+ local_valid = selected & valid
+ if np.any(local_valid):
+ ylo, yhi = np.nanpercentile(clean_flux[local_valid], [1.0, 99.0])
+ ymargin = 0.15 * max(yhi - ylo, 1e-6)
+ axis.set_ylim(ylo - ymargin, yhi + ymargin)
+ axis.set_xlim(observed_lo, observed_hi)
+ axis.set_title(title)
+ axis.set_xlabel("Observed vacuum wavelength (Angstrom)")
+ axis.set_ylabel("Flux")
+ for key in keys:
+ line = line_map[key]
+ detection = line["detection_class"]
+ color = {"detected": "tab:red", "tentative": "tab:orange"}.get(detection, "tab:blue")
+ expected = line["expected_observed_wavelength_angstrom"]
+ axis.axvline(expected, color=color, alpha=0.65, linewidth=0.9)
+ axis.text(
+ expected,
+ 0.98,
+ line["label"],
+ rotation=90,
+ va="top",
+ ha="right",
+ color=color,
+ fontsize=7,
+ transform=axis.get_xaxis_transform(),
+ )
+ centroid = line.get("observed_centroid_angstrom")
+ if centroid is not None:
+ axis.axvline(centroid, color=color, linestyle=":", alpha=0.8)
+ masked = selected & (~valid)
+ if np.any(masked):
+ axis.scatter(
+ wavelength[masked],
+ np.full(np.count_nonzero(masked), axis.get_ylim()[0]),
+ marker="x",
+ color="magenta",
+ s=12,
+ label="masked/zero ivar",
+ )
+ axis.legend(loc="lower right", fontsize=7)
+
+ figure.suptitle(
+ "AstroPRISM spectrum-caption pilot: gray=raw, black=7-pixel smooth, "
+ "red=detected, orange=tentative, blue=not detected",
+ fontsize=10,
+ )
+ figure.savefig(path, dpi=150)
+ plt.close(figure)
+
+
+def _format_number(value: Any, digits: int = 2) -> str:
+ if value is None or not np.isfinite(value):
+ return "—"
+ return f"{value:.{digits}f}"
+
+
+def _review_markdown(records: list[dict[str, Any]]) -> str:
+ lines = [
+ "# AstroPRISM spectrum-caption review pilot",
+ "",
+ (
+ "This review bundle follows the grounded-fact pipeline in "
+ f"[AstroSpecLM Figure 1 and Appendix A.5]({PAPER_URL}). It is a "
+ "local pilot, not an accepted training artifact."
+ ),
+ "",
+ (
+ "Important split warning: held-out objects are quarantined from training, "
+ "and using review feedback to tune generation contaminates them for clean "
+ "benchmark claims."
+ if any(record["split"] != "train" for record in records)
+ else "All selected candidates are from the training split."
+ ),
+ "",
+ "| Candidate | Split | TARGETID | Status |",
+ "|---:|---|---|---|",
+ ]
+ for record in records:
+ role = record["training_policy"]["role"]
+ lines.append(
+ f"| {record['candidate_index']} | {record['split']} | "
+ f"`{record['paired_targetid']}` | {role} |"
+ )
+ lines.extend(["", "All captions remain `pending_review`; none is training-eligible yet.", ""])
+
+ for record in records:
+ measurement = record["teacher_reference"]["measurement_summary"]
+ continuum_ratio = _format_number(
+ measurement["continuum"].get("red_to_blue_ratio"), 3
+ )
+ continuum_uncertainty = _format_number(
+ measurement["continuum"].get("red_to_blue_ratio_uncertainty"), 3
+ )
+ dn4000_value = _format_number(measurement["dn4000"].get("value"), 3)
+ dn4000_uncertainty = _format_number(
+ measurement["dn4000"].get("uncertainty"), 3
+ )
+ lines.extend(
+ [
+ f"## Candidate {record['candidate_index']} ({record['split']})",
+ "",
+ f"> {record['spectrum_caption']['original']}",
+ "",
+ f"![Candidate {record['candidate_index']} spectrum](plots/candidate_{record['candidate_index']:04d}.png)",
+ "",
+ "Summary facts:",
+ "",
+ f"- External catalog redshift: `{measurement['catalog_redshift']:.7f}`.",
+ f"- Median per-pixel S/N: `{measurement['quality']['median_pixel_snr_rest_3850_7000']:.2f}`.",
+ f"- Continuum: `{measurement['continuum']['label']}`; red/blue f_lambda ratio "
+ f"`{continuum_ratio} ± {continuum_uncertainty}`.",
+ f"- Dn4000: `{dn4000_value} ± {dn4000_uncertainty}` "
+ f"(`{measurement['dn4000']['strength_label']}`).",
+ "",
+ "Detected or tentative line measurements (review flags still apply):",
+ "",
+ "| Feature | Class | Sigma | Observed centroid (A) | Signed rest EW (A) | Review flags |",
+ "|---|---|---:|---:|---:|---|",
+ ]
+ )
+ for line in measurement["lines"]:
+ if line["detection_class"] not in {"detected", "tentative"}:
+ continue
+ flags = ", ".join(line.get("quality_flags", [])) or "—"
+ lines.append(
+ f"| {line['label']} | {line['detection_class']} | "
+ f"{_format_number(line.get('significance_sigma'))} | "
+ f"{_format_number(line.get('observed_centroid_angstrom'), 1)} | "
+ f"{_format_number(line.get('equivalent_width_signed_rest_angstrom'))} | "
+ f"{flags} |"
+ )
+ lines.extend(
+ [
+ "",
+ "Claim links:",
+ "",
+ ]
+ )
+ for claim in record["audit_sidecar"]["claims"]:
+ fact_ids = ", ".join(f"`{item}`" for item in claim["fact_ids"])
+ lines.append(f"- {claim['text']} → {fact_ids}")
+ lines.append("")
+ return "\n".join(lines) + "\n"
+
+
+def parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument("--indices", type=int, nargs="+", default=None)
+ parser.add_argument("--train-limit", type=int, help="select the first N training candidates")
+ parser.add_argument(
+ "--raw-manifest",
+ type=Path,
+ default=REPO_ROOT / "_local/astroprism/raw_spectra/manifest.json",
+ )
+ parser.add_argument(
+ "--accepted-candidates",
+ type=Path,
+ default=REPO_ROOT / "_local/astroprism/stage1a_accepted/candidates.jsonl",
+ )
+ parser.add_argument(
+ "--paired-root",
+ type=Path,
+ default=REPO_ROOT / "_local/astroprism/hf/datasets/desi_x_ls_aion",
+ )
+ parser.add_argument(
+ "--drafts",
+ type=Path,
+ default=REPO_ROOT / "_local/astroprism/stage1b_spectrum_caption_pilot_drafts.json",
+ )
+ parser.add_argument(
+ "--output",
+ type=Path,
+ default=REPO_ROOT / "_local/astroprism/stage1b_spectrum_caption_pilot",
+ )
+ parser.add_argument(
+ "--requests-only",
+ action="store_true",
+ help="write compact generation requests and audit facts without requiring drafts",
+ )
+ parser.add_argument("--overwrite", action="store_true")
+ return parser.parse_args()
+
+
+def main() -> int:
+ args = parse_args()
+ raw_manifest = _load_json(args.raw_manifest)
+ accepted_rows = _load_jsonl(args.accepted_candidates)
+ records = _select_candidate_records(
+ accepted_rows,
+ requested_indices=args.indices,
+ train_limit=args.train_limit,
+ )
+ indices = [int(record["candidate_index"]) for record in records]
+
+ raw_candidates = raw_manifest["candidates"]
+ if len(raw_candidates) != len(accepted_rows):
+ raise ValueError("raw and accepted candidate counts differ")
+ for record in records:
+ index = int(record["candidate_index"])
+ raw = raw_candidates[index]
+ for key in ("candidate_index", "paired_targetid", "caption_object_id"):
+ if raw[key] != record[key]:
+ raise ValueError(f"candidate {index}: raw/accepted identity mismatch for {key}")
+
+ raw_dir = args.raw_manifest.parent
+ arrays: dict[str, np.ndarray] = {}
+ for key in ("spectrum_flux", "spectrum_ivar", "spectrum_mask", "spectrum_wavelength"):
+ output = raw_manifest["outputs"][key]
+ path = raw_dir / output["path"]
+ if sha256_file(path) != output["sha256"]:
+ raise ValueError(f"raw array checksum mismatch: {path}")
+ arrays[key] = np.load(path, mmap_mode="r")
+
+ catalog_rows = _load_catalog_rows(records, args.paired_root)
+ drafts_payload: dict[str, Any] | None = None
+ drafts: dict[str, Any] = {}
+ if not args.requests_only:
+ drafts_payload = _load_json(args.drafts)
+ loaded_drafts = drafts_payload.get("captions")
+ if not isinstance(loaded_drafts, dict):
+ raise ValueError("draft file must contain a captions object")
+ drafts = loaded_drafts
+
+ output = args.output.resolve()
+ if output.exists() and not args.overwrite:
+ raise FileExistsError(f"output already exists: {output}")
+ output.parent.mkdir(parents=True, exist_ok=True)
+ staging = Path(tempfile.mkdtemp(prefix=f".{output.name}.tmp-", dir=output.parent))
+ try:
+ if not args.requests_only:
+ (staging / "plots").mkdir()
+ packaged: list[dict[str, Any]] = []
+ generation_requests: list[dict[str, Any]] = []
+ request_audits: list[dict[str, Any]] = []
+ for accepted in records:
+ index = int(accepted["candidate_index"])
+ catalog = catalog_rows[index]
+ measurement = measure_spectrum_facts(
+ arrays["spectrum_wavelength"][index],
+ arrays["spectrum_flux"][index],
+ arrays["spectrum_ivar"][index],
+ arrays["spectrum_mask"][index],
+ catalog_redshift=catalog["redshift"],
+ )
+ measurement = _json_safe(measurement)
+ facts = _structured_facts(measurement, catalog, candidate_index=index)
+ compact_payload = _compact_caption_payload(measurement, catalog)
+ model_input = {
+ "system_prompt": ALIGNMENT_CAPTION_PROMPT,
+ "compact_spectral_facts": compact_payload,
+ }
+ compact_payload_sha256 = sha256_json(compact_payload)
+ model_input_sha256 = sha256_json(model_input)
+
+ split = accepted["split"]
+ is_held_out = split != "train"
+ generation_requests.append(
+ {
+ "schema_version": 1,
+ "candidate_index": index,
+ "paired_targetid": accepted["paired_targetid"],
+ "split": split,
+ "model_input": model_input,
+ "compact_payload_sha256": compact_payload_sha256,
+ "model_input_sha256": model_input_sha256,
+ }
+ )
+ exposure_policy = _compact_policy_metadata(measurement)
+ request_audits.append(
+ {
+ "schema_version": 1,
+ "artifact_type": "astroprism_spectrum_caption_request_audit",
+ "candidate_index": index,
+ "paired_targetid": accepted["paired_targetid"],
+ "caption_object_id": accepted["caption_object_id"],
+ "source_group_id": accepted["source_group_id"],
+ "split": split,
+ "model_input_sha256": model_input_sha256,
+ "compact_payload_sha256": compact_payload_sha256,
+ "exposure_policy": exposure_policy,
+ "teacher_reference": {
+ "facts": facts,
+ "measurement_summary": measurement,
+ "catalog_context": catalog,
+ },
+ "training_policy": {
+ "eligible_for_training": False,
+ "role": (
+ "held_out_inspection_only" if is_held_out else "generation_request_only"
+ ),
+ },
+ }
+ )
+ if args.requests_only:
+ continue
+
+ draft = drafts.get(str(index))
+ if not isinstance(draft, dict):
+ raise ValueError(f"missing draft caption for candidate {index}")
+ _validate_draft(draft, facts, candidate_index=index)
+ assert drafts_payload is not None
+
+ packaged.append(
+ {
+ "schema_version": 1,
+ "artifact_type": "astroprism_spectrum_caption_candidate",
+ "candidate_index": index,
+ "paired_targetid": accepted["paired_targetid"],
+ "caption_object_id": accepted["caption_object_id"],
+ "source_group_id": accepted["source_group_id"],
+ "split": split,
+ "student_input": {
+ "modalities": ["spectrum"],
+ "raw_manifest": str(args.raw_manifest.resolve()),
+ "array_index": index,
+ "wavelength_frame": "observed_vacuum",
+ "spectrum_length": int(arrays["spectrum_flux"].shape[1]),
+ },
+ "teacher_reference": {
+ "facts": facts,
+ "measurement_summary": measurement,
+ "catalog_context": catalog,
+ },
+ "spectrum_caption": {
+ "original": draft["student_caption"],
+ "revised": None,
+ "status": "pending_review",
+ },
+ "audit_sidecar": {
+ "active_modalities": ["spectrum"],
+ "claims": draft["claims"],
+ "model_exposure_policy": exposure_policy,
+ "limitations": measurement["limitations"],
+ "spectral_grounding_status": "candidate_unverified",
+ },
+ "generation": {
+ "method": "interactive_language_generation_from_compact_facts",
+ "generator": drafts_payload["generator"],
+ "decoding_settings": drafts_payload.get("decoding_settings"),
+ "paper_source": PAPER_URL,
+ "paper_version": PAPER_VERSION,
+ "prompt_adaptation": "AstroPRISM concise A.5 adaptation",
+ "prompt_sha256": sha256_text(ALIGNMENT_CAPTION_PROMPT),
+ "compact_payload_schema_version": COMPACT_PAYLOAD_SCHEMA_VERSION,
+ "compact_policy_version": COMPACT_POLICY_VERSION,
+ "compact_payload_sha256": compact_payload_sha256,
+ "model_input_sha256": model_input_sha256,
+ },
+ "training_policy": {
+ "eligible_for_training": False,
+ "eligible_after_review": not is_held_out,
+ "role": (
+ "held_out_inspection_only"
+ if is_held_out
+ else "training_candidate_pending_review"
+ ),
+ "restrictions": (
+ [
+ "Never include in alignment, SFT, reward, or prompt-tuning data.",
+ "Review feedback contaminates this object for clean benchmark claims.",
+ ]
+ if is_held_out
+ else ["Requires explicit human acceptance before loader integration."]
+ ),
+ },
+ }
+ )
+ _plot_record(
+ staging / "plots" / f"candidate_{index:04d}.png",
+ arrays["spectrum_wavelength"][index],
+ arrays["spectrum_flux"][index],
+ arrays["spectrum_ivar"][index],
+ arrays["spectrum_mask"][index],
+ measurement,
+ candidate_index=index,
+ target_id=accepted["paired_targetid"],
+ )
+
+ _write_jsonl(staging / "generation_requests.jsonl", generation_requests)
+ _write_jsonl(staging / "request_audit.jsonl", request_audits)
+ if not args.requests_only:
+ _write_jsonl(staging / "captions.jsonl", packaged)
+ (staging / "review.md").write_text(_review_markdown(packaged), encoding="utf-8")
+ output_entries = {}
+ for path in sorted(item for item in staging.rglob("*") if item.is_file()):
+ relative = path.relative_to(staging).as_posix()
+ output_entries[relative] = {
+ "sha256": sha256_file(path),
+ "size_bytes": path.stat().st_size,
+ }
+ manifest = {
+ "schema_version": 1,
+ "artifact_type": (
+ "astroprism_spectrum_caption_requests"
+ if args.requests_only
+ else "astroprism_spectrum_caption_pilot"
+ ),
+ "role": (
+ "generation_requests_only_not_training_eligible"
+ if args.requests_only
+ else "review_only_not_training_eligible"
+ ),
+ "created_at_utc": datetime.now(timezone.utc).isoformat(),
+ "candidate_indices": indices,
+ "row_count": len(records),
+ "split_counts": {
+ split: sum(record["split"] == split for record in records)
+ for split in sorted({record["split"] for record in records})
+ },
+ "held_out_candidate_indices": [
+ record["candidate_index"] for record in records if record["split"] != "train"
+ ],
+ "sources": {
+ "raw_manifest": {
+ "path": str(args.raw_manifest.resolve()),
+ "sha256": sha256_file(args.raw_manifest),
+ },
+ "accepted_candidates": {
+ "path": str(args.accepted_candidates.resolve()),
+ "sha256": sha256_file(args.accepted_candidates),
+ },
+ **(
+ {
+ "drafts": {
+ "path": str(args.drafts.resolve()),
+ "sha256": sha256_file(args.drafts),
+ }
+ }
+ if not args.requests_only
+ else {}
+ ),
+ "paired_metadata_root": str(args.paired_root.resolve()),
+ "paper": {"url": PAPER_URL, "version": PAPER_VERSION},
+ },
+ "implementation": {
+ "fact_extractor_version": FACT_EXTRACTOR_VERSION,
+ "compact_payload_schema_version": COMPACT_PAYLOAD_SCHEMA_VERSION,
+ "compact_policy_version": COMPACT_POLICY_VERSION,
+ "tentative_min_median_pixel_snr": TENTATIVE_MIN_MEDIAN_PIXEL_SNR,
+ "dn4000_class_boundaries": list(DN4000_CLASS_BOUNDARIES),
+ "dn4000_class_boundary_sigma": DN4000_CLASS_BOUNDARY_SIGMA,
+ "fact_extractor_sha256": sha256_file(
+ REPO_ROOT / "src/data/astroprism_spectrum_facts.py"
+ ),
+ "packager_sha256": sha256_file(Path(__file__)),
+ "prompt_sha256": sha256_text(ALIGNMENT_CAPTION_PROMPT),
+ },
+ "outputs": output_entries,
+ }
+ _write_json(staging / "manifest.json", manifest)
+
+ if output.exists():
+ if not args.overwrite:
+ raise FileExistsError(f"output already exists: {output}")
+ shutil.rmtree(output)
+ os.replace(staging, output)
+ staging = None
+ finally:
+ if staging is not None and staging.exists():
+ shutil.rmtree(staging)
+
+ print(
+ json.dumps(
+ {
+ "status": "requests_prepared" if args.requests_only else "prepared",
+ "output": str(output),
+ "candidate_indices": indices,
+ "generation_requests": str(output / "generation_requests.jsonl"),
+ "request_audit": str(output / "request_audit.jsonl"),
+ "review": None if args.requests_only else str(output / "review.md"),
+ "captions": (None if args.requests_only else str(output / "captions.jsonl")),
+ },
+ indent=2,
+ sort_keys=True,
+ )
+ )
+ return 0
+
+
+if __name__ == "__main__":
+ raise SystemExit(main())
diff --git a/tools/prepare_astroprism_stage1a.py b/tools/prepare_astroprism_stage1a.py
new file mode 100644
index 0000000..f8a4838
--- /dev/null
+++ b/tools/prepare_astroprism_stage1a.py
@@ -0,0 +1,76 @@
+#!/usr/bin/env python3
+"""Validate and extract the pinned AstroPRISM AION feature parity reference.
+
+The released image embeddings are a parity/regression oracle only, not Stage 1A
+raw-image model inputs.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(REPO_ROOT) not in sys.path:
+ sys.path.insert(0, str(REPO_ROOT))
+
+from src.data.astroprism_stage1a import Stage1AConfig, prepare_stage1a # noqa: E402
+
+
+def parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument(
+ "--roster",
+ type=Path,
+ default=REPO_ROOT / "docs/applications/astroprism_data/overlap_candidates_646.csv.gz",
+ )
+ parser.add_argument(
+ "--paired-root",
+ type=Path,
+ default=REPO_ROOT / "_local/astroprism/hf/datasets/desi_x_ls_aion",
+ )
+ parser.add_argument(
+ "--caption-root",
+ type=Path,
+ default=REPO_ROOT / "_local/astroprism/hf/datasets/galaxy-descriptions",
+ )
+ parser.add_argument(
+ "--output",
+ type=Path,
+ default=REPO_ROOT / "_local/astroprism/stage1a_candidates",
+ )
+ parser.add_argument(
+ "--validate-only",
+ action="store_true",
+ help="Run all checksum/schema/content checks without writing output.",
+ )
+ parser.add_argument(
+ "--overwrite",
+ action="store_true",
+ help="Atomically replace an existing output after validation succeeds.",
+ )
+ parser.add_argument("--batch-size", type=int, default=16)
+ return parser.parse_args()
+
+
+def main() -> int:
+ args = parse_args()
+ summary = prepare_stage1a(
+ Stage1AConfig(
+ roster_path=args.roster,
+ paired_root=args.paired_root,
+ caption_root=args.caption_root,
+ output_dir=args.output,
+ validate_only=args.validate_only,
+ overwrite=args.overwrite,
+ batch_size=args.batch_size,
+ )
+ )
+ print(json.dumps(summary, indent=2, sort_keys=True))
+ return 0
+
+
+if __name__ == "__main__":
+ raise SystemExit(main())
diff --git a/tools/promote_astroprism_spectrum_captions.py b/tools/promote_astroprism_spectrum_captions.py
new file mode 100644
index 0000000..5535765
--- /dev/null
+++ b/tools/promote_astroprism_spectrum_captions.py
@@ -0,0 +1,835 @@
+#!/usr/bin/env python3
+"""Promote reviewed AstroPRISM spectrum captions into one portable artifact.
+
+The source generation bundles are immutable review artifacts. This tool
+validates their manifests and sidecars, overlays the user's review decision,
+and atomically writes a single accepted bundle in candidate-roster order.
+Only the train split is eligible for gradient updates; validation and test are
+accepted inputs for their assigned evaluation splits.
+"""
+
+from __future__ import annotations
+
+import argparse
+import copy
+import hashlib
+import json
+import os
+import re
+import shutil
+import tempfile
+import uuid
+from collections import Counter, defaultdict
+from collections.abc import Iterable
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path, PurePosixPath
+from typing import Any
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+
+EXPECTED_CANDIDATES = 646
+EXPECTED_SPLIT_COUNTS = {"train": 516, "validation": 65, "test": 65}
+REVIEW_DECISION_SOURCE = "user_instruction"
+REVIEW_INSTRUCTION = (
+ "Generate the 130 held-out spectrum captions separately, add them as validation and test, "
+ "and mark all spectrum captions as reviewed and ready for the alignment pipeline."
+)
+ROLE_BY_SPLIT = {
+ "train": "alignment_training",
+ "validation": "alignment_validation",
+ "test": "alignment_test",
+}
+PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
+WINDOWS_ABSOLUTE_PATH = re.compile(r"^[A-Za-z]:[\\/]")
+
+
+class PromotionValidationError(ValueError):
+ """Raised when promotion would violate the accepted-caption contract."""
+
+
+@dataclass(frozen=True)
+class PromotionConfig:
+ """Inputs for one atomic Stage 1B spectrum-caption promotion."""
+
+ output_dir: Path
+ accepted_candidates_path: Path = (
+ REPO_ROOT / "_local/astroprism/stage1a_accepted/candidates.jsonl"
+ )
+ accepted_manifest_path: Path = REPO_ROOT / "_local/astroprism/stage1a_accepted/manifest.json"
+ train_bundle: Path = REPO_ROOT / "_local/astroprism/stage1b_spectrum_caption_train_all"
+ heldout_bundle: Path = REPO_ROOT / "_local/astroprism/stage1b_spectrum_caption_heldout_all"
+ expected_count: int = EXPECTED_CANDIDATES
+ expected_split_counts: dict[str, int] | None = None
+ raw_manifest_reference: str = "../raw_spectra/manifest.json"
+ review_instruction: str = REVIEW_INSTRUCTION
+ link_mode: str = "auto"
+ overwrite: bool = False
+
+
+@dataclass(frozen=True)
+class _Bundle:
+ label: str
+ root: Path
+ manifest: dict[str, Any]
+ manifest_sha256: str
+ records: dict[int, dict[str, Any]]
+ requests: dict[int, dict[str, Any]]
+ audits: dict[int, dict[str, Any]]
+ plots: dict[int, Path]
+
+
+def sha256_file(path: Path, chunk_size: int = 8 * 1024 * 1024) -> str:
+ digest = hashlib.sha256()
+ with path.open("rb") as handle:
+ while chunk := handle.read(chunk_size):
+ digest.update(chunk)
+ return digest.hexdigest()
+
+
+def _sha256_json(value: Any) -> str:
+ payload = json.dumps(
+ value, sort_keys=True, separators=(",", ":"), ensure_ascii=False, allow_nan=False
+ )
+ return hashlib.sha256(payload.encode("utf-8")).hexdigest()
+
+
+def _read_json(path: Path, *, label: str) -> dict[str, Any]:
+ if not path.is_file():
+ raise PromotionValidationError(f"{label} does not exist: {path}")
+ try:
+ value = json.loads(path.read_text(encoding="utf-8"))
+ except (OSError, json.JSONDecodeError) as exc:
+ raise PromotionValidationError(f"cannot read {label} {path}: {exc}") from exc
+ if not isinstance(value, dict):
+ raise PromotionValidationError(f"{label} must contain a JSON object: {path}")
+ return value
+
+
+def _read_jsonl(path: Path, *, label: str) -> list[dict[str, Any]]:
+ if not path.is_file():
+ raise PromotionValidationError(f"{label} does not exist: {path}")
+ result: list[dict[str, Any]] = []
+ try:
+ with path.open("r", encoding="utf-8") as handle:
+ for line_number, line in enumerate(handle, start=1):
+ if not line.strip():
+ raise PromotionValidationError(
+ f"{label} contains a blank line at line {line_number}"
+ )
+ value = json.loads(line)
+ if not isinstance(value, dict):
+ raise PromotionValidationError(
+ f"{label} line {line_number} must be a JSON object"
+ )
+ result.append(value)
+ except json.JSONDecodeError as exc:
+ raise PromotionValidationError(
+ f"cannot parse {label} {path} at line {exc.lineno}: {exc.msg}"
+ ) from exc
+ return result
+
+
+def _write_json(path: Path, value: Any) -> None:
+ path.parent.mkdir(parents=True, exist_ok=True)
+ path.write_text(
+ json.dumps(value, indent=2, sort_keys=True, ensure_ascii=False, allow_nan=False) + "\n",
+ encoding="utf-8",
+ )
+
+
+def _write_jsonl(path: Path, rows: Iterable[dict[str, Any]]) -> None:
+ path.parent.mkdir(parents=True, exist_ok=True)
+ with path.open("w", encoding="utf-8", newline="\n") as handle:
+ for row in rows:
+ handle.write(
+ json.dumps(row, sort_keys=True, ensure_ascii=False, allow_nan=False) + "\n"
+ )
+
+
+def _require_text(value: Any, *, label: str) -> str:
+ if not isinstance(value, str) or not value.strip():
+ raise PromotionValidationError(f"{label} must be a non-empty string")
+ return value
+
+
+def _safe_member(root: Path, relative: str, *, label: str) -> Path:
+ pure = PurePosixPath(relative)
+ if pure.is_absolute() or not pure.parts or ".." in pure.parts:
+ raise PromotionValidationError(f"unsafe {label} path: {relative!r}")
+ resolved_root = root.resolve()
+ path = (resolved_root / Path(*pure.parts)).resolve()
+ if not path.is_relative_to(resolved_root):
+ raise PromotionValidationError(f"{label} escapes its bundle: {relative!r}")
+ if not path.is_file():
+ raise PromotionValidationError(f"{label} does not exist: {path}")
+ return path
+
+
+def _validate_output_entry(
+ root: Path, manifest: dict[str, Any], relative: str, *, label: str
+) -> Path:
+ outputs = manifest.get("outputs")
+ if not isinstance(outputs, dict):
+ raise PromotionValidationError(f"{label} manifest has no outputs object")
+ entry = outputs.get(relative)
+ if not isinstance(entry, dict):
+ raise PromotionValidationError(f"{label} manifest does not declare {relative}")
+ path = _safe_member(root, relative, label=f"{label} output")
+ if entry.get("size_bytes") != path.stat().st_size:
+ raise PromotionValidationError(f"{label} size mismatch for {relative}")
+ digest = sha256_file(path)
+ if entry.get("sha256") != digest:
+ raise PromotionValidationError(f"{label} SHA-256 mismatch for {relative}")
+ return path
+
+
+def _validate_all_manifest_outputs(
+ root: Path, manifest: dict[str, Any], *, label: str
+) -> None:
+ outputs = manifest.get("outputs")
+ if not isinstance(outputs, dict) or not outputs:
+ raise PromotionValidationError(f"{label} manifest has no outputs")
+ for relative in outputs:
+ if not isinstance(relative, str):
+ raise PromotionValidationError(f"{label} manifest output names must be strings")
+ _validate_output_entry(root, manifest, relative, label=label)
+
+
+def _index_rows(
+ rows: list[dict[str, Any]], expected_indices: set[int], *, label: str
+) -> dict[int, dict[str, Any]]:
+ indexed: dict[int, dict[str, Any]] = {}
+ for row in rows:
+ try:
+ index = int(row["candidate_index"])
+ except (KeyError, TypeError, ValueError) as exc:
+ raise PromotionValidationError(f"{label} has an invalid candidate_index") from exc
+ if index in indexed:
+ raise PromotionValidationError(f"{label} repeats candidate {index}")
+ if index not in expected_indices:
+ raise PromotionValidationError(f"{label} has unexpected candidate {index}")
+ indexed[index] = row
+ if set(indexed) != expected_indices:
+ missing = sorted(expected_indices - set(indexed))
+ raise PromotionValidationError(f"{label} is missing candidates {missing[:10]}")
+ return indexed
+
+
+def _validate_accepted_candidates(
+ rows: list[dict[str, Any]],
+ manifest: dict[str, Any],
+ candidates_path: Path,
+ *,
+ expected_count: int,
+ expected_split_counts: dict[str, int],
+) -> None:
+ if len(rows) != expected_count:
+ raise PromotionValidationError(
+ f"accepted candidate roster has {len(rows)} rows, expected {expected_count}"
+ )
+ outputs = manifest.get("outputs")
+ candidate_entry = outputs.get("candidates") if isinstance(outputs, dict) else None
+ if not isinstance(candidate_entry, dict):
+ raise PromotionValidationError("accepted manifest has no outputs.candidates")
+ if candidate_entry.get("row_count") != len(rows):
+ raise PromotionValidationError("accepted candidate row count does not match its manifest")
+ if candidate_entry.get("size_bytes") != candidates_path.stat().st_size:
+ raise PromotionValidationError("accepted candidate size does not match its manifest")
+ if candidate_entry.get("sha256") != sha256_file(candidates_path):
+ raise PromotionValidationError("accepted candidate SHA-256 does not match its manifest")
+ declared_splits = manifest.get("splits", {}).get("row_counts")
+ if declared_splits != expected_split_counts:
+ raise PromotionValidationError("accepted manifest split counts do not match the contract")
+
+ observed: Counter[str] = Counter()
+ paired_ids: set[str] = set()
+ caption_ids: set[str] = set()
+ group_splits: dict[str, set[str]] = defaultdict(set)
+ for expected_index, row in enumerate(rows):
+ if row.get("candidate_index") != expected_index:
+ raise PromotionValidationError(
+ f"accepted row {expected_index} is not in candidate roster order"
+ )
+ if row.get("acceptance_status") != "accepted":
+ raise PromotionValidationError(f"accepted candidate {expected_index} is not accepted")
+ paired_id = _require_text(
+ row.get("paired_targetid"), label=f"candidate {expected_index} paired_targetid"
+ )
+ caption_id = _require_text(
+ row.get("caption_object_id"), label=f"candidate {expected_index} caption_object_id"
+ )
+ if paired_id in paired_ids or caption_id in caption_ids:
+ raise PromotionValidationError(
+ f"accepted candidate {expected_index} repeats an object identity"
+ )
+ paired_ids.add(paired_id)
+ caption_ids.add(caption_id)
+ split = row.get("split")
+ if split not in expected_split_counts:
+ raise PromotionValidationError(
+ f"accepted candidate {expected_index} has unsupported split {split!r}"
+ )
+ observed[split] += 1
+ group = _require_text(
+ row.get("source_group_id"), label=f"candidate {expected_index} source_group_id"
+ )
+ group_splits[group].add(split)
+ if dict(observed) != expected_split_counts:
+ raise PromotionValidationError(
+ f"accepted split counts are {dict(observed)}, expected {expected_split_counts}"
+ )
+ leaking = sorted(group for group, splits in group_splits.items() if len(splits) != 1)
+ if leaking:
+ raise PromotionValidationError(f"source groups cross splits: {leaking[:5]}")
+
+
+def _validate_bundle_record(
+ record: dict[str, Any],
+ request: dict[str, Any],
+ audit: dict[str, Any],
+ accepted: dict[str, Any],
+ *,
+ label: str,
+) -> None:
+ index = int(accepted["candidate_index"])
+ split = accepted["split"]
+ for kind, item in (("caption", record), ("request", request), ("audit", audit)):
+ if item.get("split") != split:
+ raise PromotionValidationError(f"{label} {kind} split mismatch for candidate {index}")
+ if item.get("paired_targetid") != accepted["paired_targetid"]:
+ raise PromotionValidationError(
+ f"{label} {kind} paired_targetid mismatch for candidate {index}"
+ )
+ for field in ("caption_object_id", "source_group_id"):
+ for kind, item in (("caption", record), ("audit", audit)):
+ if item.get(field) != accepted[field]:
+ raise PromotionValidationError(
+ f"{label} {kind} {field} mismatch for candidate {index}"
+ )
+
+ caption = record.get("spectrum_caption")
+ if not isinstance(caption, dict) or caption.get("status") != "pending_review":
+ raise PromotionValidationError(
+ f"{label} spectrum caption {index} must still be pending_review"
+ )
+ _require_text(caption.get("original"), label=f"candidate {index} spectrum caption")
+ policy = record.get("training_policy")
+ if not isinstance(policy, dict) or policy.get("eligible_for_training") is not False:
+ raise PromotionValidationError(
+ f"{label} candidate {index} must be training-ineligible before promotion"
+ )
+ expected_after_review = split == "train"
+ if policy.get("eligible_after_review") is not expected_after_review:
+ raise PromotionValidationError(
+ f"{label} candidate {index} has an invalid pre-promotion split policy"
+ )
+
+ generation = record.get("generation")
+ if not isinstance(generation, dict):
+ raise PromotionValidationError(f"{label} candidate {index} has no generation object")
+ for digest_field in ("model_input_sha256", "compact_payload_sha256"):
+ values = (
+ generation.get(digest_field),
+ request.get(digest_field),
+ audit.get(digest_field),
+ )
+ if not isinstance(values[0], str) or len(set(values)) != 1:
+ raise PromotionValidationError(
+ f"{label} {digest_field} provenance mismatch for candidate {index}"
+ )
+ if record.get("teacher_reference") != audit.get("teacher_reference"):
+ raise PromotionValidationError(
+ f"{label} teacher-reference provenance mismatch for candidate {index}"
+ )
+ claims = record.get("audit_sidecar", {}).get("claims")
+ if not isinstance(claims, list):
+ raise PromotionValidationError(f"{label} candidate {index} has no audit claims")
+
+
+def _load_bundle(
+ root: Path,
+ accepted_rows: list[dict[str, Any]],
+ *,
+ label: str,
+ expected_indices: set[int],
+ expected_split_counts: dict[str, int],
+) -> _Bundle:
+ root = root.resolve()
+ manifest_path = root / "manifest.json"
+ manifest = _read_json(manifest_path, label=f"{label} manifest")
+ if manifest.get("artifact_type") != "astroprism_spectrum_caption_pilot":
+ raise PromotionValidationError(f"{label} has unsupported artifact_type")
+ if manifest.get("role") != "review_only_not_training_eligible":
+ raise PromotionValidationError(f"{label} is not an immutable review bundle")
+ if manifest.get("row_count") != len(expected_indices):
+ raise PromotionValidationError(f"{label} manifest row count mismatch")
+ if manifest.get("split_counts") != expected_split_counts:
+ raise PromotionValidationError(f"{label} manifest split counts mismatch")
+
+ _validate_all_manifest_outputs(root, manifest, label=label)
+ caption_path = _validate_output_entry(root, manifest, "captions.jsonl", label=label)
+ request_path = _validate_output_entry(
+ root, manifest, "generation_requests.jsonl", label=label
+ )
+ audit_path = _validate_output_entry(root, manifest, "request_audit.jsonl", label=label)
+ _validate_output_entry(root, manifest, "review.md", label=label)
+
+ caption_rows = _read_jsonl(caption_path, label=f"{label} captions")
+ request_rows = _read_jsonl(request_path, label=f"{label} generation requests")
+ audit_rows = _read_jsonl(audit_path, label=f"{label} request audits")
+ records = _index_rows(caption_rows, expected_indices, label=f"{label} captions")
+ requests = _index_rows(request_rows, expected_indices, label=f"{label} generation requests")
+ audits = _index_rows(audit_rows, expected_indices, label=f"{label} request audits")
+
+ ordered_indices = [int(row["candidate_index"]) for row in caption_rows]
+ if manifest.get("candidate_indices") != ordered_indices:
+ raise PromotionValidationError(f"{label} candidate_indices do not match captions")
+ expected_heldout = sorted(
+ index for index in expected_indices if accepted_rows[index]["split"] != "train"
+ )
+ if sorted(manifest.get("held_out_candidate_indices", [])) != expected_heldout:
+ raise PromotionValidationError(f"{label} held-out candidate list mismatch")
+
+ plots: dict[int, Path] = {}
+ for index in sorted(expected_indices):
+ record = records[index]
+ request = requests[index]
+ audit = audits[index]
+ _validate_bundle_record(
+ record, request, audit, accepted_rows[index], label=label
+ )
+ relative = f"plots/candidate_{index:04d}.png"
+ plot = _validate_output_entry(root, manifest, relative, label=label)
+ with plot.open("rb") as handle:
+ if handle.read(len(PNG_SIGNATURE)) != PNG_SIGNATURE:
+ raise PromotionValidationError(f"{label} candidate {index} plot is not PNG")
+ plots[index] = plot
+
+ return _Bundle(
+ label=label,
+ root=root,
+ manifest=manifest,
+ manifest_sha256=sha256_file(manifest_path),
+ records=records,
+ requests=requests,
+ audits=audits,
+ plots=plots,
+ )
+
+
+def _is_absolute_local_path(value: str) -> bool:
+ return value.startswith("/") or WINDOWS_ABSOLUTE_PATH.match(value) is not None
+
+
+def _portable_copy(
+ value: Any, *, key: str | None = None, parent: dict[str, Any] | None = None
+) -> Any:
+ """Deep-copy JSON data while replacing machine-local absolute paths."""
+
+ if isinstance(value, dict):
+ return {
+ child_key: _portable_copy(child, key=child_key, parent=value)
+ for child_key, child in value.items()
+ }
+ if isinstance(value, list):
+ return [_portable_copy(item) for item in value]
+ if isinstance(value, str) and _is_absolute_local_path(value):
+ if key == "raw_manifest":
+ return "../raw_spectra/manifest.json"
+ if parent is not None and isinstance(parent.get("shard"), str):
+ return parent["shard"]
+ return Path(value).name
+ return copy.deepcopy(value)
+
+
+def _promoted_record(
+ source: dict[str, Any], *, split: str, source_bundle: str, review_instruction: str
+) -> dict[str, Any]:
+ record = _portable_copy(source)
+ record["spectrum_caption"]["status"] = "accepted"
+ record["review"] = {
+ "status": "accepted",
+ "decision_source": REVIEW_DECISION_SOURCE,
+ "instruction": review_instruction,
+ "source_bundle": source_bundle,
+ "source_caption_status": "pending_review",
+ }
+ is_train = split == "train"
+ record["training_policy"] = {
+ "dataset_ready": True,
+ "gradient_training_allowed": is_train,
+ "eligible_for_training": is_train,
+ "eligible_for_evaluation": not is_train,
+ "role": ROLE_BY_SPLIT[split],
+ "restrictions": (
+ []
+ if is_train
+ else [
+ "Use only in the assigned evaluation split.",
+ "Do not use for gradient updates, SFT, reward training, or prompt tuning.",
+ ]
+ ),
+ }
+ return record
+
+
+def _copy_or_link(source: Path, destination: Path, *, mode: str) -> str:
+ destination.parent.mkdir(parents=True, exist_ok=True)
+ if mode == "copy":
+ shutil.copy2(source, destination)
+ return "copy"
+ if mode not in {"auto", "hardlink"}:
+ raise ValueError("link_mode must be 'auto', 'hardlink', or 'copy'")
+ try:
+ os.link(source, destination)
+ return "hardlink"
+ except OSError:
+ if mode == "hardlink":
+ raise
+ shutil.copy2(source, destination)
+ return "copy"
+
+
+def _portable_lineage(bundle: _Bundle) -> dict[str, Any]:
+ outputs = bundle.manifest["outputs"]
+ return {
+ "artifact_type": bundle.manifest["artifact_type"],
+ "role": bundle.manifest["role"],
+ "row_count": bundle.manifest["row_count"],
+ "split_counts": bundle.manifest["split_counts"],
+ "manifest_sha256": bundle.manifest_sha256,
+ "captions_sha256": outputs["captions.jsonl"]["sha256"],
+ "generation_requests_sha256": outputs["generation_requests.jsonl"]["sha256"],
+ "request_audit_sha256": outputs["request_audit.jsonl"]["sha256"],
+ "implementation": _portable_copy(bundle.manifest.get("implementation", {})),
+ "paper": _portable_copy(bundle.manifest.get("sources", {}).get("paper", {})),
+ }
+
+
+def _mapping_sha256(records: list[dict[str, Any]]) -> str:
+ payload = "".join(
+ f"{row['candidate_index']}\t{row['source_group_id']}\t{row['split']}\n"
+ for row in records
+ )
+ return hashlib.sha256(payload.encode("utf-8")).hexdigest()
+
+
+def _roster_provenance(accepted_manifest: dict[str, Any], *, row_count: int) -> dict[str, Any]:
+ roster = accepted_manifest.get("roster")
+ if not isinstance(roster, dict):
+ roster = {}
+ digest = roster.get("sha256")
+ if not isinstance(digest, str) or not digest:
+ digest = accepted_manifest.get("pins", {}).get("roster_sha256")
+ if not isinstance(digest, str) or not digest:
+ raise PromotionValidationError("accepted manifest does not pin the candidate roster")
+ declared_count = roster.get("row_count", row_count)
+ if declared_count != row_count:
+ raise PromotionValidationError("accepted manifest roster row count mismatch")
+ return {"row_count": row_count, "sha256": digest}
+
+
+def _review_markdown(split_counts: dict[str, int], instruction: str) -> str:
+ return f"""# AstroPRISM accepted spectrum captions
+
+All {sum(split_counts.values())} spectrum captions were explicitly accepted by user instruction.
+The source captions and their generation/audit provenance are preserved in this bundle.
+
+Review decision: `accepted`
+
+Decision source: `{REVIEW_DECISION_SOURCE}`
+
+Instruction: {instruction}
+
+| Split | Rows | Gradient updates | Pipeline role |
+|---|---:|---|---|
+| train | {split_counts['train']} | yes | `{ROLE_BY_SPLIT['train']}` |
+| validation | {split_counts['validation']} | no | `{ROLE_BY_SPLIT['validation']}` |
+| test | {split_counts['test']} | no | `{ROLE_BY_SPLIT['test']}` |
+
+Validation and test captions are reviewed and dataset-ready, but remain evaluation-only.
+"""
+
+
+def _output_entries(root: Path, *, exclude: set[str]) -> dict[str, dict[str, Any]]:
+ result: dict[str, dict[str, Any]] = {}
+ for path in sorted(item for item in root.rglob("*") if item.is_file()):
+ relative = path.relative_to(root).as_posix()
+ if relative in exclude:
+ continue
+ result[relative] = {
+ "size_bytes": path.stat().st_size,
+ "sha256": sha256_file(path),
+ }
+ return result
+
+
+def _write_sha256sums(root: Path) -> None:
+ entries = _output_entries(root, exclude={"SHA256SUMS"})
+ (root / "SHA256SUMS").write_text(
+ "".join(f"{entry['sha256']} {relative}\n" for relative, entry in entries.items()),
+ encoding="utf-8",
+ )
+
+
+def _assert_portable_json(value: Any, *, location: str) -> None:
+ if isinstance(value, dict):
+ for key, child in value.items():
+ _assert_portable_json(child, location=f"{location}.{key}")
+ elif isinstance(value, list):
+ for index, child in enumerate(value):
+ _assert_portable_json(child, location=f"{location}[{index}]")
+ elif isinstance(value, str) and _is_absolute_local_path(value):
+ raise PromotionValidationError(f"absolute local path remains at {location}: {value}")
+
+
+def _validate_portability(root: Path) -> None:
+ for path in sorted(root.rglob("*.json")):
+ _assert_portable_json(
+ json.loads(path.read_text(encoding="utf-8")), location=path.relative_to(root).as_posix()
+ )
+ for path in sorted(root.rglob("*.jsonl")):
+ for line_number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1):
+ _assert_portable_json(
+ json.loads(line),
+ location=f"{path.relative_to(root).as_posix()}:{line_number}",
+ )
+
+
+def _commit_staging(staging: Path, output: Path, *, overwrite: bool) -> None:
+ if not output.exists() and not output.is_symlink():
+ os.replace(staging, output)
+ return
+ if not overwrite:
+ raise FileExistsError(f"output already exists: {output}")
+ backup = output.with_name(f".{output.name}.backup-{uuid.uuid4().hex}")
+ os.replace(output, backup)
+ try:
+ os.replace(staging, output)
+ except BaseException:
+ os.replace(backup, output)
+ raise
+ if backup.is_dir() and not backup.is_symlink():
+ shutil.rmtree(backup)
+ else:
+ backup.unlink(missing_ok=True)
+
+
+def promote_spectrum_captions(config: PromotionConfig) -> dict[str, Any]:
+ """Validate, promote, and atomically publish the accepted caption bundle."""
+
+ split_counts = dict(config.expected_split_counts or EXPECTED_SPLIT_COUNTS)
+ if set(split_counts) != set(ROLE_BY_SPLIT):
+ raise ValueError("expected_split_counts must define train, validation, and test")
+ if any(not isinstance(count, int) or count < 0 for count in split_counts.values()):
+ raise ValueError("split counts must be nonnegative integers")
+ if sum(split_counts.values()) != config.expected_count:
+ raise ValueError("split counts must sum to expected_count")
+ if config.link_mode not in {"auto", "hardlink", "copy"}:
+ raise ValueError("link_mode must be 'auto', 'hardlink', or 'copy'")
+ review_instruction = _require_text(
+ config.review_instruction, label="review_instruction"
+ ).strip()
+ raw_reference = _require_text(
+ config.raw_manifest_reference, label="raw_manifest_reference"
+ )
+ if _is_absolute_local_path(raw_reference):
+ raise ValueError("raw_manifest_reference must be portable, not absolute")
+
+ output = Path(config.output_dir).resolve()
+ if (output.exists() or output.is_symlink()) and not config.overwrite:
+ raise FileExistsError(f"output already exists: {output}")
+ accepted_path = Path(config.accepted_candidates_path).resolve()
+ accepted_manifest_path = Path(config.accepted_manifest_path).resolve()
+ train_root = Path(config.train_bundle).resolve()
+ heldout_root = Path(config.heldout_bundle).resolve()
+ for source_root in (train_root, heldout_root):
+ if output == source_root or output.is_relative_to(source_root):
+ raise ValueError("output_dir must not be inside a source bundle")
+
+ accepted_manifest = _read_json(accepted_manifest_path, label="accepted manifest")
+ accepted_rows = _read_jsonl(accepted_path, label="accepted candidates")
+ _validate_accepted_candidates(
+ accepted_rows,
+ accepted_manifest,
+ accepted_path,
+ expected_count=config.expected_count,
+ expected_split_counts=split_counts,
+ )
+ train_indices = {
+ int(row["candidate_index"]) for row in accepted_rows if row["split"] == "train"
+ }
+ heldout_indices = set(range(config.expected_count)) - train_indices
+ train = _load_bundle(
+ train_root,
+ accepted_rows,
+ label="train bundle",
+ expected_indices=train_indices,
+ expected_split_counts={"train": split_counts["train"]},
+ )
+ heldout = _load_bundle(
+ heldout_root,
+ accepted_rows,
+ label="held-out bundle",
+ expected_indices=heldout_indices,
+ expected_split_counts={
+ "test": split_counts["test"],
+ "validation": split_counts["validation"],
+ },
+ )
+
+ output.parent.mkdir(parents=True, exist_ok=True)
+ staging = Path(tempfile.mkdtemp(prefix=f".{output.name}.tmp-", dir=output.parent))
+ committed = False
+ try:
+ promoted: list[dict[str, Any]] = []
+ requests: list[dict[str, Any]] = []
+ audits: list[dict[str, Any]] = []
+ link_actions: Counter[str] = Counter()
+ for accepted in accepted_rows:
+ index = int(accepted["candidate_index"])
+ source_bundle = train if index in train.records else heldout
+ record = _promoted_record(
+ source_bundle.records[index],
+ split=accepted["split"],
+ source_bundle="train" if source_bundle is train else "heldout",
+ review_instruction=review_instruction,
+ )
+ record["student_input"]["raw_manifest"] = raw_reference
+ promoted.append(record)
+ requests.append(_portable_copy(source_bundle.requests[index]))
+ audits.append(_portable_copy(source_bundle.audits[index]))
+ action = _copy_or_link(
+ source_bundle.plots[index],
+ staging / f"plots/candidate_{index:04d}.png",
+ mode=config.link_mode,
+ )
+ link_actions[action] += 1
+
+ if [row["candidate_index"] for row in promoted] != list(range(config.expected_count)):
+ raise AssertionError("internal promotion ordering invariant failed")
+ _write_jsonl(staging / "captions.jsonl", promoted)
+ _write_jsonl(staging / "generation_requests.jsonl", requests)
+ _write_jsonl(staging / "request_audit.jsonl", audits)
+ (staging / "review.md").write_text(
+ _review_markdown(split_counts, review_instruction), encoding="utf-8"
+ )
+
+ outputs = _output_entries(staging, exclude={"manifest.json", "SHA256SUMS"})
+ manifest = {
+ "schema_version": 1,
+ "artifact_type": "astroprism_spectrum_captions_accepted",
+ "role": "reviewed_alignment_dataset",
+ "created_at_utc": datetime.now(timezone.utc).isoformat(),
+ "row_count": len(promoted),
+ "split_counts": split_counts,
+ "candidate_indices": list(range(config.expected_count)),
+ "roster": _roster_provenance(
+ accepted_manifest, row_count=len(accepted_rows)
+ ),
+ "review": {
+ "status": "accepted",
+ "decision_source": REVIEW_DECISION_SOURCE,
+ "instruction": review_instruction,
+ "accepted_count": len(promoted),
+ },
+ "training_policy": {
+ "dataset_ready": True,
+ "gradient_training_splits": ["train"],
+ "evaluation_splits": ["validation", "test"],
+ "roles": ROLE_BY_SPLIT,
+ },
+ "splits": {
+ "row_counts": split_counts,
+ "mapping_sha256": _mapping_sha256(promoted),
+ "source_groups_cross_splits": False,
+ },
+ "lineage": {
+ "accepted_candidates": {
+ "row_count": len(accepted_rows),
+ "sha256": sha256_file(accepted_path),
+ "manifest_sha256": sha256_file(accepted_manifest_path),
+ },
+ "train_bundle": _portable_lineage(train),
+ "heldout_bundle": _portable_lineage(heldout),
+ },
+ "storage": {
+ "plot_actions": dict(sorted(link_actions.items())),
+ "raw_manifest_reference": raw_reference,
+ },
+ "outputs": outputs,
+ }
+ _write_json(staging / "manifest.json", manifest)
+ _validate_portability(staging)
+ _write_sha256sums(staging)
+ _commit_staging(staging, output, overwrite=config.overwrite)
+ committed = True
+ finally:
+ if not committed and staging.exists():
+ shutil.rmtree(staging)
+
+ return {
+ "status": "accepted",
+ "output_dir": str(output),
+ "candidate_count": len(promoted),
+ "split_counts": split_counts,
+ "gradient_training_count": split_counts["train"],
+ "evaluation_count": split_counts["validation"] + split_counts["test"],
+ "captions_sha256": manifest["outputs"]["captions.jsonl"]["sha256"],
+ "manifest": str(output / "manifest.json"),
+ "sha256sums": str(output / "SHA256SUMS"),
+ }
+
+
+def parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument(
+ "--output",
+ type=Path,
+ default=REPO_ROOT / "_local/astroprism/stage1b_spectrum_caption_accepted",
+ )
+ parser.add_argument(
+ "--accepted-candidates",
+ type=Path,
+ default=REPO_ROOT / "_local/astroprism/stage1a_accepted/candidates.jsonl",
+ )
+ parser.add_argument(
+ "--accepted-manifest",
+ type=Path,
+ default=REPO_ROOT / "_local/astroprism/stage1a_accepted/manifest.json",
+ )
+ parser.add_argument(
+ "--train-bundle",
+ type=Path,
+ default=REPO_ROOT / "_local/astroprism/stage1b_spectrum_caption_train_all",
+ )
+ parser.add_argument(
+ "--heldout-bundle",
+ type=Path,
+ default=REPO_ROOT / "_local/astroprism/stage1b_spectrum_caption_heldout_all",
+ )
+ parser.add_argument("--link-mode", choices=("auto", "hardlink", "copy"), default="auto")
+ parser.add_argument("--overwrite", action="store_true")
+ return parser.parse_args()
+
+
+def main() -> int:
+ args = parse_args()
+ summary = promote_spectrum_captions(
+ PromotionConfig(
+ output_dir=args.output,
+ accepted_candidates_path=args.accepted_candidates,
+ accepted_manifest_path=args.accepted_manifest,
+ train_bundle=args.train_bundle,
+ heldout_bundle=args.heldout_bundle,
+ link_mode=args.link_mode,
+ overwrite=args.overwrite,
+ )
+ )
+ print(json.dumps(summary, indent=2, sort_keys=True))
+ return 0
+
+
+if __name__ == "__main__":
+ raise SystemExit(main())