Single Time Integration Until Convergence

This tutorial shows the smallest complete attractor-search workflow in kinamax.

It mirrors the batched example, but restricts the computation to one problem instance and one initial condition:

  • one driven oscillator, kinamax.integration.models.H46Problem

  • one drive frequency, fd

  • one initial state, X0

The goal is to make each moving part visible before introducing batching and orbit clustering.

What The Script Does

The workflow has four stages:

  1. define the physical problem and the attractor-finder settings,

  2. generate one initial condition,

  3. integrate until one candidate subharmonic satisfies the shooting residual tolerance,

  4. flatten the structured result into a table that is easy to inspect and save.

Because this file uses Jupytext percent format, it remains a normal Python script while also behaving like a notebook in the documentation.

from __future__ import annotations

from pathlib import Path

import jax
from jax import config
import jax.numpy as jnp
import numpy as np
import polars as pl
from diffrax import PIDController, Tsit5

from kinamax.integration.core import (
    AttractorFinder,
    AttractorFinderConfig,
    post_process_attractor_finder_results,
)
from kinamax.integration.models import H46Problem

config.update("jax_enable_x64", True)  # Use double precision for improved accuracy

Set Up The Output Directory

The example writes its results next to the source file so the generated parquet file stays attached to the tutorial.

example_dir = Path(__file__).resolve().parent if "__file__" in globals() else Path.cwd()
output_dir = example_dir / "outputs"
output_dir.mkdir(parents=True, exist_ok=True)

Define The Physical Problem

In this first parameter block we choose the forcing frequency and instantiate one H46 oscillator. The problem namespace also exposes the state ordering, which is reused later when we format the output table.

fd = 50.0
problem_class = H46Problem
state_vector_labels = problem_class.state_vector_labels
problem = H46Problem.Params(fd=jnp.array(fd), Ad=jnp.array(2.5))

Define The Attractor-Finder Settings

The attractor search has two levels of configuration:

  • AttractorFinderConfig controls the integration start time, initial time step, target frequency, and convergence threshold.

  • AttractorFinder.Params controls the overall shooting strategy: how many residual samples are used per period, which subharmonics are tested, and how much total integration effort is allowed.

finder_config = AttractorFinderConfig(
    convergence_tol=jnp.array(1.0e-10),
    target_frequency=fd,
    init_time=jnp.array(0.0),
    init_time_step=jnp.array(1.0e-3),
    subharmonic_factor=10.0,
)

solver = Tsit5()
controller = PIDController(rtol=1e-8, atol=1e-9)
target_subharmonics = np.array([1, 2, 3, 5], dtype=int)
attractor_finder = AttractorFinder.Params(
    residuals_per_period=np.array(20, int),
    targetted_subharmonics=target_subharmonics,
    max_periods=2000,
    controller=controller,
    solver=solver,
)

Generate One Initial Condition

The state is sampled in a box scaled by the characteristic displacement and velocity of the oscillator. The last state corresponds to cumulative energy, so it starts from zero.

key = jax.random.PRNGKey(758493)
init_condition = (
    (jax.random.uniform(key, shape=(len(state_vector_labels),)) - 0.5)
    * 2.0
    * jnp.array([5.0 * problem.xw, 10.0 * problem.xw * problem.w0, 0.0])
)

Integrate Until Convergence

AttractorFinder.find_attractors(...) returns the original inputs together with a structured solution object. The call is JIT-compiled here because even a single search performs repeated ODE solves internally.

find_one = jax.jit(
    lambda p, x0, cfg: AttractorFinder.find_attractors(
        attractor_finder, problem_class, p, x0, cfg
    )
)
problem, finder_config, init_condition, solution = find_one(
    problem, init_condition, finder_config
)

Why The Post-Processing Step Matters

AttractorFinder.find_attractors(...) returns structured arrays designed for computation. post_process_attractor_finder_results(...) converts that result into a tidy table with one row per detected attractor sample and with the problem/configuration metadata duplicated alongside it. That table is much easier to save, filter, plot, and compare across examples.

processed = post_process_attractor_finder_results(
    problem_class=problem_class,
    problems=problem,
    finder_configs=finder_config,
    init_conditions=init_condition,
    solutions=solution,
    target_subharmonics=target_subharmonics,
    solution_state_labels=[lab + "a" for lab in state_vector_labels],
)

out_path = output_dir / "simulations.parquet"
processed.write_parquet(out_path)
pl.Config.set_tbl_rows(8)
print(processed)
print(f"Wrote: {out_path}")

shape: (3, 26)
┌───────────┬────────┬──────┬───────┬───┬────────────────┬────────────┬────────────────┬───────────┐
│ sim_label ┆ xw     ┆ fd   ┆ w0    ┆ … ┆ simulated_time ┆ final_flag ┆ simulated_iter ┆ converged │
│ ---       ┆ ---    ┆ ---  ┆ ---   ┆   ┆ ---            ┆ ---        ┆ ations         ┆ ---       │
│ f64       ┆ f32    ┆ f64  ┆ f32   ┆   ┆ f64            ┆ i64        ┆ ---            ┆ bool      │
│           ┆        ┆      ┆       ┆   ┆                ┆            ┆ i64            ┆           │
╞═══════════╪════════╪══════╪═══════╪═══╪════════════════╪════════════╪════════════════╪═══════════╡
│ 0.0       ┆ 0.0005 ┆ 50.0 ┆ 121.0 ┆ … ┆ 17.0           ┆ 1          ┆ 85             ┆ true      │
│ 0.0       ┆ 0.0005 ┆ 50.0 ┆ 121.0 ┆ … ┆ 17.0           ┆ 1          ┆ 85             ┆ true      │
│ 0.0       ┆ 0.0005 ┆ 50.0 ┆ 121.0 ┆ … ┆ 17.0           ┆ 1          ┆ 85             ┆ true      │
└───────────┴────────┴──────┴───────┴───┴────────────────┴────────────┴────────────────┴───────────┘
Wrote: /Users/lcharleux/Documents/GitHub/kinamax/docs/examples/integration/single_time_integration/outputs/simulations.parquet