Skip to content

About

physics-simulation, jax, llm

Resources

Stars

0 stars

Watchers

0 watching

Forks

Repository files navigation

Space Sandbox

Engine-agnostic space simulator built around one contract: a scenario JSON document (schemas/scenario.schema.json). See CLAUDE.md for the project thesis and architecture principles.

Project philosophy

Everything in this repo produces or consumes one JSON document, validated against schemas/scenario.schema.json before it runs; nothing hard-codes a physics implementation — every backend (NumPy reference, JAX differentiable) implements the same Backend interface and is chosen by a string in the scenario. Generated numerical code — integrators, force models, LLM-authored JSON, gradients — is treated as a hypothesis, not an answer: it ships with an empirical check (energy conservation, a circular-orbit sanity test, a finite-difference gradient comparison) rather than a claim of correctness.

Status

Phases 0 through 3 are done. The full loop runs end to end: describe a scenario in English, simulate it, solve for the parts you could not specify, and get an explanation of the result.

  • Phase 0 — repo scaffold, validating scenario loader, Backend interface.
  • Phase 1a — reference backend: NumPy velocity-Verlet, validated by energy-conservation and circular-orbit tests.
  • Phase 1b — LLM scenario authoring (English -> schema-valid JSON) via the Anthropic API.
  • Phase 1c — local LLM fallback for authoring: author_scenario now talks to either the Anthropic API or a local Ollama server through a shared provider interface, with the same retry/validation loop for both.
  • Phase 2 — jax backend: the same physics written so jax.grad can differentiate through it, plus inverse design (solve for the initial conditions a scenario leaves free). Every gradient is validated against finite differences before any optimizer is allowed to use it.
  • Phase 3a — the LLM authors optimization blocks, not just bodies, and --solve runs them. A semantic gate rejects blocks that are schema-valid but unrunnable before they reach the solver, since the schema's enums are deliberately larger than what the code implements.
  • Phase 3b — grounded narration: --narrate explains a run in plain English, using numbers computed in Python and a small hand-written knowledge base, never the model's own recollection.

Results

Concrete numbers from the current test suite / example runs, not estimates:

  • Energy conservation (reference backend, examples/two_body.json, 5000 steps / 5 orbital periods): relative drift 1.676715e-13. Circular-orbit radius stability over 2 periods: max relative deviation 2.467645e-06.
  • Gradient check (examples/solve_circular_orbit.json, jax backend): finite-difference vs. autodiff agree to 8.194e-10 relative (eps=1e-6), and the disagreement shrinks as O(eps^2) across three epsilons (8.4e-05 → 8.4e-09 → 8.2e-10) — the signature of a correct gradient, not a coincidence.
  • Solved vs. analytic: gradient descent (Adam) recovers a circular-orbit velocity of 4.443267283233834, against the analytic sqrt(G*M/r) = 4.443269032129325 — relative error 3.936056e-07, matching the integrator's own O(dt²) truncation floor (confirmed separately by halving dt and watching the error quarter).

Reproduce with pytest -s tests/test_reference_backend.py and python -m spacesandbox.solve examples/solve_circular_orbit.json.

Setup

python -m venv .venv
.venv\Scripts\activate        # Windows
pip install -e ".[dev]"

Usage

python -m spacesandbox.loader examples\two_body.json          # validate + inspect
python -m spacesandbox.run    examples\circular_orbit.json    # simulate, report energy drift
python -m spacesandbox.solve  examples\solve_circular_orbit.json   # gradient check, then solve
pytest

spacesandbox.solve prints the finite-difference vs. autodiff comparison first and refuses to optimize unless the two agree, then compares the solved answer to the analytically-known one. Both numbers are printed in full; neither is reduced to a pass/fail.

Scenario authoring

python -m spacesandbox.author "a sun and an earth-like planet" --provider anthropic
python -m spacesandbox.author "a sun and an earth-like planet" --provider local

--provider anthropic (the default) needs ANTHROPIC_API_KEY set and calls the Claude API. --provider local talks to a local Ollama server (default http://localhost:11434, model qwen3:8b) instead — free and offline, but noticeably less reliable at producing valid schema JSON on the first attempt, so --max-attempts matters more here than on the API path. Each attempt's pass/fail is logged to stderr so local-vs-API reliability can be compared from real runs instead of guessed at. If Ollama isn't running, the CLI fails with a clear "is ollama serve running?" error rather than silently falling back to the API.

Authoring a scenario that solves for something

python -m spacesandbox.author "a sun and an earth-like planet" --run
python -m spacesandbox.author "a probe at 1.5 AU -- find the velocity that makes its orbit circular" --solve

A prompt that asks to find or solve for a value gets an optimization block as well as bodies, and --solve feeds it to the same gradient-checked solver a hand-written scenario goes through — there is no laxer path for generated scenarios. --run simulates instead, for scenarios with no block to solve; the two flags are mutually exclusive, and passing the one that does not match the generated scenario is an error naming the one you wanted rather than a guess.

The schema's quantity enum is deliberately larger than what the engine implements (only eccentricity today), so an authored block is checked against the implemented set before it runs: an unimplemented quantity or optimizer method, a body id that does not exist, a missing target_body or target_value, an integrator the differentiable backend cannot use, or a reference_body equal to its target_body all fail at authoring time with a message the model can act on, and drive a retry. Each corresponds to a real crash — or, in the last case, a silent NaN — inside the solver.

At solve time the validated gradient catches one more kind of nonsense: a free parameter whose derivative is exactly zero cannot influence the objective, so the optimizer could never move it. A test particle's own mass is the clean example — it passes every authoring-time check, and the gradient check itself still reports PASSED, because finite differences and autodiff agree the derivative is zero.

Read the generated objective before trusting the answer, because the model may quietly narrow an unsupported goal into the one supported quantity (e.g. "capture into a closed orbit", meaning e < 1, becomes eccentricity → 0, a strictly circular orbit), producing a numerically correct answer to a narrower question than you asked. No check catches that one; it needs a human.

Explaining a result

python -m spacesandbox.run    examples\two_body.json --narrate
python -m spacesandbox.solve  examples\solve_circular_orbit.json --narrate
python -m spacesandbox.author "a probe at 1.5 AU -- find its circular velocity" --solve --narrate

--narrate prints a plain-English explanation of what the run did. Python computes every number and every physical verdict (bound vs. unbound, converged vs. budget-spent); the LLM is handed those facts plus the one to three most relevant notes from knowledge/orbital_mechanics_notes.md and told to use nothing else. That knowledge file is hand-written ground truth — a wrong entry there does not produce an obviously wrong explanation, it produces a fluent and confident one, so edit it carefully. A narration that fails never fails the run it describes: you get a warning and your numbers.

On spacesandbox.author, --narrate accompanies --run or --solve and reuses whichever provider authored the scenario, so an offline --provider local run stays entirely offline instead of quietly reaching for the API.

About

physics-simulation, jax, llm

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages