Architecture

This section explains why the software has its current shape: JAX-native functional and PyTree design, explicit unit boundaries, one-way package dependencies, evidence ownership, and a deliberately thin foundation. Where the theory section covers mathematical method contracts, this section covers package structure and responsibility.

Foundation boundary¶

jaxstro is the ecosystem foundation. Its directly importable modules are units, quantity, constants, astrometry, coords, geometry, numerics, spatial, params, atmospheres, provenance, testing, and jaxconfig. It does not own simulations, survey rendering, stellar evolution, or inference workflows. Domain packages depend on jaxstro; jaxstro never imports package-specific code back from them. That one-way rule keeps foundation changes auditable before they reach Gravax, Progenax, Fluxax, Startrax, or later packages.

The keystone decision is ADR 0001 — Thin foundation posture. The dependency and packaging decisions are recorded in ADR 0002 — Adopt equinox as core dependency, ADR 0003 — Standalone uv project with hatchling, and ADR 0010 — Ecosystem config architecture.

Astro-first, science-general¶

The package should be marketable as evidence-first JAX infrastructure for differentiable science. Astronomy supplies the pressure tests: physical units, awkward coordinate transforms, tabulated spectra, stiff numerical ranges, and gradients that must be trusted by downstream inference. The abstractions that survive those tests are useful well beyond astronomy.

New foundation modules should pass four checks before they belong here:

1. Generic across domains. The primitive is useful without knowing about stars, galaxies, filters, surveys, or a specific simulator.

2. JAX-native by construction. The public runtime path composes with jit, vmap, and grad where differentiation is part of the contract.

3. Explicit about boundaries. Units, valid domains, clamping, saturation, static arguments, and non-differentiable preprocessing are named rather than hidden.

4. Backed by evidence. The module has focused unit tests, transform tests, finite-difference or analytic validation where relevant, and documentation that explains failure modes.

The broader product vision and future-module map are in Astro-first, science-general vision.

Quantity architecture¶

jaxstro.units remains the current canonical ecosystem contract. The jaxstro.quantity layer is implemented: it provides concrete units, dimension-safe arithmetic, exact parser/serialization, role-aware bases, versioned constants, and explicit equivalencies. However, ecosystem adoption and any replacement cutover remain deferred. The architecture is documented in Quantity system architecture, but that page is an evaluation surface rather than a migration instruction. It extends the decision recorded in ADR 0006 — Build own Quantity, not unxt/quax.

Numerical shape¶

Array kernels are JAX-native and use jit, vmap, or grad only where each public method’s contract supports that transform. Fixed-shape numerical iterations and sanitized branches protect traced execution, while tests check forward values and gradients on the smooth domains named by each method. This does not make discrete indexing, sorting, catalog lookup, or every branch differentiable. The detailed mathematical contracts live in Writing AD-safe scientific numerics, and quantitative checks live in Validation.

Units policy¶

jaxstro defaults to CGS through DEFAULT_UNITS, because it is the domain-agnostic base layer. Domain packages choose their own package-level defaults. Core APIs either accept explicit units or explicit physical constants; convenience wrappers may resolve units=None to the package default. ADR 0007 — CGS as default unit system records this policy.

Provenance ownership¶

Two provenance surfaces answer different questions. The source-backed provenance cards connect constants, transforms, and capability boundaries to references, code symbols, validation paths, and evidence states. They describe why a public scientific claim is trustworthy.

The runtime artifact manifests instead record what a particular computation consumed: configuration, environment, seeds, hashes, and method identifiers. They make a run reproducible. A card does not replace a manifest, and a manifest does not establish the scientific source behind a constant or transform.

Differentiable and discrete boundaries¶

spatial indexing and candidate construction, together with atmosphere catalog selection, are host-side, discrete preprocessing. Their selected arrays can feed JAX-native kernels, but the selection decisions themselves are not advertised as differentiable. This division keeps runtime shapes and ownership explicit rather than hiding Python decisions inside traced code.

Data-layer boundary¶

Large third-party scientific data is not vendored into the package. Foundation data adapters may expose local discovery, provenance, processed-artifact validation, and catalog-first runtime selection, but raw products remain in user cache locations or explicitly gitignored local mirrors. The first large example is jaxstro.atmospheres: it can process and index local NewEra, BOSZ, Sonora, and TLUSTY atmosphere spectra, while filters, photometry, bolometric corrections, survey rendering, and physical interpretation remain downstream. Atmosphere support remains in progress: the capability map distinguishes staged data, processed artifacts, and implemented runtime backends rather than treating every catalog entry as available interpolation support.

Test layers¶

The suite is organized by risk:

Every release-facing claim should point either to the test that enforces it or to a decision record that explains why the boundary exists.

Spectra data architecture¶

The local atmosphere capability map is in Atmosphere capabilities. It explains which libraries are processed, which have runtime backends, and why TLUSTY uses ragged frequency-grid subgroups.

The runtime boundary is documented in Spectra data architecture. It defines the host-side AtmosphereQuery -> PreparationResult path, the generic jaxstro.spectra owner, and fixed-topology JAX evaluation. It also explains why surface versus observer flux, F_lambda versus F_nu, and point versus bin sampling must remain explicit. Filters, photometry, instruments, images, and survey semantics remain downstream in Fluxax.