Batched time integration demo

This example shows how to use kinamax to run many time integrations in parallel (batched over initial conditions and a frequency sweep), detect steady-state orbits/attractors, and then reconstruct representative limit cycles for post-processing and plotting.

The dynamical system used here is kinamax.integration.models.H46_EM_Problem, a driven nonlinear oscillator with an electromechanical (EM) coupling term. The workflow is generic: you can swap in another Container-style problem as long as it provides rhs() and labels.

What you get

  • A frequency sweep integrated for many random initial conditions, written to outputs/simulations.parquet

  • Orbit detection results (orbit labels + mapping back to simulations), written to:

    • outputs/orbits.parquet

    • outputs/sim_orbit.parquet

  • One-period sampled trajectories reconstructed from each attractor, written to outputs/orbits_from_attractors/

  • A compact per-orbit summary table (including average power), written to outputs/orbit_data.parquet

How it works (step-by-step)

Step 1 — Integrate until convergence

The first notebook:

  • Build an AttractorFinder (Diffrax solver + adaptive controller)

  • Create a frequency sweep (fd) and a batch of random initial conditions

  • vmap + jit the attractor search so it runs efficiently on a single device

  • Post-process the raw solver output into a tidy polars.DataFrame

  • Save the result to outputs/simulations.parquet

This is the “batched time integration” part of the example.

Step 2 — Detect orbits

The second notebook calls kinamax.integration.core.detect_orbits, which groups converged attractors into orbit labels and produces:

  • outputs/orbits.parquet: attractor/orbit metadata and representative attractor states

  • outputs/sim_orbit.parquet: mapping between simulation runs and detected orbits

Note: the EM model tracks cumulative energy states (Ev, Eh). Orbit detection focuses on the dynamical variables, so the script zeroes the energy-like states before saving.

Step 3 — Reconstruct limit cycles

The third notebook:

  • Integrate one drive period starting from each detected attractor state

  • Save a dense set of samples per period (configurable via samples_per_period)

  • Write one parquet file per (orbit_label, attractor_label) in outputs/orbits_from_attractors/

  • Recompute and store per-orbit energy/power summaries in outputs/orbit_data.parquet

Step 4 — Plot

Open step04_plots.ipynb to visualise:

  • orbit branches across the frequency sweep

  • representative state trajectories

  • energy / power metrics derived from the reconstructed cycles

Running the example

From docs/examples/integration/batched_time_integration/, run the first three stages in order:

python step01_integrate_until_convergence.py
python step02_detect_orbits.py
python step03_calculate_orbits.py

Or run the convenience script that executes steps 1–3 in order:

bash run_steps_1_to_3.sh

Then open step04_plots.ipynb in Jupyter to generate plots.

Tips for adapting

  • Sweep resolution: edit fd = jnp.linspace(...) in step01_integrate_until_convergence.py.

  • Basin exploration: edit Nstart and the initial condition scaling in step01_integrate_until_convergence.py.

  • Orbit sampling: edit OrbitCalculator(samples_per_period=...) in step03_calculate_orbits.py.

  • Model parameters: override fields when building H46_EM_Problem instances (e.g. Ad, Q, R, …).