Anna Rosen, Ph.D.Computational Astrophysicist

Software

Differentiable astrophysics, built as instruments

What can observations genuinely tell us about stellar populations?

I build JAX-native scientific software in which the path from a cluster's birth to its survey observables stays differentiable — across adaptive timesteps and stellar phase transitions, exactly where differentiating naively stops being correct. The pipeline returns not just a prediction, but how that prediction responds to each physical parameter.

Jaxstro is the spine of that program: a family of packages that each own one scientific stage and remain independently useful. It powers current work on star-cluster identifiability, and it is built to outlast any single project — a foundation for questions I haven't asked yet.

The source for each package opens as its paper lands. Until then the pages below show what is working and how far along it is.

Six panels in two columns. Top row: a simulated g-band image of a crowded cluster, first as an expected source-limited scene and then as a nominal 30-second Rubin-like exposure with sky and noise. Middle row: maps of how much information each pixel carries about the IMF slope. Bottom row: the same for the binary mass-ratio slope. The information maps are concentrated on a sparse subset of stars, and differ between the two parameters.
PreliminaryThis is what differentiable codes buy you. Fisher information per pixel — IMF slope (middle), mass-ratio slope (bottom).fluxax

The Jaxstro Ecosystem

Jaxstro turns stellar-population simulations from forward calculators into models we can differentiate, fit, interrogate, and use to design observations.

  1. Birth populationsprogenax
  2. Stellar evolutionstartraxCollisional dynamicsgravax
  3. Observablesfluxax
  4. Inference & designinformax
The Jaxstro ecosystem: a differentiable model from a cluster's birth to its survey observables, returning the gradients that reveal what the data can recover. Stellar evolution and collisional dynamics are one coupled step by design — a star loses mass while it is being moved, and an encounter changes the path it evolves along, so neither finishes before the other begins. Every stage composes on the shared jaxstro foundation. Select any stage for detail.

Inference is where the chain pays out: informax works out what an observation can recover before it is taken — and says so even when the answer is that it cannot.

The packages

  1. 00Foundation

    jaxstro

    The shared numerical substrate.

    Units and constants, coordinate transforms, numerical methods, derivative contracts, and provenance — so every package composes without silent inconsistencies.

    Mature
  2. 01Birth populations

    progenax

    Truth-known cluster birth conditions.

    Differentiable initial conditions for star clusters: true-equilibrium King, EFF and LIMEPY structure with anisotropy and primordial mass segregation, mass-dependent multiplicity, and IMFs you can take gradients through.

    MatureMethods paperin preparation
  3. 02Stellar evolution

    startrax

    From birth mass to stellar state and remnants.

    A differentiable map from initial mass, age and metallicity to full single-star state — winds, mass loss, lifetimes, remnants — with gradients that survive the phase transitions, where differentiating naively does not.

    In developmentMethods paperplanned
  4. 03Collisional dynamics

    gravax

    Differentiable collisional dynamics.

    A JAX-native, hierarchical collisional N-body engine: high-order exact direct gravity with explicit ownership of binaries and compact subsystems, and gradients that survive the adaptive timestep — so a cluster model can be fitted to data, not only run forward.

    In developmentMethods paperplanned
  5. 04Observables

    fluxax

    Physical states → survey observables.

    A differentiable Gaia and Rubin/LSST catalogue forward model — photometry, astrometry, selection, crowding, blending and extinction inside one gradient-safe likelihood — with a differentiable pixel renderer alongside it.

    Far alongMethods paperplanned
  6. 05Inference & design

    informax

    What the data can actually recover.

    Exact Fisher-information geometry and local optimal experimental design, with simulation-based calibration, coverage and TARP gates — and refusal contracts that make a result decline itself when the data cannot support it.

    Far alongMethods paperplanned

No package is publicly released yet; methods papers come first. The readiness meter is the honest signal — everything here is under active development, and they differ in how far along they are. Each package has its own page with the demos that show it working, and what its methods paper will establish.