What’s Next

spectrex is under active development. This page describes the planned directions for the next phase of work. None of these are committed to a release schedule; they reflect current thinking and may evolve.

Regularisation improvements

The most concrete near-term item is an update in the regularisation.

The original motivation for regularisation in spectrex is that a small number of PCA basis components should suffice to represent any stellar spectrum — sparsity in the spectral coefficients within each source. That assumption maps to the lasso penalty \(\sum_k \|\mathbf{a}_k\|_1\).

The current JAXProximalSolver implements the group lasso \(\sum_k \|\mathbf{a}_k\|_2\) instead, which promotes source-level sparsity — driving the entire coefficient vector of an absent source to zero. This is useful for deblending in crowded fields but does not enforce sparsity in the spectral basis per source.

We will therefore update our solver to use sparsity of the spectral dimension per source (L1 penalty) and group lasso will remain for the spatial distribution of sources. (More on that later).

Planned additions to JAXProximalSolver:

Penalty	Effect
Lasso \(\sum_k |\mathbf{a}_k|_1\)	Few basis components active per source
Sparse group lasso \(\alpha \sum_k |\mathbf{a}_k|_2 + (1{-}\alpha)\sum_k |\mathbf{a}_k|_1\)	Both: few active sources and few components per source

See Solvers for a full explanation of the three penalties and their proximal operators.

Uncertainty quantification

FISTA currently returns a MAP point estimate, i.e no error bars. For downstream science (stellar populations, redshifts, variability) credible intervals on the reconstructed spectra are essential.

Planned approaches, from simplest to most rigorous:

• Support-conditioned Gaussian — fix the active set (non-zero coefficients), compute a Gaussian posterior on that reduced problem. Cheap; ignores support uncertainty.

• MCMC / variational inference — full posterior via NumPyro, directly on the JAX compute graph. This is best for scientific outcomes.

Scaling to the full detector

The current implementation targets stamp-sized cutouts (~500 × 20 px) primarily for the mathematical developments. A full NIRISS detector is 2048 × 2048 px with potentially thousands of sources and heavily overlapping Order A traces.

Key challenges:

• Memory — the compact trace representation in JAXOperator is image-size-independent, but the stacked coefficient matrix grows with source count K.

• Overlapping traces — two sources whose Order A streaks overlap must be solved jointly; independent tile solves introduce contamination bias at tile edges.

• Order B cross-contamination — a source’s undispersed zeroth-order spot lands ~216 rows below the source position; it can contaminate a tile that does not contain the source itself.

The JAXOperator compact trace design was chosen to accommodate this path; no fundamental redesign is anticipated. The development is primarily benchmarking and putting guardrails on the Scipy based operator.

Source detection with a weak or absent positional prior

The current pipeline requires a source catalog with known positions (e.g. from a direct image). While this is not a challenge to detect sources on the direct image, it would be good to account for situations when positions are uncertain or unavailable, unmodelled sources produce residuals that bias neighbour reconstructions.

Planned approaches in order of increasing complexity:

1. Iterative catalog refinement — solve → detect residuals → add sources → re-solve. The undispersed Order B spots (compact, ~4 px, ~18× brighter per pixel than Order A) are reliable position anchors via the known ~216-row offset.

2. Joint position and spectrum optimisation — treat source positions as continuous latent variables; optimise jointly via JAX autodiff through get_trace.

3. Sparse recovery in source-position space — L0/L1 penalty on a dense grid of candidate positions; research-grade (a paper on its own).

Full-detector scaling (above) is a practical prerequisite for field-wide source detection.