Jaxonomy — hybrid dynamical systems on JAX

---

Jaxonomy とは？

Jaxonomy は、ハイブリッド力学系（連続的な物理現象、離散的な制御則、イベント駆動ロジックを単一のモデルに組み合わせたシステム）をモデル化、シミュレーション、最適化するための Python フレームワークです。すべてのシミュレーションは JAX 上で実行されるため、JIT コンパイルや vmap によるバッチ処理に対応し、エンドツーエンドで完全に微分可能です。

import jax.numpy as jnp
import jaxonomy as jx

# Double integrator: A, B, C, D define the plant; Q, R weight the LQR cost.
A, B = jnp.array([[0., 1.], [0., 0.]]), jnp.array([[0.], [1.]])
C, D = jnp.eye(2), jnp.zeros((2, 1))
Q, R = jnp.eye(2), jnp.array([[1.]])

builder = jx.DiagramBuilder()
# Start displaced from the origin, so the regulator has something to do.
plant      = builder.add(jx.library.LTISystem(A, B, C, D,
                                              initialize_states=jnp.array([1.0, 0.0])))
controller = builder.add(jx.library.LinearQuadraticRegulator(A, B, Q, R))
builder.connect(plant.output_ports[0],      controller.input_ports[0])
builder.connect(controller.output_ports[0], plant.input_ports[0])

diagram = builder.build()
results = jx.simulate(
    diagram,
    diagram.create_context(),
    (0.0, 10.0),
    # Nothing is stored unless you name it here: without recorded_signals,
    # results.time and results.outputs come back as None.
    recorded_signals={"x": plant.output_ports[0],
                      "u": controller.output_ports[0]},
)

results.outputs["x"]    # (T, 2) — position and velocity, driven back to zero

---

🔥 JAX を選ぶ理由

計算基盤として JAX を選ぶことで、NumPy ベースのシミュレーターでは実現が難しい機能を利用できます。

Traditional simulator          Jaxonomy / JAX
──────────────────────         ─────────────────────────────────────────
simulate(params)          →    jit(simulate)(params)          10–100× faster
for p in param_grid: …    →    vmap(simulate)(param_grid)     embarrassingly parallel
finite_diff_gradient(…)   →    grad(simulate)(params)         exact gradients, free

機能	SciPy／NumPy	Julia／DiffEq	Modelica	MathWorks¹	Jaxonomy
Python ネイティブ	✓	✗	✗	✗	✓
JIT／コード生成	✗	✓	✓ (C++)	✓ (C/C++)	✓
ODE を通じた完全な自動微分	✗	一部対応	✗	一部対応²	✓
ハイブリッドイベントとゼロクロッシング	一部対応	✓	✓	✓	✓
非因果的／方程式ベース	✗	✗	✓	✓ (Simscape)	✓
ブロック線図による構成	✗	一部対応	一部対応	✓ (Simulink)	✓
状態機械のモデル化	✗	✗	✗	✓ (Stateflow)	✓
LQR／MPC／カルマンフィルターを標準搭載	✗	一部対応	ライブラリ経由	✓ (Toolboxes)	✓
ニューラル ODE／SINDy	✗	✓	✗	✗	✓
低次元化モデリング（balred／POD-DEIM／DMD／Koopman）	✗	一部対応	✗	✓ (Toolboxes)	✓
バッチ／アンサンブル（vmap）	✗	✗	✗	✗	✓
オープンソース（MIT）	✓	✓	一部対応	✗	✓

_{¹ Simulink + Simscape + Stateflow + Control System Toolbox  ·  ² Simulink Design Optimization 経由。エンドツーエンドの AD には非対応}

---

⚡ 主な機能

機能	実現できること
⚡ JAX ネイティブエンジン	シミュレーションを JIT コンパイルし、vmap でアンサンブルを実行できます。イベント処理を含む ODE ソルバーを通じて微分できます
🔀 ハイブリッドダイナミクス＋ステートマシン	連続 ODE、周期的な離散更新、ゼロ交差イベント、および StateMachineBuilder で記述した有限ステートマシンを、1 つのモデルに組み合わせられます。ハイブリッド軌道最適化では、イベント時刻を通じて jax.grad の勾配が伝播します
🔌 因果関係を前提としないモデリング	Modelica に着想を得たマルチドメインコンポーネント（電気、機械、熱、流体、バッテリー）を、Pantelides インデックス削減および BDF 質量行列 DAE ソルバーとともに利用できます
🎯 制御と推定	LQR（連続、離散、有限時間、LQG）、線形 MPC（ネイティブ＋OSQP）、非線形 MPC（シューティング／転写／Hermite-Simpson）、Kalman／EKF／UKF／RLS／Luenberger、古典的なチューニングヘルパー付き 2 自由度 PID
🧮 単位を考慮した接続	ポートとシグナルに BusUnit アノテーションを任意で付与できます。ダイアグラムコンパイラーが、実行時に気付かれないバグとして発生する代わりに、ビルド時に次元の不一致を検出します
🧠 データ駆動型モデリング	Neural ODE、Universal Differential Equations、SINDy シンボリック回帰、ニューラルネットワークブロック（MLP／PyTorch／TensorFlow／ONNX）、微分可能なルックアップテーブルのフィッティング、統計的サロゲート（ガウス過程、多項式カオス、RBF）
📉 低次元化モデリング	jaxonomy.library.rom：線形 MOR（バランス打ち切り、minreal、モード／残差化）、DEIM によるハイパーリダクションを用いた POD–Galerkin、データ駆動型オペレーター ROM（DMD／DMDc／ERA、Koopman／eDMD リフト線形予測器）。単一の reduce(...) インターフェースで利用でき、すべての低次元化モデルは微分可能かつシミュレーション可能なブロックです
🎲 不確実性と感度	主要機能として提供する jaxonomy.uq ワークフロー：パラメーター分布を用いたモンテカルロ法、ラテン超方格法＋準モンテカルロサンプリング、Sobol 感度分解、Morris スクリーニング
🤝 FMI 2.0／3.0 相互運用性	FMI 協調シミュレーション FMU（ModelicaFMU、混合型および配列 I/O）またはモデル交換 FMU（ModelicaFMUME。FMI イベントインジケーターをゼロ交差として扱い、Jaxonomy 独自のソルバーで統合）をインポートできます。また、build_fmu でダイアグラムをバイナリ .fmu としてエクスポートできます。CI で Reference FMUs、OpenModelica、fmusim リファレンスシミュレーターとの動作を検証済みです
🧩 150 以上のライブラリブロック	積分器、フィルター、ステートマシン、ルックアップテーブル、座標変換、コンテナーブロック、バス／マルチプレクサーファミリー、確率的ソースなど

---

📦 インストール

Python 3.10 以降が必要です。

# Create and activate a virtual environment (recommended)
python -m venv .venv
source .venv/bin/activate        # Windows: .venv\Scripts\activate

# Install
pip install jaxonomy             # core
pip install jaxonomy[safe]       # + SciPy, Matplotlib, control, jaxopt
pip install jaxonomy[nmpc]       # + nonlinear MPC (requires IPOPT on PATH)
pip install jaxonomy[all]        # + everything

ソースからインストールする場合：

git clone https://github.com/machinavitalis/jaxonomy
cd jaxonomy
pip install -e .

CLI ランナー：

jaxonomy_cli run --model path/to/model.json

---

🚀 クイックスタート

数行で最初のシミュレーションを実行できます。カスタムブロックを作成し、ダイアグラムに組み込んで、その ODE を積分します。

import jaxonomy as jx
import jax.numpy as jnp

# Van der Pol oscillator as a custom block
class VanDerPol(jx.LeafSystem):
    def __init__(self, mu=1.0, **kwargs):
        super().__init__(**kwargs)
        self.declare_dynamic_parameter("mu", mu)
        self.declare_continuous_state(
            default_value=jnp.array([0.0, 2.0]), ode=self._ode
        )
        self.declare_continuous_state_output(name="x")

    def _ode(self, time, state, *inputs, **params):
        x, mu = state.continuous_state, params["mu"]
        return jnp.array([x[1],  mu * (1 - x[0]**2) * x[1] - x[0]])

builder = jx.DiagramBuilder()
vdp = builder.add(VanDerPol(mu=2.0, name="vdp"))
diagram = builder.build()

results = jx.simulate(
    diagram, diagram.create_context(), (0.0, 20.0),
    options=jx.SimulatorOptions(buffer_length=4000),  # room for adaptive steps
    recorded_signals={"x": vdp.output_ports[0]},
)
# results.outputs["x"] → time-series of shape (T, 2)

---

📚 ドキュメント

• オンラインドキュメントとチュートリアル： py.jaxonomy.com
• ローカルドキュメント：

  pip install -r requirements.docs.txt
  mkdocs serve   # → http://127.0.0.1:8000
  

• サンプルノートブック： docs/examples/
• 対象範囲に関する注記： PINNs と PDE サロゲート — 古典的な PDE PINN は対象外です。物理法則に基づく ダイナミクス 学習（UDE／ Neural DAE／Neural ODE／SINDy）は中核機能です。

---

🤖 AI エージェントから Jaxonomy を操作する（MCP）

Jaxonomy には MCP サーバーが同梱されており、エンジンを AI エージェントが直接呼び出せるツールとして公開します。ライブラリのブロックを列挙したり、モデルを構築して検証したり、シミュレーションを実行したり、データにパラメーターを適合させたり、システムを線形化したりして、実際の結果に基づいて推論できます。エージェント（Claude Desktop/Code、Cursor など）から Jaxonomy を操作する場合は、接続しておくと便利です。Python を手動で記述する場合は、上記の pip install だけで十分なので、この手順は省略できます。

pip install jaxonomy[mcp]

次に、エージェントのクライアントにサーバーを登録します。Claude Desktop の場合は、claude_desktop_config.json に追加します。

{
  "mcpServers": {
    "jaxonomy": {
      "command": "python",
      "args": ["-m", "jaxonomy.mcp.server"]
    }
  }
}

jaxonomy[mcp] がインストールされているインタープリター（または jaxonomy-mcp エントリーポイント）を使用してください。このサーバーは MCP Registry に io.github.machinavitalis/jaxonomy として掲載されています。ツールの詳細なリファレンス、クライアントの設定、制限事項については、py.jaxonomy.com/mcp をご覧ください。

ドキュメントをエージェントに参照させるだけでよければ、SKILL.md を読み込ませてください。Jaxonomy を使うべき状況、別の手段を選ぶべき状況、コア API、最初のスクリプトでよく起きる問題について説明しています。ドキュメントサイトでは、llms.txt 形式のページも公開しています。すべてのページのインデックスは py.jaxonomy.com/llms.txt、ドキュメント一式をまとめた Markdown ファイルは py.jaxonomy.com/llms-full.txt をご覧ください。

---

📖 使用例

1 · ハイブリッドダイナミクス：跳ねるボール

Jaxonomy はハイブリッドシステム向けに設計されています。これは、連続的な物理現象と瞬時の離散的なリセットが相互作用するモデルです。跳ねるボールは典型的な例です。自由落下の常微分方程式が、速度を反転させる衝突イベントによって中断されます。

支配方程式

跳ね返りの間のダイナミクスは、次の式に従います。

$$ \dot{x} = v, \qquad \dot{v} = -g $$

ボールが地面に当たる（$x = 0$、$v < 0$）と、ゼロ交差イベントが発生し、状態がリセットされます。

$$ x^+ = 0, \qquad v^+ = -e \cdot v \quad (e \in [0, 1] \text{ — 反発係数}) $$

import jaxonomy as jx
import jax.numpy as jnp

class BouncingBall(jx.LeafSystem):
    def __init__(self, g=9.81, e=0.8, **kwargs):
        super().__init__(**kwargs)
        self.declare_dynamic_parameter("g", g)
        self.declare_dynamic_parameter("e", e)
        # [height, velocity]
        self.declare_continuous_state(
            default_value=jnp.array([5.0, 0.0]), ode=self._ode
        )
        self.declare_continuous_state_output(name="state")

        # Zero-crossing: fires as height crosses zero from above
        self.declare_zero_crossing(
            guard=self._hit_ground,
            reset_map=self._bounce,
            direction="positive_then_non_positive",
        )

    def _ode(self, time, state, *inputs, **params):
        v = state.continuous_state[1]
        return jnp.array([v, -params["g"]])

    def _hit_ground(self, time, state, *inputs, **params):
        return state.continuous_state[0]  # guard on height

    def _bounce(self, time, state, *inputs, **params):
        x, v = state.continuous_state
        new_state = jnp.array([0.0, -params["e"] * v])
        return state.with_continuous_state(new_state)

builder = jx.DiagramBuilder()
ball = builder.add(BouncingBall(g=9.81, e=0.85, name="ball"))
diagram  = builder.build()
results  = jx.simulate(
    diagram, diagram.create_context(), (0.0, 8.0),
    recorded_signals={"state": ball.output_ports[0]},
)

ゼロ交差イベントは、ソルバーの密出力多項式に対する40 ステップの二分法で特定されます。ユーザーが許容誤差を指定しなくても、マイクロ秒未満の時間精度が得られます。

---

2 · 最適制御：LQR 振り子

状態が $x = [\theta, \dot\theta]^\top$ の線形化された振り子の場合：

$$ \dot{x} = Ax + Bu, \qquad A = \begin{bmatrix} 0 & 1 \ g/L & 0 \end{bmatrix}, \quad B = \begin{bmatrix} 0 \ 1/mL^2 \end{bmatrix} $$

線形二次レギュレーターは、次の無限時間範囲のコストを最小化します。

$$ J = \int_0^\infty \bigl( x^\top Q, x + u^\top R, u \bigr), dt $$

代数リカッチ方程式 $A^\top P + PA - PBR^{-1}B^\top P + Q = 0$ を解いて最適ゲイン $K = R^{-1}B^\top P$ を求めることで、$u^* = -Kx$ となります。

import jaxonomy as jx
from jaxonomy.library import LTISystem, LinearQuadraticRegulator
import jax.numpy as jnp

g, L, m = 9.81, 1.0, 1.0
A = jnp.array([[0, 1], [g/L, 0]])
B = jnp.array([[0], [1 / (m * L**2)]])
C, D = jnp.eye(2), jnp.zeros((2, 1))

Q = jnp.diag(jnp.array([10.0, 1.0]))  # penalise angle more than rate
R = jnp.array([[0.1]])                 # control effort cost

builder    = jx.DiagramBuilder()
plant      = builder.add(LTISystem(A, B, C, D, name="pendulum"))
controller = builder.add(LinearQuadraticRegulator(A, B, Q, R, name="lqr"))

builder.connect(plant.output_ports[0],      controller.input_ports[0])
builder.connect(controller.output_ports[0], plant.input_ports[0])

diagram = builder.build()
context = diagram.create_context()

# Perturb the pendulum's initial angle by 15° (set one block's sub-state)
context = context.with_subcontext(
    plant.system_id,
    context[plant.system_id].with_continuous_state(jnp.array([jnp.pi / 12, 0.0])),
)

results = jx.simulate(
    diagram, context, (0.0, 5.0),
    recorded_signals={"x": plant.output_ports[0]},
)

---

3 · 微分可能なパラメーター同定

Jaxonomy では、シミュレーション全体を通して微分し、モデルのパラメーターをデータに適合させられます。有限差分近似も、手書きの随伴コードも必要ありません。ここでは、ロールアウト全体を微分し、Optax で最適化することで、ばね・ダンパーの剛性と減衰を求めます。

import jax
import jax.numpy as jnp
import optax
import jaxonomy as jx

class SpringDamper(jx.LeafSystem):
    def __init__(self, k=1.0, c=0.3, **kwargs):
        super().__init__(**kwargs)
        self.declare_dynamic_parameter("k", k)
        self.declare_dynamic_parameter("c", c)
        self.declare_continuous_state(
            default_value=jnp.array([1.0, 0.0]), ode=self._ode
        )

    def _ode(self, time, state, *inputs, **params):
        x, v = state.continuous_state
        return jnp.array([v, -params["k"] * x - params["c"] * v])

sd   = SpringDamper(name="sd")
opts = jx.SimulatorOptions(enable_autodiff=True, max_major_steps=200)

def final_state(theta):                       # theta = {"k": ..., "c": ...}
    ctx = sd.create_context()
    ctx.parameters["k"], ctx.parameters["c"] = theta["k"], theta["c"]
    res = jx.simulate(sd, ctx, (0.0, 1.5), options=opts)
    return res.context.continuous_state       # differentiable final [x, v]

target = jax.lax.stop_gradient(final_state({"k": 4.0, "c": 0.5}))  # "measured"
loss   = lambda theta: jnp.sum((final_state(theta) - target) ** 2)

theta   = {"k": 1.0, "c": 0.1}
opt     = optax.adam(2e-1)
state   = opt.init(theta)
grad_fn = jax.jit(jax.grad(loss))             # gradient through the ODE solver
for _ in range(300):
    updates, state = opt.update(grad_fn(theta), state)
    theta = optax.apply_updates(theta, updates)
# theta → {"k": 4.00, "c": 0.50}

jax.grad(loss) はすべての ODE ステップを通して自動的に勾配を伝播するため、同じ手法をバッテリー ECM（上記）、パワートレイン、その他あらゆるパラメトリックモデルに適用できます。また、境界、変換、制約が必要な場合は、jaxonomy.optimization がより高水準の Optimizable API でこの機能をラップします。

---

4 · 非因果的な物理モデリング

Jaxonomy には、マルチドメイン物理システム向けの Modelica に着想を得た非因果モデリングレイヤーが含まれています。コンポーネントの接続を記号的に記述すると、コンパイラーが支配 DAE を自動的に導出してインデックスを下げ、あらゆるダイアグラムに追加できる LeafSystem を生成します。

初期条件付きRC回路：

$$ C,\dot{V}C = I, \qquad V_C(0) = 0,\text{V}, \qquad V{\text{src}} = 1,\text{V} $$

import jaxonomy as jx
from jaxonomy.acausal import AcausalCompiler, AcausalDiagram, EqnEnv
from jaxonomy.acausal import electrical as elec

ev = EqnEnv()
ad = AcausalDiagram()

vs  = elec.VoltageSource(ev, name="vs", v=1.0)
r   = elec.Resistor(ev,      name="r",  R=1.0)
c   = elec.Capacitor(ev,     name="c",  C=1.0,
                              initial_voltage=0.0, initial_voltage_fixed=True)
gnd = elec.Ground(ev,        name="gnd")

ad.connect(vs, "p", r,  "n")
ad.connect(r,  "p", c,  "p")
ad.connect(c,  "n", vs, "n")
ad.connect(vs, "n", gnd, "p")

compiler = AcausalCompiler(ev, ad)
rc_block = compiler()         # → LeafSystem, JIT-compiled ODE

builder = jx.DiagramBuilder()
rc = builder.add(rc_block)
diagram = builder.build()
results = jx.simulate(
    diagram,
    diagram.create_context(),
    (0.0, 5.0),
    recorded_signals={"rc": rc.output_ports[0]},
)

results.outputs["rc"]   # capacitor charges 0 → 0.993 V over 5 τ (R = C = 1)

同じパイプラインでマルチドメインシステムにも対応します。electrical、rotational、translational ドメインを接続する電気機械アクチュエーターを、単一の最適化済みシステムにコンパイルできます。

利用可能な非因果ドメイン：electrical · rotational · translational · thermal · fluid

---

5 · 微分可能な感度解析とバッチシミュレーション

jax.grad、jax.jit、jax.vmap はすべて simulate と組み合わせられるため、最小限の定型コードで強力なワークフローを実現できます。

import jax
import jax.numpy as jnp
import jaxonomy as jx

class SpringMass(jx.LeafSystem):
    def __init__(self, mass=1.0, damping=0.1, stiffness=10.0, **kwargs):
        super().__init__(**kwargs)
        self.declare_dynamic_parameter("mass", mass)
        self.declare_dynamic_parameter("damping", damping)
        self.declare_dynamic_parameter("stiffness", stiffness)
        self.declare_continuous_state(
            default_value=jnp.array([1.0, 0.0]), ode=self._ode
        )
        self.declare_continuous_state_output(name="x")

    def _ode(self, time, state, *inputs, **params):
        x, v = state.continuous_state
        a = -(params["stiffness"] * x + params["damping"] * v) / params["mass"]
        return jnp.array([v, a])

# ── Sensitivity: ∂(final position)/∂(all parameters) in one reverse pass ─────
plant = SpringMass(name="plant")
grad_opts = jx.SimulatorOptions(enable_autodiff=True, max_major_steps=200)

def final_position(theta):
    ctx = plant.create_context()
    for name, value in theta.items():
        ctx.parameters[name] = value
    res = jx.simulate(plant, ctx, (0.0, 5.0), options=grad_opts)
    return res.context.continuous_state[0]

grads = jax.grad(final_position)({"mass": 1.0, "damping": 0.1, "stiffness": 10.0})

# ── Monte Carlo ensemble: 1000 trajectories over a mass sweep (vmap) ─────────
builder = jx.DiagramBuilder()
plant_b = builder.add(SpringMass(name="plant"))
diagram = builder.build()

results = jx.simulate_batch(
    diagram, t_span=(0.0, 5.0),
    param_batches={"plant.mass": jnp.linspace(0.5, 2.0, 1000)},
    options=jx.SimulatorOptions(math_backend="jax", max_major_steps=200),
    recorded_signals={"x": plant_b.output_ports[0]},
)
# results.outputs["x"] shape: (1000, T, 2)

---

🧩 ライブラリの概要

信号処理と制御に必要なツールキット全体を網羅する、150 種類以上の組み込みブロックを備えています。代表的な一部を以下に示します。すべてを網羅したものではありません。

カテゴリ	ブロック
ソース	Constant、Step、Ramp、Chirp、Pulse、Sawtooth、Clock、DataSource
確率的処理	RandomNumber、UniformRandomNumber、WhiteNoise、BandLimitedNoise、PRBS、RandomSource — すべて with_key に対応しており、vmap の下で独立したノイズストリームを使用できます
算術演算と非線形要素	Adder、Gain、Product、Abs、Power、Trigonometric、Saturate / SoftSaturate、DeadZone、RateLimiter / SoftRateLimiter、Quantizer、Backlash
連続系のダイナミクス	Integrator、Derivative、TransferFunction、LTISystem、PID / PIDContinuous
離散系のダイナミクス	IntegratorDiscrete、UnitDelay、FilterDiscrete、LowPassDiscrete、LeadLag、Notch、PIDDiscrete、PIDController2DOF、Decimator、RateTransition
ルーティング、バス、行列	Mux / Demux、BusCreator / BusSelector / BusUpdate（BusUnit アノテーションに対応）、Slice、Stack、Switch / MultiPortSwitch、IfThenElse、MatrixMultiplication、MatrixInversion、DotProduct、CrossProduct
論理とステートマシン	TruthTable、LogicalOperator、Comparator、Relay、EdgeDetection、StateMachine（StateMachineBuilder DSL で作成）
ルックアップテーブル	LookupTable1d / LookupTable2d / LookupTableND、Prelookup / InterpolationUsingPrelookup、TableSearch — すべて微分可能で、fit_lookup_table_* を使ってデータからフィッティングできます
遅延とコンテナ	TransportDelay、VariableTransportDelay、EnabledSubsystem、TriggeredSubsystem、ForEach、Conditional
制御	LinearQuadraticRegulator / DiscreteTimeLinearQuadraticRegulator / FiniteHorizonLinearQuadraticRegulator / LinearQuadraticGaussian、LinearDiscreteTimeMPC（ネイティブ版と LinearDiscreteTimeMPC_OSQP）、DirectShootingNMPC / DirectTranscriptionNMPC / HermiteSimpsonNMPC
推定	KalmanFilter、ExtendedKalmanFilter、UnscentedKalmanFilter、InfiniteHorizonKalmanFilter、RecursiveLeastSquares、AugmentedStateEKF、Luenberger
物理と座標	CoordinateRotation、RigidBody、BatteryCell
ML / データ	MLP（Equinox）、Sindy、PyTorch、TensorFlow、ONNX / ONNXJax
ROM とサロゲート	reduce(...) → ReducedOrderModel、線形 MOR（balred / minreal / modal_truncation / residualize）、galerkin_reduce + deim（POD–DEIM）、dmd / dmdc / era、DMDForecaster / KoopmanPredictor、およびサロゲートブロック GaussianProcess / PolynomialChaos / RadialBasisSurrogate
相互運用	ModelicaFMU / ModelicaFMUME（FMI 2.0 / 3.0 のコシミュレーションとモデル交換のインポート）+ build_fmu による FMU エクスポート、MuJoCo / MJX、Ros2Publisher / Ros2Subscriber、QuanserHAL、PyTwin
カスタム	インスタンスごとに状態を保持する、ユーザー作成アルゴリズム向けの CustomPythonBlock、CustomJaxBlock

上記のラッパーに限らず、カスタムブロックも第一級の機能として扱われます。LeafSystem をサブクラス化してポートと状態を宣言すると、JIT、 autodiff、イベント検出を含むフレームワーク全体にブロックを統合できます。

---

🏗️ システムアーキテクチャ

Jaxonomy のモデルは、相互接続されたブロックからなる有向グラフ、Diagram です。ブロック（LeafSystem）では、ポート、連続／離散状態、イベントガードを宣言します。シミュレーターは、ODE の積分、離散更新、イベント検出を自動的に統括します。

graph LR
    subgraph Diagram["DiagramBuilder.build()"]
        direction LR
        r[/"r(t)\nReference"/] --> sum(("Σ"))
        sum -->|"e(t)"| ctrl["Controller\nLQR · MPC · PID"]
        ctrl -->|"u(t)"| plant["Plant\nLTI · Nonlinear · Acausal"]
        plant -->|"x(t)"| obs["State Estimator\nKalman · EKF · UKF"]
        obs --> sum
        plant -->|"y(t)"| out[/"y(t)\nOutput"/]
    end

    style Diagram fill:#f5f8ff,stroke:#4a6fa5,stroke-width:2px
    style ctrl fill:#dbeafe,stroke:#2563eb
    style plant fill:#dcfce7,stroke:#16a34a
    style obs fill:#fef9c3,stroke:#ca8a04

JAX バックエンドはすべてのシミュレーションを支え、XLA コンパイルに基づくシンプルな Python API を提供します。

graph TD
    A["Python API\nDiagram · LeafSystem · simulate()"] --> B["JAX Backend\nJIT · vmap · grad"]
    B --> C1["Dopri5\nAdaptive RK45\nnon-stiff ODEs"]
    B --> C2["BDF\nImplicit solver\nstiff systems"]
    B --> C3["Event Handler\nZero-crossing\nbisection"]

    style A fill:#e0f2fe,stroke:#0284c7
    style B fill:#fdf4ff,stroke:#9333ea
    style C1 fill:#f0fdf4,stroke:#16a34a
    style C2 fill:#f0fdf4,stroke:#16a34a
    style C3 fill:#f0fdf4,stroke:#16a34a

---

Jaxonomy スタック

Jaxonomy は、3 つのパッケージからなるスタックの基盤となるエンジンです。各パッケージはそれぞれ独立したリポジトリで、MIT ライセンスのもとで提供され、直上のパッケージのみに依存します。

• Jaxonomy — このパッケージです。 ハイブリッド力学系向けの汎用 JAX ネイティブシミュレーションエンジンです。ロボティクスに特化しておらず、スタック内の他のパッケージには依存しません。
• Jaxterity — Jaxonomy の上に構築されたロボティクスレイヤーです。URDF/MJCF のインポート、MJX を利用した多関節ダイナミクス、校正済みのアクチュエーター／センサー、システム同定、全身制御を提供します。Jaxonomy をインポートし、そのプリミティブを再実装することはありません。
• Jaxility — デプロイ用アーティファクトのファクトリーです。校正済みロボットを、署名付きアテステーションマニフェストを含む Arm SoC（Cortex-A / Cortex-M）向け組み込み C にコンパイルします。Jaxterity を利用します。

境界のルール：ロボティクス以外（HVAC、バッテリー、航空宇宙、エネルギー）で制御エンジニアに役立つ機能は Jaxonomy に属します。関節、アクチュエーター、接触、キネマティックチェーンに特化した機能は Jaxterity に属します。組み込みコード生成、ターゲット、アテステーションは Jaxility に属します。

Jaxonomy がまだ対応していない（または部分的にしか対応していない）項目は、CLAIMS.md の内部証拠台帳を公開向けに反転した KNOWN_GAPS.md で管理されています。

ライセンス

バージョン 2.2.0 以降、MIT ライセンスのもとでリリースされています。

MIT ライセンスのオープンソースパッケージ pycollimator（Collimator, Inc.）を派生元としています。