Future modules and capabilities roadmap

This is the living build note for Jaxstro’s future numerical-methods, scientific-ML, and supporting infrastructure. It preserves the package-wide assessment and architectural advice that motivated the roadmap, then turns the recommended build order into checklists.

The roadmap is not a promise to implement every named method. A new capability still has to satisfy the thin-foundation admission rule: a generic mathematical contract, a reusable public API, explicit JAX and AD semantics, independent validation, and evidence that Jaxstro is the right owner.

Summary¶

• Jaxstro is already a strong shared numerical-science foundation: units, quantities, coordinates, robust numerics, spatial algorithms, spectra, atmospheres, parameter transforms, scientific contracts, validation, and provenance.

• Its numerical breadth is substantial; its distinguishing strength is not merely algorithm count, but explicit JAX, AD, boundary, failure, and evidence contracts.

• It is not yet a general scientific-ML toolkit. It has the necessary foundations—Equinox parameterization, autodiff products, operators, losses, RNG, and provenance—but lacks a coherent data, training, and checkpoint layer.

• The target is to make Jaxstro the primary reusable numerical and scientific-ML substrate while keeping inference, information geometry, calibration, SBI, and experimental design in Informax.

• The highest-value additions are a narrow jaxstro.ml layer, matrix-free iterative solvers, vector nonlinear and implicit solves, quasi-Monte Carlo, adaptive quadrature, and signal-processing primitives.

• Jaxstro should integrate with Equinox, Optax, Diffrax, NumPyro, and specialized solver libraries where appropriate—not attempt to replace all of them.

Current package scale¶

These are exported names, not algorithm counts; exports also include types, constants, and status identifiers. The values record the inventory audited on 2026-07-12 and must be refreshed when the public API or generated registries change.

The important interpretation is that Jaxstro has broad implementation coverage, but callable-level contract classification has not yet caught up with that breadth.

Existing methods and features¶

Checked items describe implemented capability families. Their detailed maturity and evidence remain governed by the Jaxstro package assessment scorecard, Scientific contract registry, and Validation.

Numerical accuracy and defensive primitives¶

• Neumaier compensated scalar, vector, array, and dot-product summation.

• Safe logarithm, exponential, division, log-sum-exp, relative error, and convergence helpers.

• Finiteness, positivity, range, and monotonicity validation.

• Concrete-input validation that does not silently fabricate traced assertions.

• Explicit float64 and highest-matmul-precision configuration.

This is valuable infrastructure because numerical stability is handled once rather than rediscovered inside every downstream physics package.

Automatic differentiation¶

• JVP and VJP.

• Jacobian-vector and vector-Jacobian products.

• Hessian-vector products.

• Gauss–Newton products.

• Generic empirical Fisher-style products.

• Independent finite-difference gradients and Jacobians.

• Directional-derivative comparison.

• Explicit gradient contracts: smooth, known-zero, blocked, surrogate, and validation-only.

This is one of Jaxstro’s strongest areas. It distinguishes “JAX can differentiate this expression” from “this derivative means what the scientific claim says it means.”

Scalar rootfinding and inversion¶

• Bracket expansion.

• Bisection and batched bisection.

• Newton and Newton-with-gradient execution.

• Monotone inverse interpolation and PPF inversion.

• Safeguarded inverse-quadratic, secant, and midpoint rootfinding.

• Low-level bracket initialization, proposal, update, and checkpointable advancement.

• Fixed-shape root telemetry with typed termination status.

• Explicit invalid-trial exclusion through valid=False.

• Physical per-lane evaluation skipping through the mapped solver.

• Separately gated implicit-function derivatives.

• Fail-closed certificates for convergence, residual, bracket width, uniqueness assumptions, finite slope, and conditioning.

This supports equilibria, event locations, inverse constitutive relations, timestep controllers, and differentiable implicit scientific models.

Interpolation and approximation¶

• Linear one-dimensional interpolation.

• Cubic Hermite interpolation.

• PCHIP-style monotone cubic interpolation.

• Natural cubic splines.

• Tabulated-function PyTrees.

• Static-rank regular-grid interpolation.

• Bilinear and trilinear interpolation.

• Clamp, fill, and eager-rejection boundary policies.

• B-spline bases and design matrices.

• Direct and de Boor spline evaluation.

• Derivatives, antiderivatives, and definite integrals.

• Roughness penalties.

• Fixed-knot least-squares fitting.

• Quantile-based adaptive knot placement.

• Tensor-product design matrices.

This is already a strong foundation for stellar tracks, atmosphere grids, tabulated equations of state, calibration surfaces, and surrogate representations.

Integration and quadrature¶

• Trapezoidal and cumulative trapezoidal integration.

• Simpson and cumulative Simpson integration.

• Gauss–Legendre quadrature.

• Probabilists’ Gauss–Hermite quadrature.

• Gauss–Laguerre quadrature.

• Clenshaw–Curtis quadrature.

• Hermite basis and expansion coefficients.

Fixed-node construction and differentiable integrand evaluation have deliberately separate contracts.

Linear algebra and operators¶

• Norms, projection, and condition-number diagnostics.

• Weighted least squares.

• QR and SVD solves.

• Covariance and correlation matrices.

• Positive-definiteness tests and diagonal jitter selection.

• Dense, diagonal, scaled, sum, product, transpose, and block-diagonal operators.

• matvec, reverse multiplication, composition, and dense realization.

This is good for small dense scientific problems and matrix-free composition. It is not yet sufficient for large sparse or Krylov problems.

Optimization helpers¶

• Squared, Huber, and pseudo-Huber losses.

• Weighted objective summaries.

• Fixed-iteration Armijo backtracking.

• Gradient and relative-step norms.

• Convergence summaries and typed line-search results.

The boundary is appropriately narrow: these are trustworthy optimizer ingredients, not a homegrown replacement for Optax or specialized optimization libraries.

Differential equations¶

• Euler.

• Midpoint/RK2.

• RK4.

• Fixed-step dispatch using lax.scan.

• Velocity-Verlet integration.

• Typed result records.

This is useful for transparent fixed-step algorithms and validation fixtures. It is not an adaptive ODE/SDE solver stack.

Probability, sampling, and random methods¶

• Normal log-density, CDF, and PPF.

• Lognormal log-density, CDF, and PPF.

• Truncated-normal log-density, CDF, and PPF.

• Finite power-law normalization, log-density, CDF, and PPF.

• Smooth removable-singularity behavior through the logarithmic power-law limit.

• Differentiable inverse-CDF sampling.

• Stratified uniforms.

• Explicit PRNG key streams and fold-in streams.

• Seed manifests.

• Systematic, stratified, and residual resampling.

These support population synthesis and Monte Carlo mechanics without turning Jaxstro into a probabilistic-programming framework.

Grids, conservative methods, and meshes¶

• Logarithmic grids.

• Geometric bin edges and centers.

• Conservative rebinning of integrated values.

• Structured one-dimensional meshes.

• Cell and face geometry.

• Neighbor stencils.

• One-dimensional finite-volume divergence.

• Cell-to-face averaging.

• Conservative remapping.

This is a credible first finite-volume layer, but it remains intentionally one-dimensional.

Special functions¶

• Planck functions in wavelength and frequency form.

• Stable logarithmic Planck kernels.

• Log-weight normalization.

• Legendre, Chebyshev, Laguerre, and Hermite polynomial bases.

Geometry, astrometry, and coordinates¶

• Vector normalization and angular separation.

• Rotation matrices.

• Axis-angle quaternions.

• Quaternion multiplication and rotation.

• Rigid transforms, inversion, and composition.

• Cartesian and spherical transformations.

• Galactic and equatorial transformations.

• Sky-tangent projection.

• Parallax and proper motions.

• Zenith and parallactic geometry.

• Explicit singular-domain derivative behavior.

Spatial algorithms¶

• Morton encoding and decoding.

• Deterministic hashing.

• Particle-to-bin and linear-cell assignment.

• Capacity-limited and exact bin filling.

• Neighbor-stencil candidate gathering.

• Approximate nearest-neighbor candidate generation.

• Exact fixed-radius pair acceptance.

• Explicit overflow, capacity, recall, symmetry, and cutoff contracts.

Units and quantities¶

Two related surfaces currently coexist:

• jaxstro.units: mature named unit systems, CGS default, astrophysical scale systems, photometric units, and explicit gravitational constants.

• jaxstro.quantity: dimensional algebra, units, conversion, parsing, registries, serialization, unit bases, constants, dimension-aware math, and opt-in spectral, temperature–energy, and mass–energy equivalencies.

The quantity layer is implemented and extensively tested, but ecosystem-wide adoption remains an evidence-gated decision.

Spectra and atmospheres¶

• Typed spectral coordinates, sampling semantics, flux semantics, provenance, and statuses.

• Frequency and wavelength conversion.

• $F_\lambda/F_\nu$ transformation.

• Surface flux to luminosity and observer flux.

• Point and conservative-bin resampling.

• Prepared rectilinear and simplex interpolation stencils.

• Exact-product atmosphere adapters for NewEra, BOSZ, Sonora, and TLUSTY.

• Catalog discovery, topology selection, acquisition planning, artifact coverage, and fail-closed policy selection.

• Host-side preparation followed by fixed-shape JAX evaluation.

This is unusually mature infrastructure for sharing atmosphere and spectral mechanics without absorbing Fluxax’s photometric interpretation.

Parameters and constrained spaces¶

• Selective Equinox PyTree-to-vector parameterization.

• Static selection of free and fixed leaves.

• Identity, exponential, softplus, and bounded sigmoid bijectors.

• Analytic log-Jacobian terms.

• Reconstruction of structured models from unconstrained vectors.

This is a parameter-coordinate bridge, not an inference engine.

Scientific contracts, evidence, and provenance¶

• Module and callable contracts.

• Maturity, support, execution-boundary, transform, AD, failure, and evidence classifications.

• Runtime export auditing.

• Typed computational-evidence artifacts.

• Metric and comparison records with units.

• Environment records and content digests.

• Deterministic JSON and Markdown.

• Artifact hashing and method manifests.

• Source-backed provenance cards.

• Freshness and assertion-bearing evidence gates.

This predict → compute → audit → state the warranted claim architecture is Jaxstro’s most distinctive advantage over a miscellaneous utility library.

What Jaxstro should become¶

The target identity is:

Jaxstro is the reusable JAX-native substrate for scientifically trustworthy numerical computation and scientific machine learning: mathematical primitives, transforms, execution telemetry, dimensional semantics, and evidence—not domain policy or inference conclusions.

That is broader than “astrophysics helpers,” but still disciplined.

The ownership split should remain:

Informax can consume Jaxstro’s linear operators, differentiable solvers, training mechanics, parameter transforms, and provenance without surrendering ownership of what inference or experimental design means.

Recommended additions¶

Unchecked items are proposals. Completing a checkbox requires the public API, tests, theory, validation, contract-registry entry, and relevant evidence—not only an implementation commit.

Priority 1: a narrow jaxstro.ml scientific-ML substrate¶

This is the most direct missing layer if Jaxstro is to become the main ML foundation.

• Add immutable Standardizer, RobustScaler, and WhiteningTransform PyTrees.

• Separate host fitting from JAX application with masks and named feature axes.

• Add deterministic train, validation, and test splits plus keyed fixed-shape batching.

• Define optimizer-agnostic TrainState, TrainingTrace, and fit_fixed_steps contracts.

• Add a lax.scan compiled training path with explicit executed masks and checkpoints.

• Add masked and weighted regression losses.

• Add scale-aware and relative-error losses.

• Add Gaussian NLL and heteroscedastic-regression mechanics.

• Add gradient clipping and finite-update rejection with telemetry.

• Record dataset, normalization, seed, model-structure, and checkpoint manifests.

• Add reproducibility and data-leakage audits.

• Add Jacobian, divergence, curl, Laplacian, and residual construction for scientific ML.

• Evaluate optional neural-tangent and Gauss–Newton operator diagnostics.

It should not include a model zoo. Equinox already owns neural modules; Optax owns optimizers. Jaxstro should own the repeatable, audited scientific execution around them.

Priority 2: matrix-free iterative linear solvers¶

• Add conjugate gradient for symmetric positive-definite systems.

• Add MINRES for symmetric indefinite systems after a real consumer exists.

• Add GMRES for general systems after a real consumer exists.

• Define reusable preconditioner protocols.

• Return fixed-shape iteration traces.

• Record residual, conditioning, breakdown, and stagnation statuses.

• Gate derivatives through jax.lax.custom_linear_solve only where the mathematical and numerical assumptions are validated.

• Add dense and analytic parity cases plus independent finite-difference gates.

This would unlock larger inverse problems, implicit layers, stellar-structure systems, PDE discretizations, and scalable Informax geometry without requiring every project to invent solver telemetry.

Priority 3: vector nonlinear systems and fixed points¶

Build this only around concrete consumers, but it is a natural extension of the scalar-root work.

• Identify and validate at least one downstream vector-root or fixed-point consumer before generalizing the API.

• Add Newton–Krylov or damped Newton for $F(\mathbf{x})=0$.

• Evaluate safeguarded fixed-point acceleration such as Anderson acceleration.

• Add trust-region or line-search globalization.

• Return typed primal and linear-solve certificates.

• Add strictly gated implicit derivatives.

• Surface uniqueness, conditioning, residual, and branch assumptions as evidence.

Potential science includes hydrostatic stellar structure, equilibrium chemistry, coupled feedback balances, differentiable steady states, and constrained orbital solutions.

Priority 4: quasi-Monte Carlo and variance reduction¶

• Add scrambled Sobol sequences.

• Evaluate Halton sequences only with their limitations documented.

• Add Latin-hypercube sampling.

• Add antithetic sampling.

• Add control-variate primitives.

• Add effective sample size and weight-degeneracy diagnostics.

• Add replicated-scramble uncertainty estimates.

• Record seed and sequence provenance.

This would benefit Progenax population integration, Fluxax synthetic surveys, simulation ensembles, and Informax nested expectation calculations.

Priority 5: adaptive quadrature¶

• Add Gauss–Kronrod pairs.

• Add error-estimated adaptive interval subdivision.

• Define fixed-capacity or explicitly staged execution suitable for JAX.

• Return typed exhaustion and nonfinite-integrand outcomes.

• Separate value and derivative contracts.

• Support vector-valued integrands and batched domains.

Avoid building a competing adaptive ODE ecosystem. Diffrax is the better owner for general event-driven ODE and SDE solving.

Priority 6: signal and time-series primitives¶

• Define FFT normalization conventions.

• Add convolution and correlation with explicit boundary policy.

• Add window functions and equivalent-noise bandwidth.

• Add power and cross-spectral density.

• Add phase and time-delay estimation.

• Evaluate irregular-sampling primitives only after a concrete science use case exists.

• Add unit-aware frequency and time-axis handling.

These would support variability, spectra, detector response, synthetic observations, and time-domain astronomy.

Priority 7: generic uncertainty propagation¶

Keep inferential interpretation in Informax, but Jaxstro can provide mechanics for:

• First-order delta-method propagation.

• Covariance pushforward through Jacobians.

• Unscented and sigma-point transforms.

• Ensemble propagation with deterministic keying.

• Sensitivity decomposition.

• Conditioning and rank diagnostics.

Jaxstro would return mathematical propagation results; Informax would decide whether they justify an inferential claim.

Priority 8: multidimensional field methods¶

Only add these once Stellax, Nebulax, or another consumer supplies concrete contracts.

• Identify at least two consumers and ratify their shared topology and conservation boundary.

• Add structured two- and three-dimensional mesh geometry.

• Add gradient, divergence, curl, and Laplacian stencils.

• Add conservative flux divergence.

• Define explicit boundary-condition objects.

• Add prolongation and restriction.

• Add matrix-free elliptic operators.

Do not grow into a full CFD framework without multiple consumers and strong conservation evidence.

What not to add¶

Avoid turning Jaxstro into a collection of inferior replacements:

• No homegrown neural-network framework; use Equinox.

• No full optimizer zoo; use Optax or specialized optimization packages.

• No general MCMC, VI, NPE, SBI, or posterior orchestration; that belongs in Informax.

• No Fisher-OED criteria, calibration gates, claim/refusal policy, or learned experimental-design proposals; those belong in Informax.

• No full adaptive ODE/SDE stack; use Diffrax.

• No generic duplicate of every NumPyro or Distrax distribution.

• No domain-specific timestep, particle, photometry, stellar-evolution, or atmosphere-acceptance policy.

• No algorithm merely because SciPy has it. Add methods when at least two projects need the same mathematical contract or when one foundational need is exceptionally general.

Build checklist¶

This is the recommended order. The capability-level checklists above define the work inside each rung.

• Generate the callable-level JAX transform and maturity matrix.

• Decide the ecosystem future of jaxstro.quantity from downstream parity, performance, serialization, ergonomics, and migration evidence.

• Add the minimal jaxstro.ml vertical slice: normalization, deterministic batching, fixed-step training, typed trace, and checkpoint provenance.

• Add conjugate gradient plus the linear-solver protocol and evidence model.

• Validate one real vector-root or fixed-point consumer before generalizing the solver.

• Add scrambled Sobol sampling and replicated uncertainty evidence.

• Add Gauss–Kronrod adaptive quadrature.

• Add signal-processing primitives driven by Fluxax or time-domain use.

• Expand meshes only when multiple field-based packages can define the shared boundary.

• Generate downstream symbol-to-project compatibility records throughout.

That sequence makes Jaxstro genuinely useful as the daily numerical and ML foundation without sacrificing the thin-foundation discipline that currently makes it trustworthy.

Definition of done for a roadmap item¶

A checkbox moves to complete only when the relevant layer includes:

1. a mathematical or semantic contract;

2. a public API with typed boundary and failure behavior;

3. supported JAX transforms and cost semantics;

4. independent validation rather than self-consistency alone;

5. reproducible artifacts where metrics matter;

6. accessible teaching that states what the evidence does not prove;

7. downstream adoption evidence when reuse justifies ownership; and

8. an updated contract registry, validation index, scorecard, and roadmap.

The highest-impact direction is not adding the most algorithms. It is making the existing breadth uniformly inspectable from equation to execution to claim, while admitting new primitives only when reuse and evidence justify them.

Evidence and companion pages¶

• API reference — current public modules and execution boundaries.

• Scientific contract registry — generated ownership, maturity, transform, AD, failure, and evidence contracts.

• Validation — independent validation anchors.

• Jaxstro package assessment scorecard — living grades and promotion gates.

• Package and documentation SOTA assessment — ranked infrastructure and pedagogy investments.

• Numerical methods roadmap — completed first-generation numerical expansion checklist.

• Astro-first, science-general vision — science-general ecosystem boundary.