Compile

Diagram compiler: topology analysis and execution plan construction.

This module provides the complete compilation pipeline:

  1. evaluation_order_of() — Depth-first search over the port dependencies (implemented in wiring.py): the topological execution order of the output ports, an algebraic loop refused. The standalone check after manual wiring is check_algebraic_loops().

  2. build_execution_plan() — Constructs an immutable ExecutionPlan from a diagram. Calls evaluation_order_of() internally.

  3. compile_diagram() — Top-level entry point that builds the plan and wraps it in a backend evaluator.

minilink.core.compile.compiler.validate_equation_shapes(system, label=None)[source]

Probe the textbook hooks once and check their declared shapes.

f (continuous leaves), step (discrete leaves), and the standard output ports created by the constructor (y -> h, x -> state) are evaluated at (x0, u_nominal, t=0); a wrong-length result raises ValueError naming the block, the hook, and both shapes. Custom port computes are not probed (they may be expensive or stateful).

minilink.core.compile.compiler.compile(system, backend='numpy', verbose=False)[source]

Compile a system into a typed evaluator.

Returns StaticEvaluator for static System leaves (n=0), StepEvaluator for StepSystem leaves, DynamicsEvaluator for DynamicSystem leaves and diagrams.

Parameters:
  • system (System or DiagramSystem) – The system to compile.

  • backend (str) – 'numpy' or 'jax'.

  • verbose (bool) – If True, print timed compilation steps.

minilink.core.compile.compiler.compile_auto(system, verbose=False)[source]

Compile with JAX when it is installed and system traces, else with NumPy.

Returns (backend, evaluator). Only two verdicts fall back to NumPy: JAX cannot be imported, or the system is “not JAX-traceable”; any other error from the JAX compile surfaces.

minilink.core.compile.compiler.compile_diagram(diagram, backend='numpy', *, bind_params=False, verbose=False)[source]

Compile a DiagramSystem into a backend evaluator.

Parameters:
  • diagram (DiagramSystem) – The wired diagram to compile.

  • backend (str) – 'numpy' or 'jax'.

  • bind_params (bool, optional) – If True, each plan operation stores a deep copy of that subsystem’s params and evaluators pass it into f / port compute. If False (default), None is passed and blocks use live self.params.

  • verbose (bool) – If True, print timed compilation steps.

Returns:

A stateless evaluator that can compute state derivatives and outputs.

Return type:

NumpyDiagramEvaluator or JaxDiagramEvaluator

Examples

>>> from minilink.core.compile.compiler import compile_diagram
>>> evaluator = compile_diagram(diagram)
>>> dx = evaluator.f(x, u, t)

Notes

bind_params=True snapshots only each subsystem’s params dict into the plan. It does not make user f / port compute implementations pure if they still read or mutate other instance state; see minilink.core.system.System.

minilink.core.compile.compiler.build_execution_plan(diagram, *, bind_params=False)[source]

Build an immutable ExecutionPlan from a DiagramSystem.

This is the core compilation step. It:

  1. Orders the output ports, refusing an algebraic loop, via evaluation_order_of().

  2. Assigns each output port a slice in the flat signal buffer.

  3. Builds the gather-source recipes for each port and state operation.

  4. Returns a frozen ExecutionPlan.

Parameters:
Return type:

ExecutionPlan

minilink.core.compile.compiler.warn_shared_seeds(diagram)[source]

Warn when two random blocks hold one seed: they would draw the same signal.

Execution plan data structures for compiled diagrams.

An ExecutionPlan is the output of the compilation step: it captures, in an immutable and flattened form, all the information needed to evaluate a DiagramSystem without any dictionary lookups or recursive calls.

The plan is consumed by evaluator backends (NumpyDiagramEvaluator, JaxDiagramEvaluator) which walk through the operation lists in topological order to compute state derivatives and outputs.

minilink.core.compile.execution_plan.NOMINAL = 0

Source is a constant value (port not connected, or not a dependency).

minilink.core.compile.execution_plan.EXTERNAL_INPUT = 1

Source is an external input to the diagram (from the diagram’s u vector).

minilink.core.compile.execution_plan.INTERNAL_SIGNAL = 2

Source is an output from another subsystem (from the internal signal buffer).

class minilink.core.compile.execution_plan.PortOperation(compute_func, local_x_slice, gather_sources, out_slice, u_dim, bound_params=None, label='', sys_id='')[source]

One output-port evaluation step in topological order.

When a diagram is compiled, every output port becomes a PortOperation. During evaluation, operations execute sequentially in topological order: for each operation, gather the local input vector u from the signal buffer, call the compute function, and write the result into the buffer.

Parameters:
  • compute_func (Callable[[...], ndarray])

  • local_x_slice (slice)

  • gather_sources (tuple[tuple[int, object, int], ...])

  • out_slice (slice)

  • u_dim (int)

  • bound_params (dict | None)

  • label (str)

  • sys_id (str)

compute_func

The port’s compute function (x, u, t, params) → y. params may be None (use live self.params in the block) or a bound snapshot when bind_params=True at compile time.

Type:

Callable[…, np.ndarray]

local_x_slice

Index slice into the global state vector x that extracts this subsystem’s local state. For static systems (n = 0): slice(0, 0).

Type:

slice

gather_sources

Recipe for assembling the local input vector u from various sources. Each entry is (source_type, source_value, dim) where:

  • source_type = NOMINAL (0): Port not connected or not a dependency. source_value is the constant np.ndarray (nominal value).

  • source_type = EXTERNAL_INPUT (1): From the diagram’s external input vector. source_value is a slice into the diagram’s u.

  • source_type = INTERNAL_SIGNAL (2): From another subsystem’s output. source_value is a slice into the internal signal buffer.

Type:

tuple[tuple[int, Any, int], …]

out_slice

Where to write this port’s output in the internal signal buffer.

Type:

slice

u_dim

Total dimension of the assembled local input vector. Equal to the sum of all dim values in gather_sources.

Type:

int

bound_params

If set, a deep copy of the subsystem params at compile time; passed as the fourth argument to compute_func. If None, evaluators pass None so blocks use their default self.params. This does not freeze other instance attributes; purity of compute_func is not enforced.

Type:

dict | None

sys_id

Id of the owning subsystem in the diagram. Used by the parametric evaluator tier to route nested diagram params (params[sys_id]).

Type:

str

class minilink.core.compile.execution_plan.StateOperation(f_func, local_x_slice, gather_sources, u_dim, bound_params=None, label='', sys_id='')[source]

One state-derivative evaluation step.

Similar to PortOperation, but calls f(x, u, t) → dx instead of h(x, u, t) → y. The result is written into the global derivative vector dx at local_x_slice.

Parameters:
  • f_func (Callable[[...], ndarray])

  • local_x_slice (slice)

  • gather_sources (tuple[tuple[int, object, int], ...])

  • u_dim (int)

  • bound_params (dict | None)

  • label (str)

  • sys_id (str)

f_func

The subsystem’s state-derivative function (x, u, t, params) → dx. params may be None or a bound snapshot (see bound_params).

Type:

Callable[…, np.ndarray]

local_x_slice

Index slice into the global state vector x and derivative vector dx for this subsystem’s local state.

Type:

slice

gather_sources

Recipe for assembling the local input vector u. Same format as PortOperation.gather_sources — the f function always needs all inputs (no dependency filtering).

Type:

tuple[tuple[int, Any, int], …]

u_dim

Total dimension of the assembled local input vector.

Type:

int

bound_params

Same semantics as PortOperation.bound_params for f_func (params dict snapshot only; see minilink.core.system.System).

Type:

dict | None

sys_id

Same semantics as PortOperation.sys_id.

Type:

str

class minilink.core.compile.execution_plan.ExecutionPlan(state_dim, signal_dim, port_ops, state_ops, output_slices, external_output_slices)[source]

Immutable flattened execution schedule for a compiled diagram.

Built by build_execution_plan(). Consumed by evaluator backends to compute state derivatives and outputs without any dictionary lookups or recursive traversals.

Parameters:
  • state_dim (int)

  • signal_dim (int)

  • port_ops (tuple[PortOperation, ...])

  • state_ops (tuple[StateOperation, ...])

  • output_slices (dict[tuple[str, str], slice])

  • external_output_slices (dict[str, slice])

state_dim

Total state dimension of the diagram (sum of all subsystem n values).

Type:

int

signal_dim

Total size of the internal signal buffer (sum of all output port dimensions across all subsystems).

Type:

int

port_ops

Output-port operations in topological order. Evaluating them sequentially fills the internal signal buffer.

Type:

tuple[PortOperation, …]

state_ops

State-derivative operations. Evaluating them sequentially (after the port operations) fills the derivative vector dx.

Type:

tuple[StateOperation, …]

output_slices

Map from (subsystem_id, port_id) to the corresponding slice in the internal signal buffer. Used to extract specific port outputs after evaluation.

Type:

dict[tuple[str, str], slice]

external_output_slices

Diagram boundary outputs only: maps diagram output port id (keys of outputs on the diagram) to the slice in the internal buffer that holds the connected source subsystem port’s signal. Empty when the diagram has no external outputs.

Type:

dict[str, slice]

NumPy/JAX backend policy: the backend strings, array_module, and the lazy JAX imports.

minilink.core.backends.COMPILE_BACKENDS: tuple[str, ...] = ('numpy', 'jax')

Backends accepted by minilink.core.compile.compiler.compile().

minilink.core.backends.SIMULATOR_BACKENDS: tuple[str, ...] = ('numpy', 'jax', 'auto')

Backends accepted by Simulator.

minilink.core.backends.TRANSCRIPTION_BACKENDS: tuple[str, ...] = ('numpy', 'jax', 'direct')

Backends accepted by JAX trajectory-optimization transcriptions.

minilink.core.backends.normalize_backend(name, *, allow_auto=False, allow_direct=False)[source]

Validate and lowercase a compile-backend string.

"numpy"

NumPy evaluator. Always available.

"jax"

JAX evaluator. Hard-requires JAX (pip install jax jaxlib).

"auto"

Try JAX, fall back to NumPy. Used by the simulator.

"direct"

Use f() directly (no compiled evaluator); valid for some trajectory-optimization transcriptions on DynamicSystem or DiagramSystem.

"auto" and "direct" are simulator / transcription concepts: the low-level minilink.core.compile.compiler.compile() only accepts "numpy" and "jax"; allow_auto / allow_direct admit them at the layer that understands them.

Parameters:
  • name (str or None) – Backend identifier. None is treated as "numpy".

  • allow_auto (bool, optional) – Allow the "auto" value (Simulator-level concept).

  • allow_direct (bool, optional) – Allow the "direct" value (transcription-level concept).

Returns:

Normalized lowercase backend string.

Return type:

str

minilink.core.backends.array_module(*values)[source]

Return jax.numpy if any value is a JAX array/tracer, else numpy.

This is the hybrid-backend idiom for equation paths: bind xp = array_module(x) once at the top of f / h, then write the math with xp so the same code runs under NumPy and traces under JAX.

Return type:

ModuleType

minilink.core.backends.jax_installed()[source]

Return True when the jax package is importable (without importing it).

Return type:

bool

minilink.core.backends.require_jax()[source]

Return jax (lazy, cached) under the float64 policy. Raises if JAX is not installed.

Return type:

ModuleType

minilink.core.backends.require_jax_numpy()[source]

Return jax.numpy (lazy, cached). Raises if JAX is not installed.

Return type:

ModuleType

minilink.core.backends.jax_x64_policy()[source]

Library policy: 64-bit floats in JAX unless MINILINK_JAX_X64=0.

MINILINK_JAX_X64=0 keeps JAX’s float32 default, for GPU or reinforcement-learning workloads.

Return type:

bool

minilink.core.backends.ensure_jax_x64()[source]

Apply jax_x64_policy() to the JAX config (idempotent; JAX required).

Called by every JAX evaluator constructor and by require_jax_numpy(), so tools built on JAX evaluators never need the caller to enable x64. Without JAX it raises require_jax()’s error, with the install hint.

Return type:

ModuleType

minilink.core.backends.configure_jax(*, enable_x64=None)[source]

Configure process-wide JAX options used by Minilink JAX code paths.

Parameters:

enable_x64 (bool | None)

Return type:

ModuleType

minilink.core.backends.is_jax_key(key)[source]

True for a JAX PRNG key (concrete or traced).

Return type:

bool

minilink.core.backends.numpy_generator(key)[source]

A NumPy generator from a generator, an integer seed, or None.

Return type:

Generator

Compiled evaluator ABCs — output, dynamics, step, and static contracts.

Concrete backends live in numpy_evaluators and jax_evaluators.

minilink.core.compile.evaluators.evaluators.outputs_from_ports(system, x, u, t, params)[source]

Evaluate every output port on system.

Return type:

dict

class minilink.core.compile.evaluators.evaluators.OutputEvaluator[source]

Base class for compiled evaluators — output tier only.

Boundary signals are returned as a dict keyed by output port id.

abstract outputs(x, u, t=0.0)[source]

All boundary output ports, frozen params.

Return type:

dict

abstract outputs_p(x, u, t, params)[source]

All boundary output ports; caller-supplied params.

Return type:

dict

class minilink.core.compile.evaluators.evaluators.DynamicsEvaluator[source]

Base class for compiled continuous-time evaluators.

Subclasses implement f, f_p, outputs, and outputs_p. Integration helpers are mixed in by the NumPy / JAX backend modules.

abstract f(x, u, t=0.0)[source]

dx = f(x, u, t) with frozen params.

abstract f_p(x, u, t, params)[source]

dx = f(x, u, t, p) with caller-supplied params.

f_ivp(x, t=0.0)[source]

dx = f(x, u_nom, t) with nominal input fixed at compile.

f_ivp_p(x, t, params)[source]

dx = f(x, u_nom, t, p) — nominal u, runtime params.

class minilink.core.compile.evaluators.evaluators.StepEvaluator[source]

Base class for compiled discrete-time evaluators.

Subclasses implement step, step_p, outputs, and outputs_p. Rollout helpers live on the NumPy / JAX backend mixins.

abstract step(x, u, k=0)[source]

x_{k+1} = step(x, u, k) with frozen params.

abstract step_p(x, u, k, params)[source]

x_{k+1} = step(x, u, k, p) with caller-supplied params.

class minilink.core.compile.evaluators.evaluators.StaticEvaluator[source]

Compiled evaluator for stateless IO blocks — no evolution map.

Step rollouts and diagram input gathering for compiled evaluators.

Shared by NumPy and JAX step evaluators and by Computer.

minilink.core.compile.evaluators.step_rollout.gather_u(gather_sources, u_dim, signals, u)[source]

Assemble a subsystem local input vector from a compiled gather plan.

Each entry in gather_sources is (source_type, source_value, dim) as stored on PortOperation and StateOperation. Source types are INTERNAL_SIGNAL, NOMINAL, and EXTERNAL_INPUT. For INTERNAL_SIGNAL and EXTERNAL_INPUT, source_value is a slice into signals or the diagram boundary input u respectively.

Parameters:
  • gather_sources (tuple[tuple[int, object, int], ...]) – Flattened wiring plan for one subsystem input port bundle.

  • u_dim (int) – Expected local input dimension (sum of dim over sources).

  • signals (ndarray) – Internal signal buffer (1-D) or port-signal workspace used for INTERNAL_SIGNAL lookups.

  • u (ndarray) – Diagram boundary input vector.

Returns:

Local input u_local with shape (u_dim,).

Return type:

ndarray

class minilink.core.compile.evaluators.step_rollout.StepRolloutMixin[source]

Discrete-time rollouts for compiled StepSystem evaluators.

Subclasses must implement step, step_p, outputs, and outputs_p and expose n, m, p, and _u_nominal.

rollout(x0, *, n_steps, u=None)[source]

Roll out x_{k+1} = step(x_k, u_k, k) for n_steps transitions.

Returns a state-only StepRollout with k.shape == (n_steps + 1,), x.shape == (n, n_steps + 1), and u.shape == (m, n_steps + 1). Signal histories are owned by Computer and HybridSimulator.

Return type:

StepRollout

rollout_p(x0, *, n_steps, u=None, params)[source]

Parametric twin of rollout().

Return type:

StepRollout

NumPy compiled evaluators — dynamics, diagrams, step, and static.

Integration methods show textbook RK4 / Euler formulas inline.

class minilink.core.compile.evaluators.numpy_evaluators.IntegrationMixin[source]

RK4 / Euler integration on f(), f_p(), f_ivp(), f_ivp_p().

rk4_integrate_zoh(x0, u_sequence, t0, dt)[source]

ZOH RK4 rollout: u_sequence shape (N, m) → states (N+1, n).

rk4_integrate_linear(x0, u_knots, t0, dt)[source]

Linear-u RK4 rollout: u_knots shape (K, m) → states (K, n).

euler_integrate_zoh(x0, u_sequence, t0, dt)[source]

ZOH Euler rollout: u_sequence shape (N, m) → states (N+1, n).

integrate_zoh_rollout(x0, u_hold, t0, dt_hold, *, dt_inner=None)[source]

Piecewise-constant u_hold over [t0, t0 + dt_hold].

Returns:

  • t_samples (ndarray, shape (n_sub + 1,))

  • x_samples (ndarray, shape (n_sub + 1, n))

integrate_zoh(x0, u_hold, t0, dt_hold, *, dt_inner=None)[source]

Advance state with piecewise-constant u_hold; return final state.

integrate_zoh_p(x0, u_hold, t0, dt_hold, params, *, dt_inner=None)[source]

Parametric twin of integrate_zoh().

class minilink.core.compile.evaluators.numpy_evaluators.NumpyDynamicEvaluator(system)[source]

Compiled evaluator for a DynamicSystem using NumPy.

Parameters:

system (DynamicSystem)

f(x, u, t=0.0)[source]

dx = f(x, u, t) with frozen params.

f_p(x, u, t, params)[source]

dx = f(x, u, t, p) with caller-supplied params.

outputs(x, u, t=0.0)[source]

All boundary output ports, frozen params.

outputs_p(x, u, t, params)[source]

All boundary output ports; caller-supplied params.

class minilink.core.compile.evaluators.numpy_evaluators.NumpyDiagramEvaluator(plan, diagram)[source]

Stateless NumPy evaluator for a compiled diagram.

Parameters:

plan (ExecutionPlan)

f(x, u, t=0.0)[source]

dx = f(x, u, t) with frozen params.

Parameters:
  • x (ndarray)

  • u (ndarray)

  • t (float)

Return type:

ndarray

f_p(x, u, t, params)[source]

dx = f(x, u, t, p) with caller-supplied params.

outputs(x, u, t=0.0)[source]

All boundary output ports, frozen params.

outputs_p(x, u, t, params)[source]

All boundary output ports; caller-supplied params.

class minilink.core.compile.evaluators.numpy_evaluators.NumpyStaticEvaluator(system)[source]

NumPy evaluator for a static System leaf (n == 0).

Parameters:

system (System)

outputs(x, u, t=0.0)[source]

All boundary output ports, frozen params.

outputs_p(x, u, t, params)[source]

All boundary output ports; caller-supplied params.

class minilink.core.compile.evaluators.numpy_evaluators.NumpyStepEvaluator(system)[source]

Compiled evaluator for a StepSystem using NumPy.

Parameters:

system (StepSystem)

step(x, u, k=0)[source]

x_{k+1} = step(x, u, k) with frozen params.

step_p(x, u, k, params)[source]

x_{k+1} = step(x, u, k, p) with caller-supplied params.

outputs(x, u, k=0)[source]

All boundary output ports, frozen params.

outputs_p(x, u, k, params)[source]

All boundary output ports; caller-supplied params.

class minilink.core.compile.evaluators.numpy_evaluators.NumpyStepDiagramEvaluator(plan, diagram)[source]

Stateless NumPy evaluator for a compiled step diagram.

Parameters:

plan (StepExecutionPlan)

step(x, u, k=0)[source]

x_{k+1} = step(x, u, k) with frozen params.

step_p(x, u, k, params)[source]

x_{k+1} = step(x, u, k, p) with caller-supplied params.

outputs(x, u, k=0)[source]

All boundary output ports, frozen params.

outputs_p(x, u, k, params)[source]

All boundary output ports; caller-supplied params.

step_block(sys_id, x, u, k=0)[source]

Advance one step-diagram subsystem (subset-firing hook for Computer).

JAX compiled evaluators — dynamics, diagrams, step, and static.

Integration methods show textbook RK4 / Euler formulas inline.

class minilink.core.compile.evaluators.jax_evaluators.JaxIntegrationMixin[source]

RK4 / Euler integration on trace and JIT f tiers.

rk4_step_trace(x, u, t, dt)[source]

Pre-JIT RK4 step for JAX composition.

rk4_step_trace_p(x, u, t, dt, params)[source]

Pre-JIT parametric RK4 step for JAX composition.

euler_step_trace(x, u, t, dt)[source]

Pre-JIT Euler step for JAX composition.

euler_step_trace_p(x, u, t, dt, params)[source]

Pre-JIT parametric Euler step for JAX composition.

rk4_step_ivp_trace(x, t, dt)[source]

Pre-JIT IVP RK4 step for JAX composition.

rk4_step_ivp_trace_p(x, t, dt, params)[source]

Pre-JIT parametric IVP RK4 step for JAX composition.

euler_step_ivp_trace(x, t, dt)[source]

Pre-JIT IVP Euler step for JAX composition.

euler_step_ivp_trace_p(x, t, dt, params)[source]

Pre-JIT parametric IVP Euler step for JAX composition.

rk4_integrate_zoh_trace(x0, u_sequence, t0, dt)[source]

Pre-JIT ZOH-input rollout for JAX composition.

rk4_integrate_zoh_trace_p(x0, u_sequence, t0, dt, params)[source]

Pre-JIT parametric ZOH-input rollout for JAX composition.

rollout_batch(x0s, u_sequences=None, *, t0=0.0, dt, n_steps=None, params=None)[source]

Roll out a family of RK4 / ZOH trajectories in one jax.vmap call.

Parameters:
  • x0s (array, shape (B, n)) – One initial state per family member.

  • u_sequences (None, (N, m), or (B, N, m)) – Held inputs per step: the compiled nominal input (n_steps required), one sequence shared by the batch, or one per member.

  • t0 (float) – Start time and step, shared by the batch.

  • dt (float) – Start time and step, shared by the batch.

  • params (None or pytree) – None uses the compiled values. Otherwise a params pytree whose leaves are either shared by the batch or carry a leading axis of size B — a leaf with one more dimension than the compiled value is a family ({"l": jnp.linspace(0.5, 2.0, B)} sweeps a length).

Returns:

Sampled states, x0 first, for every member of the family.

Return type:

jax.Array, shape (B, N + 1, n)

rk4_integrate_linear_trace(x0, u_knots, t0, dt)[source]

Pre-JIT linear-u rollout for JAX composition.

rk4_integrate_linear_trace_p(x0, u_knots, t0, dt, params)[source]

Pre-JIT parametric linear-u rollout for JAX composition.

rk4_integrate_ivp_trace(x0, t0, dt, n_steps)[source]

Pre-JIT IVP rollout for JAX composition.

rk4_integrate_ivp_trace_p(x0, t0, dt, n_steps, params)[source]

Pre-JIT parametric IVP rollout for JAX composition.

euler_integrate_zoh_trace(x0, u_sequence, t0, dt)[source]

Pre-JIT ZOH Euler rollout for JAX composition.

euler_integrate_zoh_trace_p(x0, u_sequence, t0, dt, params)[source]

Pre-JIT parametric ZOH Euler rollout for JAX composition.

euler_integrate_ivp_trace(x0, t0, dt, n_steps)[source]

Pre-JIT IVP Euler rollout for JAX composition.

euler_integrate_ivp_trace_p(x0, t0, dt, n_steps, params)[source]

Pre-JIT parametric IVP Euler rollout for JAX composition.

integrate_zoh_rollout(x0, u_hold, t0, dt_hold, *, dt_inner=None)[source]

Piecewise-constant u_hold over [t0, t0 + dt_hold].

Returns:

  • t_samples (ndarray, shape (n_sub + 1,))

  • x_samples (ndarray, shape (n_sub + 1, n))

integrate_zoh(x0, u_hold, t0, dt_hold, *, dt_inner=None)[source]

Advance state with piecewise-constant u_hold; return final state.

integrate_zoh_p(x0, u_hold, t0, dt_hold, params, *, dt_inner=None)[source]

Parametric twin of integrate_zoh().

class minilink.core.compile.evaluators.jax_evaluators.JaxDynamicEvaluator(system, verbose=False)[source]

Compiled evaluator for a DynamicSystem using JAX.

Parameters:

system (DynamicSystem)

f(x, u, t=0.0)[source]

dx = f(x, u, t) with frozen params.

f_trace(x, u, t=0.0)[source]

Pre-JIT flat callable for JAX composition.

outputs(x, u, t=0.0)[source]

All boundary output ports, frozen params.

outputs_trace(x, u, t=0.0)[source]

Pre-JIT boundary outputs for JAX composition.

f_p(x, u, t, params)[source]

dx = f(x, u, t, p) with caller-supplied params.

f_trace_p(x, u, t, params)[source]

Pre-JIT parametric dynamics for JAX composition.

outputs_p(x, u, t, params)[source]

All boundary output ports; caller-supplied params.

outputs_trace_p(x, u, t, params)[source]

Pre-JIT parametric boundary outputs for JAX composition.

f_jit(x, u, t=0.0)

dx = f(x, u, t) with frozen params.

f_jit_p(x, u, t, params)

dx = f(x, u, t, p) with caller-supplied params.

outputs_jit(x, u, t=0.0)

All boundary output ports, frozen params.

outputs_jit_p(x, u, t, params)

All boundary output ports; caller-supplied params.

class minilink.core.compile.evaluators.jax_evaluators.JaxDiagramEvaluator(plan, diagram, verbose=False)[source]

JAX-compatible evaluator for a compiled diagram.

Parameters:

plan (ExecutionPlan)

f(x, u, t=0.0)[source]

dx = f(x, u, t) with frozen params.

f_trace(x, u, t=0.0)[source]

Pre-JIT flat callable for JAX composition.

outputs(x, u, t=0.0)[source]

All boundary output ports, frozen params.

outputs_trace(x, u, t=0.0)[source]

Pre-JIT boundary outputs for JAX composition.

f_p(x, u, t, params)[source]

dx = f(x, u, t, p) with caller-supplied params.

f_trace_p(x, u, t, params)[source]

Pre-JIT parametric dynamics for JAX composition.

outputs_p(x, u, t, params)[source]

All boundary output ports; caller-supplied params.

outputs_trace_p(x, u, t, params)[source]

Pre-JIT parametric boundary outputs for JAX composition.

f_jit(x, u, t=0.0)

dx = f(x, u, t) with frozen params.

f_jit_p(x, u, t, params)

dx = f(x, u, t, p) with caller-supplied params.

outputs_jit(x, u, t=0.0)

All boundary output ports, frozen params.

outputs_jit_p(x, u, t, params)

All boundary output ports; caller-supplied params.

class minilink.core.compile.evaluators.jax_evaluators.JaxStaticEvaluator(system, verbose=False)[source]

JAX evaluator for a static System leaf.

Parameters:

system (System)

outputs(x, u, t=0.0)[source]

All boundary output ports, frozen params.

outputs_trace(x, u, t=0.0)[source]

Pre-JIT boundary outputs for JAX composition.

outputs_p(x, u, t, params)[source]

All boundary output ports; caller-supplied params.

outputs_trace_p(x, u, t, params)[source]

Pre-JIT parametric boundary outputs for JAX composition.

outputs_jit(x, u, t=0.0)

All boundary output ports, frozen params.

outputs_jit_p(x, u, t, params)

All boundary output ports; caller-supplied params.

class minilink.core.compile.evaluators.jax_evaluators.JaxStepEvaluator(system, verbose=False)[source]

Compiled evaluator for a StepSystem using JAX.

Parameters:

system (StepSystem)

step(x, u, k=0)[source]

x_{k+1} = step(x, u, k) with frozen params.

step_trace(x, u, k=0)[source]

Pre-JIT flat step for JAX composition.

step_p(x, u, k, params)[source]

x_{k+1} = step(x, u, k, p) with caller-supplied params.

step_trace_p(x, u, k, params)[source]

Pre-JIT parametric step for JAX composition.

outputs(x, u, k=0)[source]

All boundary output ports, frozen params.

outputs_trace(x, u, k=0)[source]

Pre-JIT boundary outputs for JAX composition.

outputs_p(x, u, k, params)[source]

All boundary output ports; caller-supplied params.

outputs_trace_p(x, u, k, params)[source]

Pre-JIT parametric boundary outputs for JAX composition.

rollout(x0, *, n_steps, u=None)[source]

Roll out x_{k+1} = step(x_k, u_k, k) for n_steps transitions.

Returns a state-only StepRollout with k.shape == (n_steps + 1,), x.shape == (n, n_steps + 1), and u.shape == (m, n_steps + 1). Signal histories are owned by Computer and HybridSimulator.

outputs_jit(x, u, k=0)

All boundary output ports, frozen params.

outputs_jit_p(x, u, k, params)

All boundary output ports; caller-supplied params.

step_jit(x, u, k=0)

x_{k+1} = step(x, u, k) with frozen params.

step_jit_p(x, u, k, params)

x_{k+1} = step(x, u, k, p) with caller-supplied params.

class minilink.core.compile.evaluators.jax_evaluators.JaxStepDiagramEvaluator(plan, diagram, verbose=False)[source]

JAX evaluator for a compiled step diagram.

Parameters:

plan (StepExecutionPlan)

step(x, u, k=0)[source]

x_{k+1} = step(x, u, k) with frozen params.

step_trace(x, u, k=0)[source]

Pre-JIT flat step for JAX composition.

step_p(x, u, k, params)[source]

x_{k+1} = step(x, u, k, p) with caller-supplied params.

step_trace_p(x, u, k, params)[source]

Pre-JIT parametric step for JAX composition.

outputs(x, u, k=0)[source]

All boundary output ports, frozen params.

outputs_jit(x, u, k=0)

All boundary output ports, frozen params.

outputs_jit_p(x, u, k, params)

All boundary output ports; caller-supplied params.

outputs_trace(x, u, k=0)[source]

Pre-JIT boundary outputs for JAX composition.

step_jit(x, u, k=0)

x_{k+1} = step(x, u, k) with frozen params.

step_jit_p(x, u, k, params)

x_{k+1} = step(x, u, k, p) with caller-supplied params.

outputs_p(x, u, k, params)[source]

All boundary output ports; caller-supplied params.

outputs_trace_p(x, u, k, params)[source]

Pre-JIT parametric boundary outputs for JAX composition.

rollout(x0, *, n_steps, u=None)[source]

Roll out x_{k+1} = step(x_k, u_k, k) for n_steps transitions.

Returns a state-only StepRollout with k.shape == (n_steps + 1,), x.shape == (n, n_steps + 1), and u.shape == (m, n_steps + 1). Signal histories are owned by Computer and HybridSimulator.