elodin-simulation

Elodin Simulation

Elodin is a JAX-based simulation platform for aerospace and physical systems. Simulations are Python scripts that define a World, spawn entities with components, compose systems, and run.

Installation

pip install -U elodin          # Released SDK
elodin editor sim.py           # Run with 3D visualization
elodin run sim.py              # Headless execution
python sim.py bench --profile  # Performance profiling

Simulation Structure

Every simulation follows this pattern:

import elodin as el
import jax.numpy as jnp

# 1. Create world
w = el.World()

# 2. Spawn entities with archetypes
w.spawn(el.Body(
    world_pos=el.SpatialTransform(linear=jnp.array([0.0, 0.0, 10.0])),
    inertia=el.SpatialInertia(mass=1.0),
), name="ball")

# 3. Define systems
@el.map
def gravity(f: el.Force, inertia: el.Inertia) -> el.Force:
    return f + el.SpatialForce(linear=inertia.mass() * jnp.array([0.0, 0.0, -9.81]))

# 4. Compose and run
sys = el.six_dof(sys=gravity, integrator=el.Integrator.Rk4)
w.run(sys, simulation_rate=120.0)

Core Concepts

Components

Data containers defined with typing.Annotated + el.Component:

import typing as ty

Wind = ty.Annotated[
    jax.Array,
    el.Component("wind", el.ComponentType(el.PrimitiveType.F64, (3,)),
                 metadata={"element_names": "x,y,z"}),
]

Built-in spatial types (WorldPos, WorldVel, Force, Inertia, WorldAccel) already carry component metadata — no ComponentType needed.

Archetypes

Group components into spawnable bundles:

@el.dataclass
class Satellite(el.Archetype):
    world_pos: el.WorldPos
    world_vel: el.WorldVel
    inertia: el.Inertia
    reaction_wheels: ReactionWheelCmd

el.Body is the built-in archetype providing WorldPos, WorldVel, Inertia, Force, WorldAccel.

Systems

Three decorator levels — choose the simplest that fits:

@el.map
def drag(vel: el.WorldVel) -> el.Force:
    return el.SpatialForce(linear=-0.01 * vel.linear())

@el.system
def gravity(graph: el.GraphQuery[GravityEdge],
            q: el.Query[el.WorldPos, el.Inertia]) -> el.Query[el.Force]:
    return graph.edge_fold(q, q, el.Force, el.SpatialForce(), compute_gravity)

System Composition

Chain systems with the pipe operator — order matters:

sys = sensors | kalman_filter | control | el.six_dof(sys=effectors)

6DOF Physics

el.six_dof() integrates forces/torques into position and velocity:

el.six_dof(
    sys=effectors,                    # Systems computing el.Force
    integrator=el.Integrator.Rk4,     # or Integrator.SemiImplicit
    time_step=1/300.0,                # Optional: override simulation step
)

Inertia is body-frame; all other quantities are world-frame.

Graph Queries

Model relationships (gravity, constraints, springs) between entities:

GravityEdge = ty.Annotated[el.Edge, el.Component("gravity_edge")]

w.spawn(el.Archetype(edge=GravityEdge(el.Edge(body_a_id, body_b_id))), name="a_to_b")

@el.system
def gravity(graph: el.GraphQuery[GravityEdge],
            q: el.Query[el.WorldPos, el.Inertia]) -> el.Query[el.Force]:
    return graph.edge_fold(
        left_query=q, right_query=q,
        return_type=el.Force, init_value=el.SpatialForce(),
        fold_fn=lambda acc, pos_a, m_a, pos_b, m_b: acc + compute(pos_a, m_a, pos_b, m_b),
    )

Spatial Vector Algebra

Elodin uses Featherstone spatial vectors. Key types:

Quaternion operations: Quaternion.from_axis_angle(), q @ vector (rotate), q.inverse(), q.integrate_body(omega).

Visualization

Asset root (assets/)

Simulations load meshes, schematics, skyboxes, themes, and terrains from a single asset root. Paths in KDL (glb path=…, icon path=…) are relative to that root.

Where it lives (resolution order):

1. $ELODIN_ASSETS — explicit override (absolute or cwd-relative)
2. <sim_entry_dir>/assets — next to main.py / the sim entrypoint
3. <cwd>/assets
4. Nearest ancestor of the sim entry that contains an assets/ directory

Repo examples usually run from the Elodin checkout root and use the shared tree at assets/ (many .glb files live at the root of that tree). For a standalone sim, put an assets/ folder beside main.py, or set ELODIN_ASSETS.

When the sim records with db_path=… or ELODIN_DB_PATH, that tree is ingested once into {db}/assets/ and KDL local paths become db:… for portable replay. See the elodin-db skill and DB Asset Server.

Native folders under the asset root

Only create the subtrees you need. Keys preserve subdirectories; there is no hard schema beyond these conventions:

assets/
  # GLB meshes — flat at root (repo default) or under meshes/ / models/
  f22.glb
  meshes/rocket.glb          # also fine; reference as path="meshes/rocket.glb"
  models/jet.glb             # docs/tests often use models/

  schematics/                # KDL panels / window sub-schematics
    main.kdl                 # conventional active schematic key after ingest
    telemetry.kdl            # window path="telemetry.kdl" → stored under schematics/

  skyboxes/
    manifest.ron             # required for named skybox name="…"
    desert_night.cubemap.ktx2

  color_schemes/             # optional UI themes (not copied into DB by name alone)
    default_dark.json
    default_light.json

  terrains/
    planar/<region>/         # world_mesh region="<region>"
      region.toml            # optional; built-in presets exist
      config.tc              # preprocess output
      data/height/<tile>.bin
      data/albedo/<tile>.bin
    spherical/               # globe-style terrain

Agent tips:

• Prefer paths relative to the asset root (path="edu-450-v2-drone.glb"), not filesystem paths outside it.
• Do not invent reserved keys: .elodin-ingested and __index__/ are DB-internal.
• Procedural meshes (sphere, box, …) and icon builtin=… need no asset files.
• Custom color-scheme JSON is not persisted into the DB today — only the scheme name travels in KDL.
• If the editor cannot find a mesh, check cwd / ELODIN_ASSETS before changing KDL.

KDL Schematics

Define 3D objects and camera views in KDL files or inline:

object_3d ball.world_pos {
    sphere radius=0.2 { color 25 50 255 }
}
object_3d aircraft.world_pos {
    glb path="f22.glb" scale=0.01 rotate="(0, 90, 0)"
}
viewport name=Chase pos="drone.world_pos.translate(-5, -5, 3)" look_at="drone.world_pos"

Panel Layout (Python API)

cam = el.Panel.viewport(track_entity=sat_id, fov=45.0, hdr=True, name="3D")
graph = el.Panel.graph(el.GraphEntity(sat_id, *el.Component.index(el.WorldPos)[:4]), name="Position")
w.spawn(el.Panel.vsplit(cam, graph), name="main_view")

SITL/HITL Integration

The reference SITL example is examples/betaflight-sitl/ (Betaflight flight controller in lockstep over UDP, recorded to a portable DB whose assets/ tree makes it replayable anywhere).

Use pre_step/post_step callbacks with StepContext for lockstep synchronization:

def post_step(tick: int, ctx: el.StepContext):
    data = ctx.component_batch_operation(reads=["drone.accel", "drone.gyro"])
    motors = flight_controller.step(accel=data["drone.accel"], gyro=data["drone.gyro"])
    ctx.write_component("drone.motor_command", motors)

w.run(sys, simulation_rate=1000.0, post_step=post_step, db_path="sitl_data")

Mark components as externally controlled to prevent simulation overwrite:

ThrustCmd = ty.Annotated[jax.Array,
    el.Component("thrust_cmd", el.ComponentType.F64, metadata={"external_control": "true"})]

Shared Constants And Monte Carlo

The Cranelift backend transparently interns baked StableHLO constants over 1 MB. For large immutable lookup tables, aero maps, ephemerides, terrain grids, or other constants, capture the JAX array in the system closure:

aero_table = jnp.asarray(aero_grid)

@el.map
def aero_force(v: el.WorldVel, f: el.Force) -> el.Force:
    coeff = aero_table[0, 0]
    return f + el.SpatialForce(linear=-coeff * v.linear())

During StableHLO parsing, large dense<"0x..."> blobs are moved into the content-addressed mmap cache. Multiple processes compiling the same constant map the same cache file, so Monte Carlo campaigns avoid paying N copies of the table in resident memory.

For campaign structure and memory reporting, use the native campaign runner:

elodin monte-carlo run examples/monte-carlo/main.py \
  --campaign examples/monte-carlo/campaign.toml \
  --spec examples/monte-carlo/spec.toml \
  --out dbs/monte-carlo-demo

Declare tunable parameters with el.monte_carlo.params_spec(...), read the current row with el.monte_carlo.params(...), and optionally emit scalar outputs with el.monte_carlo.result(...). Each run writes a separate DB path. The runner pins ELODIN_CACHE_DIR for all workers and auto-sizes workers/runtime threads from available CPUs when unset. For SITL campaigns, register external controllers with world.recipe(...); the campaign runner injects worker-slot ports into every process via ELODIN_MONTE_CARLO_* environment variables. Campaign startup reaps pre-existing elodin and elodin-db processes by default to avoid stale editor/database sessions colliding with worker ports; use --keep-existing only when intentionally managing those processes yourself. Per-run stdout/stderr lands in runs/<run_id>/logs/, and the runner injects ELODIN_SIM_SUMMARY_JSON so each run writes a structured timing snapshot. At campaign end, the native runner prints and writes campaign_summary.txt with an aggregated version of the standard elodin simulation summary block plus CPU/RAM/disk rollups. Use --memory-probe only for shared-constant PSS proof runs; it enables expensive /proc/<pid>/smaps sampling and writes memory.json/processes.csv.

Execution Modes

Backend selection: The backend parameter defaults to "cranelift" — a pure-Rust StableHLO JIT that runs the entire tick as a single native function call, with no Python in the hot loop. Use "jax-gpu" for high-parallelism workloads that benefit from GPU execution. For tiny worlds, the CPU cranelift backend is usually fastest because kernel launch and device-transfer overhead dominates compute.

Use examples/n-body/main.py as the canonical benchmark. It runs the supported backends (cranelift, jax-cpu, jax-gpu) side-by-side:

nix develop --command ELODIN_BACKEND=jax-gpu elodin run examples/n-body/main.py

To compare backends, run the same command with ELODIN_BACKEND set to each of: cranelift, jax-cpu, jax-gpu.

Earth Gravity Models

from elodin.j2 import J2              # Simple oblate Earth
from elodin.egm08 import EGM08        # High-fidelity spherical harmonics

model = EGM08(max_degree=64)          # <2.5ms at degree <=250
force = model.compute_field(x, y, z, mass)

Physics Regression Testing

When changes to the simulation pipeline (Noxpr graph, cranelift-mlir compilation, shape handling, etc.) might alter numeric output, use a database-export diff to detect regressions. The process:

1. Capture a baseline on main

git stash && git checkout main
nix develop
just install

# Run the sim, writing to a dedicated DB path
BALL_DB_PATH=dbs/ball-main uv run examples/ball/main.py bench --ticks 2000

# Export to flat CSVs (--flatten splits vector columns)
elodin-db export --format csv --flatten --output exports/ball-main dbs/ball-main

The BALL_DB_PATH env var is read by examples/ball/main.py and passed to world().run(..., db_path=...). Other examples can be wired the same way.

2. Capture the branch under test

git checkout <branch> && git stash pop
nix develop
just install
BALL_DB_PATH=dbs/ball-branch uv run examples/ball/main.py bench --ticks 2000
elodin-db export --format csv --flatten --output exports/ball-branch dbs/ball-branch

3. Diff component-by-component (ignoring timestamps)

# Quick pass/fail for every physics component:
for f in ball.world_pos.csv ball.world_vel.csv ball.force.csv ball.wind.csv ball.world_accel.csv; do
    echo -n "$f: "
    diff <(cut -d',' -f2- exports/ball-main/$f) \
         <(cut -d',' -f2- exports/ball-branch/$f) | wc -l
done

Zero diff lines = bit-for-bit identical physics. Non-zero tells you which component diverged. Inspect the first differing row to find the tick where divergence starts and whether it is a large discrete jump (logic bug) or gradual drift (floating-point).

4. Interpret results

Tips

• Use --ticks 2000 (not 1200) to expose drifts that accumulate.
• Run all three canonical benchmarks to cover every code path:
  • ball -- single-entity (batch1 path), uses el.Seed (U64) + random.key
  • drone -- multi-entity, uses sensor_tick (U64)
  • cube-sat -- JAX backend (backend="jax-cpu"), covers the non-IREE path
• If you need to dump the StableHLO MLIR for comparison, set ELODIN_IREE_DUMP_DIR=/tmp/debug before running.
• Clean up temp databases after: rm -rf dbs/ball-main dbs/ball-branch exports/.

Additional Resources

• For the full Python API reference, see api-reference.md
• For simulation code patterns from examples, see examples.md
• Asset / DB Asset Server docs: https://docs.elodin.systems/reference/db-asset-server/
• Online docs: https://docs.elodin.systems