Skip to content

Neural Network Module (nn)

GPU-first differentiable Kuramoto dynamics for neural network integration via JAX and equinox. Every function and layer is JIT-compilable, vmap-compatible, and fully differentiable — enabling gradient-based coupling inference, synchronisation optimisation, and physics-informed machine learning.

Requires: pip install scpn-phase-orchestrator[nn] (installs jax + equinox + optax)

Runtime API

Production ML jobs should make the accelerator contract explicit at start-up:

from scpn_phase_orchestrator.nn import (
    KuramotoLayer,
    jax_runtime_info,
    require_accelerator,
)

print(jax_runtime_info())
device = require_accelerator()

require_accelerator() raises when JAX is installed but only CPU devices are visible. Use require_accelerator(allow_cpu=True) only for CI, notebooks, and smoke tests that intentionally run without a GPU/TPU.

runtime

Runtime contract for the GPU-first differentiable nn API.

The numerical layers remain pure JAX/equinox modules. This module makes the runtime status explicit for ML users: whether JAX is installed, which backend is active, which devices are visible, and whether a non-CPU accelerator is available for production training paths.

Classes

JaxRuntimeInfo dataclass 

JaxRuntimeInfo(
    has_jax: bool,
    backend: str | None,
    devices: tuple[str, ...],
    default_device: str | None,
    device_count: int,
    accelerator_count: int,
)

Snapshot of the active JAX runtime visible to nn callers.

Attributes
has_accelerator property 
has_accelerator: bool

Return True when at least one non-CPU JAX device is visible.

Returns

bool Return True when at least one non-CPU JAX device is visible.

Functions:

require_jax 

require_jax() -> ModuleType

Return the imported JAX module or raise a clear installation error.

Returns

ModuleType Return the imported JAX module or raise a clear installation error.

Raises

RuntimeError If the JAX runtime requirement is not met.

Source code in src/scpn_phase_orchestrator/nn/runtime.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
def require_jax() -> ModuleType:
    """Return the imported JAX module or raise a clear installation error.

    Returns
    -------
    ModuleType
        Return the imported JAX module or raise a clear installation error.

    Raises
    ------
    RuntimeError
        If the JAX runtime requirement is not met.
    """
    if not HAS_JAX:
        msg = (
            "JAX is required for scpn_phase_orchestrator.nn. "
            "Install with: pip install scpn-phase-orchestrator[nn]"
        )
        raise RuntimeError(msg)

    import jax

    return jax

jax_runtime_info 

jax_runtime_info() -> JaxRuntimeInfo

Return an import-safe summary of the active JAX runtime.

The function never imports JAX when it is not installed, so base package users can still import scpn_phase_orchestrator.nn and receive an actionable runtime report.

Returns

JaxRuntimeInfo Return an import-safe summary of the active JAX runtime.

Source code in src/scpn_phase_orchestrator/nn/runtime.py 
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
def jax_runtime_info() -> JaxRuntimeInfo:
    """Return an import-safe summary of the active JAX runtime.

    The function never imports JAX when it is not installed, so base package
    users can still import ``scpn_phase_orchestrator.nn`` and receive an
    actionable runtime report.

    Returns
    -------
    JaxRuntimeInfo
        Return an import-safe summary of the active JAX runtime.
    """
    if not HAS_JAX:
        return JaxRuntimeInfo(
            has_jax=False,
            backend=None,
            devices=(),
            default_device=None,
            device_count=0,
            accelerator_count=0,
        )

    jax = require_jax()
    devices_raw = tuple(jax.devices())
    devices = tuple(_device_label(device) for device in devices_raw)
    accelerators = tuple(
        device for device in devices_raw if not _device_kind(device).startswith("cpu")
    )
    default_backend = str(jax.default_backend())
    default_device = devices[0] if devices else None
    return JaxRuntimeInfo(
        has_jax=True,
        backend=default_backend,
        devices=devices,
        default_device=default_device,
        device_count=len(devices),
        accelerator_count=len(accelerators),
    )

default_device 

default_device() -> str

Return the default JAX device label used by the nn API.

Returns

str Return the default JAX device label used by the nn API.

Raises

RuntimeError If the JAX runtime requirement is not met.

Source code in src/scpn_phase_orchestrator/nn/runtime.py 
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
def default_device() -> str:
    """Return the default JAX device label used by the ``nn`` API.

    Returns
    -------
    str
        Return the default JAX device label used by the ``nn`` API.

    Raises
    ------
    RuntimeError
        If the JAX runtime requirement is not met.
    """
    info = jax_runtime_info()
    if info.default_device is None:
        msg = (
            "JAX is required for scpn_phase_orchestrator.nn. "
            "Install with: pip install scpn-phase-orchestrator[nn]"
        )
        raise RuntimeError(msg)
    return info.default_device

require_accelerator 

require_accelerator(*, allow_cpu: bool = False) -> str

Return the production training device or fail fast on CPU-only runtimes.

Parameters

allow_cpu : bool Permit CPU-only JAX execution. This is useful for CI, documentation examples, and small smoke tests. Production ML training should keep the default False so misconfigured GPU jobs fail before expensive work starts.

Returns

str A JAX device label such as "gpu:0", "tpu:0", or "cpu:0" when allow_cpu is enabled.

Raises

RuntimeError If the JAX runtime requirement is not met.

Source code in src/scpn_phase_orchestrator/nn/runtime.py 
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
def require_accelerator(*, allow_cpu: bool = False) -> str:
    """Return the production training device or fail fast on CPU-only runtimes.

    Parameters
    ----------
    allow_cpu : bool
        Permit CPU-only JAX execution. This is useful for CI, documentation examples,
        and small smoke tests. Production ML training should keep the default ``False``
        so misconfigured GPU jobs fail before expensive work starts.

    Returns
    -------
    str
        A JAX device label such as ``"gpu:0"``, ``"tpu:0"``, or ``"cpu:0"`` when
        ``allow_cpu`` is enabled.

    Raises
    ------
    RuntimeError
        If the JAX runtime requirement is not met.
    """
    info = jax_runtime_info()
    if not info.has_jax:
        msg = (
            "JAX is required for scpn_phase_orchestrator.nn. "
            "Install with: pip install scpn-phase-orchestrator[nn]"
        )
        raise RuntimeError(msg)

    accelerator = next(
        (device for device in info.devices if not device.lower().startswith("cpu")),
        None,
    )
    if accelerator is not None:
        return accelerator
    if allow_cpu and info.default_device is not None:
        return info.default_device

    msg = (
        "No JAX GPU/TPU accelerator is visible. Install a hardware-enabled "
        "jaxlib build or configure the accelerator runtime before production "
        "training. For CI or smoke tests, call require_accelerator(allow_cpu=True)."
    )
    raise RuntimeError(msg)

Architecture

                        ┌─────────────────────────┐
                        │   Functional API (JAX)   │
                        │  kuramoto_forward()      │
                        │  winfree_forward()       │
                        │  simplicial_forward()    │
                        │  stuart_landau_forward() │
                        │  order_parameter()       │
                        └───────────┬─────────────┘
                                    │
                    ┌───────────────┼───────────────┐
                    ↓               ↓               ↓
            KuramotoLayer    SimplicialLayer   StuartLandauLayer
            (eqx.Module)    (eqx.Module)      (eqx.Module)
                    │               │               │
                    ↓               ↓               ↓
              training.py     UDEKuramotoLayer   BOLDGenerator
              (loss + optim)  (physics + MLP)    (hemodynamics)
                    │
          ┌─────────┼──────────┬───────────────┐
          ↓         ↓          ↓               ↓
     InverseKuramoto  Reservoir   OIM   DifferentiableSupervisor
     (coupling inference) (readout) (combinatorial) (closed-loop policy)

Functional API

Pure JAX functions — no state, no side effects. Each function is decorated with @jax.jit internally or designed to be JIT'd by the caller.

Kuramoto model

Function Signature
kuramoto_step (phases, omegas, K, dt) → phases
kuramoto_rk4_step (phases, omegas, K, dt) → phases
kuramoto_forward (phases, omegas, K, dt, n_steps, method="rk4") → (final, traj)

Masked (sparse) variants append _masked and take an additional mask: jax.Array parameter for selective coupling.

Winfree model

Function Signature
winfree_step (phases, omegas, K, dt) → phases
winfree_rk4_step (phases, omegas, K, dt) → phases
winfree_forward (phases, omegas, K, dt, n_steps, method="rk4") → (final, traj)

Winfree coupling: dθ_i/dt = ω_i + K · Q(θ_i) · Σ P(θ_j) where Q is the sensitivity function and P is the pulse function.

Use Winfree dynamics when the observed coupling is pulse-driven rather than a smooth sinusoidal pull. Typical cases include biological pacemakers, circadian or neural populations with event-like signalling, flashing or firing oscillator ensembles, endocrine or chemical pulse trains, and sensor networks where one unit perturbs another through brief impulses. The model is the right API choice when the phase response curve and pulse waveform are part of the hypothesis, not just incidental numerical details.

Simplicial (3-body) model

Function Signature
simplicial_step (phases, omegas, K, dt, sigma2=0.0) → phases
simplicial_rk4_step (phases, omegas, K, dt, sigma2=0.0) → phases
simplicial_forward (phases, omegas, K, dt, n_steps, sigma2=0.0, method="rk4") → (final, traj)

The sigma2 parameter controls the 3-body interaction strength. When sigma2=0, reduces to standard Kuramoto.

Use simplicial dynamics when pairwise edges are not enough to represent the interaction mechanism. The 3-body term models triadic or group constraints: neural assemblies whose synchrony depends on co-active triplets, social or multi-agent triads, reaction loops, power-network group modes, and topology extracted from simplicial complexes or hypergraphs. This surface is intended for regressions where cluster states, abrupt synchronization transitions, or learned coupling cannot be explained by pairwise Kuramoto alone.

Winfree and simplicial models answer different modelling questions. Winfree changes the timing law from smooth coupling to pulse-response coupling. Simplicial Kuramoto changes the interaction topology from pairwise edges to group terms. Combining them conceptually is useful for event-driven systems with group structure, for example spiking neural assemblies, biological tissue motifs, swarm coordination with triadic constraints, or industrial networks whose failure modes depend on triangular dependencies rather than isolated links.

Stuart-Landau model

Function Signature
stuart_landau_step (phases, amps, omegas, mu, K, K_r, dt, eps=1.0) → (phases, amps)
stuart_landau_rk4_step (phases, amps, omegas, mu, K, K_r, dt, eps=1.0) → (phases, amps)
stuart_landau_forward (phases, amps, omegas, mu, K, K_r, dt, n_steps, eps=1.0, method="rk4") → (phases, amps, phase_traj, amp_traj)

Analysis functions

Function Returns Description
order_parameter(phases) scalar Kuramoto R = |Σ exp(iθ)|/N
plv(trajectory) (N,N) array Pairwise phase-locking value
coupling_laplacian(K) (N,N) array Graph Laplacian L = D - K
saf_order_parameter(K, omegas, solver="auto") scalar Self-consistent analytical R
saf_loss(K, omegas, budget, solver="auto") scalar Differentiable SAF loss

Spectral Alignment Function use cases

SAF estimates the synchrony of a Kuramoto network directly from the coupling Laplacian and natural frequencies, without rolling out the ODE. Use it when the question is "which topology should synchronise this frequency field?" rather than "what is the phase trajectory at each time step?"

Concrete uses:

  • Coupling-topology optimisation under a wiring or energy budget.
  • Fast screening of candidate graphs before expensive time-domain simulation.
  • Differentiable regularisation for learned K matrices in neural pipelines.
  • Review of auto_initial_k matrices produced by auto-binding before runtime actuation.
  • Sensitivity analysis for which frequency modes are poorly aligned with the graph Laplacian.

Solver modes:

  • solver="eigh" computes the exact dense Laplacian eigendecomposition and is appropriate for small and medium dense systems where eigenvectors are needed for auditability.
  • solver="cg" uses the equivalent Laplacian pseudoinverse formulation and conjugate-gradient matrix-vector products. This avoids full eigendecomposition and is the preferred GPU path for large dense systems.
  • solver="auto" uses eigh up to exact_size_limit and switches to cg above that size.

Scaling limits: both paths still consume a dense (N, N) coupling matrix. The CG path removes the cubic eigensolver bottleneck, but it does not make dense memory disappear. For very sparse networks, keep a sparse or masked coupling representation upstream and materialise dense K only when the SAF audit size fits device memory.

Boundary contract: K must be a square real-valued JAX-compatible coupling matrix and omegas must be a real-valued one-dimensional vector with matching length. Boolean and complex payloads are rejected before Laplacian construction. SAF solver controls are finite positive scalars or integers, while saf_loss budget controls are finite non-negative scalars.

functional

Pure JAX functions for differentiable Kuramoto dynamics.

All functions are JIT-compilable, vmap-compatible, and differentiable via JAX autodiff. No NumPy conversions — inputs and outputs stay as JAX arrays for gradient flow.

Requires: jax>=0.4

Functions:

kuramoto_step 

kuramoto_step(
    phases: Array, omegas: Array, K: Array, dt: float
) -> jax.Array

Single Euler step of the Kuramoto model.

Parameters

phases : jax.Array (N,) oscillator phases in [0, 2pi). omegas : jax.Array (N,) natural frequencies. K : jax.Array (N, N) coupling matrix. dt : float integration timestep.

Returns

jax.Array (N,) updated phases, wrapped to [0, 2pi).

Source code in src/scpn_phase_orchestrator/nn/functional.py 
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
def kuramoto_step(
    phases: jax.Array,
    omegas: jax.Array,
    K: jax.Array,
    dt: float,
) -> jax.Array:
    """Single Euler step of the Kuramoto model.

    Parameters
    ----------
    phases : jax.Array
        (N,) oscillator phases in [0, 2pi).
    omegas : jax.Array
        (N,) natural frequencies.
    K : jax.Array
        (N, N) coupling matrix.
    dt : float
        integration timestep.

    Returns
    -------
    jax.Array
        (N,) updated phases, wrapped to [0, 2pi).
    """
    diff = phases[jnp.newaxis, :] - phases[:, jnp.newaxis]
    coupling = jnp.sum(K * jnp.sin(diff), axis=1)
    return (phases + dt * (omegas + coupling)) % TWO_PI

kuramoto_rk4_step 

kuramoto_rk4_step(
    phases: Array, omegas: Array, K: Array, dt: float
) -> jax.Array

Single RK4 step of the Kuramoto model.

Parameters

phases : jax.Array (N,) oscillator phases in [0, 2pi). omegas : jax.Array (N,) natural frequencies. K : jax.Array (N, N) coupling matrix. dt : float integration timestep.

Returns

jax.Array (N,) updated phases, wrapped to [0, 2pi).

Source code in src/scpn_phase_orchestrator/nn/functional.py 
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
def kuramoto_rk4_step(
    phases: jax.Array,
    omegas: jax.Array,
    K: jax.Array,
    dt: float,
) -> jax.Array:
    """Single RK4 step of the Kuramoto model.

    Parameters
    ----------
    phases : jax.Array
        (N,) oscillator phases in [0, 2pi).
    omegas : jax.Array
        (N,) natural frequencies.
    K : jax.Array
        (N, N) coupling matrix.
    dt : float
        integration timestep.

    Returns
    -------
    jax.Array
        (N,) updated phases, wrapped to [0, 2pi).
    """

    def deriv(p: jax.Array) -> jax.Array:
        """Evaluate the ODE right-hand side at a state."""
        diff = p[jnp.newaxis, :] - p[:, jnp.newaxis]
        return omegas + jnp.sum(K * jnp.sin(diff), axis=1)

    k1 = deriv(phases)
    k2 = deriv(phases + 0.5 * dt * k1)
    k3 = deriv(phases + 0.5 * dt * k2)
    k4 = deriv(phases + dt * k3)
    new = phases + (dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4)
    return new % TWO_PI

kuramoto_forward 

kuramoto_forward(
    phases: Array,
    omegas: Array,
    K: Array,
    dt: float,
    n_steps: int,
    method: str = "rk4",
) -> tuple[jax.Array, jax.Array]

Run N Kuramoto steps, returning final state and trajectory.

Uses jax.lax.scan for efficient compilation and autodiff.

Parameters

phases : jax.Array (N,) initial oscillator phases. omegas : jax.Array (N,) natural frequencies. K : jax.Array (N, N) coupling matrix. dt : float integration timestep. n_steps : int number of integration steps. method : str "rk4" or "euler".

Returns

tuple[jax.Array, jax.Array] : final: (N,) phases after n_steps trajectory: (n_steps, N) full phase trajectory.

Source code in src/scpn_phase_orchestrator/nn/functional.py 
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
def kuramoto_forward(
    phases: jax.Array,
    omegas: jax.Array,
    K: jax.Array,
    dt: float,
    n_steps: int,
    method: str = "rk4",
) -> tuple[jax.Array, jax.Array]:
    """Run N Kuramoto steps, returning final state and trajectory.

    Uses jax.lax.scan for efficient compilation and autodiff.

    Parameters
    ----------
    phases : jax.Array
        (N,) initial oscillator phases.
    omegas : jax.Array
        (N,) natural frequencies.
    K : jax.Array
        (N, N) coupling matrix.
    dt : float
        integration timestep.
    n_steps : int
        number of integration steps.
    method : str
        "rk4" or "euler".

    Returns
    -------
    tuple[jax.Array, jax.Array]
        : final: (N,) phases after n_steps trajectory: (n_steps, N) full phase
        trajectory.
    """
    step_fn = kuramoto_rk4_step if method == "rk4" else kuramoto_step

    def body(carry: jax.Array, _: None) -> tuple[jax.Array, jax.Array]:
        """Return the loop body for the iteration."""
        p = step_fn(carry, omegas, K, dt)
        return p, p

    final, trajectory = jax.lax.scan(body, phases, None, length=n_steps)
    return final, trajectory

kuramoto_step_masked 

kuramoto_step_masked(
    phases: Array,
    omegas: Array,
    K: Array,
    mask: Array,
    dt: float,
) -> jax.Array

Single Euler step with masked coupling.

Parameters

phases : jax.Array (N,) oscillator phases in [0, 2pi). omegas : jax.Array (N,) natural frequencies. K : jax.Array (N, N) coupling weights. mask : jax.Array (N, N) binary mask (1 = edge exists, 0 = no edge). dt : float integration timestep.

Returns

jax.Array (N,) updated phases.

Source code in src/scpn_phase_orchestrator/nn/functional.py 
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
def kuramoto_step_masked(
    phases: jax.Array,
    omegas: jax.Array,
    K: jax.Array,
    mask: jax.Array,
    dt: float,
) -> jax.Array:
    """Single Euler step with masked coupling.

    Parameters
    ----------
    phases : jax.Array
        (N,) oscillator phases in [0, 2pi).
    omegas : jax.Array
        (N,) natural frequencies.
    K : jax.Array
        (N, N) coupling weights.
    mask : jax.Array
        (N, N) binary mask (1 = edge exists, 0 = no edge).
    dt : float
        integration timestep.

    Returns
    -------
    jax.Array
        (N,) updated phases.
    """
    dphi = _kuramoto_deriv_masked(phases, omegas, K, mask)
    return (phases + dt * dphi) % TWO_PI

kuramoto_rk4_step_masked 

kuramoto_rk4_step_masked(
    phases: Array,
    omegas: Array,
    K: Array,
    mask: Array,
    dt: float,
) -> jax.Array

Single RK4 step with masked coupling.

Parameters

phases : jax.Array Oscillator phases in radians, shape (N,). omegas : jax.Array Natural frequencies in rad/s, shape (N,). K : jax.Array Coupling matrix K, shape (N, N). mask : jax.Array Boolean coupling mask, shape (N, N). dt : float Integration step size.

Returns

jax.Array The phases after one masked RK4 step.

Source code in src/scpn_phase_orchestrator/nn/functional.py 
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
def kuramoto_rk4_step_masked(
    phases: jax.Array,
    omegas: jax.Array,
    K: jax.Array,
    mask: jax.Array,
    dt: float,
) -> jax.Array:
    """Single RK4 step with masked coupling.

    Parameters
    ----------
    phases : jax.Array
        Oscillator phases in radians, shape ``(N,)``.
    omegas : jax.Array
        Natural frequencies in rad/s, shape ``(N,)``.
    K : jax.Array
        Coupling matrix ``K``, shape ``(N, N)``.
    mask : jax.Array
        Boolean coupling mask, shape ``(N, N)``.
    dt : float
        Integration step size.

    Returns
    -------
    jax.Array
        The phases after one masked RK4 step.
    """

    def deriv(p: jax.Array) -> jax.Array:
        """Evaluate the ODE right-hand side at a state."""
        return _kuramoto_deriv_masked(p, omegas, K, mask)

    k1 = deriv(phases)
    k2 = deriv(phases + 0.5 * dt * k1)
    k3 = deriv(phases + 0.5 * dt * k2)
    k4 = deriv(phases + dt * k3)
    return (phases + (dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4)) % TWO_PI

kuramoto_forward_masked 

kuramoto_forward_masked(
    phases: Array,
    omegas: Array,
    K: Array,
    mask: Array,
    dt: float,
    n_steps: int,
    method: str = "rk4",
) -> tuple[jax.Array, jax.Array]

Run N Kuramoto steps with masked coupling.

Parameters

phases : jax.Array (N,) initial phases. omegas : jax.Array (N,) natural frequencies. K : jax.Array (N, N) coupling weights. mask : jax.Array (N, N) binary mask. dt : float timestep. n_steps : int integration steps. method : str "rk4" or "euler".

Returns

tuple[jax.Array, jax.Array] (final, trajectory) — same as kuramoto_forward.

Source code in src/scpn_phase_orchestrator/nn/functional.py 
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
def kuramoto_forward_masked(
    phases: jax.Array,
    omegas: jax.Array,
    K: jax.Array,
    mask: jax.Array,
    dt: float,
    n_steps: int,
    method: str = "rk4",
) -> tuple[jax.Array, jax.Array]:
    """Run N Kuramoto steps with masked coupling.

    Parameters
    ----------
    phases : jax.Array
        (N,) initial phases.
    omegas : jax.Array
        (N,) natural frequencies.
    K : jax.Array
        (N, N) coupling weights.
    mask : jax.Array
        (N, N) binary mask.
    dt : float
        timestep.
    n_steps : int
        integration steps.
    method : str
        "rk4" or "euler".

    Returns
    -------
    tuple[jax.Array, jax.Array]
        (final, trajectory) — same as kuramoto_forward.
    """
    step_fn = kuramoto_rk4_step_masked if method == "rk4" else kuramoto_step_masked

    def body(carry: jax.Array, _: None) -> tuple[jax.Array, jax.Array]:
        """Return the loop body for the iteration."""
        p = step_fn(carry, omegas, K, mask, dt)
        return p, p

    final, trajectory = jax.lax.scan(body, phases, None, length=n_steps)
    return final, trajectory

winfree_step 

winfree_step(
    phases: Array, omegas: Array, K: float, dt: float
) -> jax.Array

Single Euler step of the Winfree model.

Parameters

phases : jax.Array (N,) oscillator phases in [0, 2pi). omegas : jax.Array (N,) natural frequencies. K : float scalar coupling strength. dt : float integration timestep.

Returns

jax.Array (N,) updated phases.

Source code in src/scpn_phase_orchestrator/nn/functional.py 
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
def winfree_step(
    phases: jax.Array,
    omegas: jax.Array,
    K: float,
    dt: float,
) -> jax.Array:
    """Single Euler step of the Winfree model.

    Parameters
    ----------
    phases : jax.Array
        (N,) oscillator phases in [0, 2pi).
    omegas : jax.Array
        (N,) natural frequencies.
    K : float
        scalar coupling strength.
    dt : float
        integration timestep.

    Returns
    -------
    jax.Array
        (N,) updated phases.
    """
    return (phases + dt * _winfree_deriv(phases, omegas, K)) % TWO_PI

winfree_rk4_step 

winfree_rk4_step(
    phases: Array, omegas: Array, K: float, dt: float
) -> jax.Array

Single RK4 step of the Winfree model.

Parameters

phases : jax.Array Oscillator phases in radians, shape (N,). omegas : jax.Array Natural frequencies in rad/s, shape (N,). K : float Coupling matrix K, shape (N, N). dt : float Integration step size.

Returns

jax.Array The phases after one Winfree RK4 step.

Source code in src/scpn_phase_orchestrator/nn/functional.py 
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
def winfree_rk4_step(
    phases: jax.Array,
    omegas: jax.Array,
    K: float,
    dt: float,
) -> jax.Array:
    """Single RK4 step of the Winfree model.

    Parameters
    ----------
    phases : jax.Array
        Oscillator phases in radians, shape ``(N,)``.
    omegas : jax.Array
        Natural frequencies in rad/s, shape ``(N,)``.
    K : float
        Coupling matrix ``K``, shape ``(N, N)``.
    dt : float
        Integration step size.

    Returns
    -------
    jax.Array
        The phases after one Winfree RK4 step.
    """

    def deriv(p: jax.Array) -> jax.Array:
        """Evaluate the ODE right-hand side at a state."""
        return _winfree_deriv(p, omegas, K)

    k1 = deriv(phases)
    k2 = deriv(phases + 0.5 * dt * k1)
    k3 = deriv(phases + 0.5 * dt * k2)
    k4 = deriv(phases + dt * k3)
    return (phases + (dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4)) % TWO_PI

winfree_forward 

winfree_forward(
    phases: Array,
    omegas: Array,
    K: float,
    dt: float,
    n_steps: int,
    method: str = "rk4",
) -> tuple[jax.Array, jax.Array]

Run N steps of Winfree dynamics.

Parameters

phases : jax.Array (N,) initial phases. omegas : jax.Array (N,) natural frequencies. K : float scalar coupling strength. dt : float timestep. n_steps : int integration steps. method : str "rk4" or "euler".

Returns

tuple[jax.Array, jax.Array] (final, trajectory).

Source code in src/scpn_phase_orchestrator/nn/functional.py 
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
def winfree_forward(
    phases: jax.Array,
    omegas: jax.Array,
    K: float,
    dt: float,
    n_steps: int,
    method: str = "rk4",
) -> tuple[jax.Array, jax.Array]:
    """Run N steps of Winfree dynamics.

    Parameters
    ----------
    phases : jax.Array
        (N,) initial phases.
    omegas : jax.Array
        (N,) natural frequencies.
    K : float
        scalar coupling strength.
    dt : float
        timestep.
    n_steps : int
        integration steps.
    method : str
        "rk4" or "euler".

    Returns
    -------
    tuple[jax.Array, jax.Array]
        (final, trajectory).
    """
    step_fn = winfree_rk4_step if method == "rk4" else winfree_step

    def body(carry: jax.Array, _: None) -> tuple[jax.Array, jax.Array]:
        """Return the loop body for the iteration."""
        p = step_fn(carry, omegas, K, dt)
        return p, p

    final, trajectory = jax.lax.scan(body, phases, None, length=n_steps)
    return final, trajectory

simplicial_step 

simplicial_step(
    phases: Array,
    omegas: Array,
    K: Array,
    dt: float,
    sigma2: float | Array = 0.0,
) -> jax.Array

Single Euler step of the simplicial (3-body) Kuramoto model.

Extends standard Kuramoto with higher-order 3-body interactions that produce explosive (first-order) synchronization transitions.

Parameters

phases : jax.Array (N,) oscillator phases in [0, 2pi). omegas : jax.Array (N,) natural frequencies. K : jax.Array (N, N) pairwise coupling matrix. dt : float integration timestep. sigma2 : float | jax.Array 3-body coupling strength (0 = standard Kuramoto).

Returns

jax.Array (N,) updated phases, wrapped to [0, 2pi).

Source code in src/scpn_phase_orchestrator/nn/functional.py 
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
def simplicial_step(
    phases: jax.Array,
    omegas: jax.Array,
    K: jax.Array,
    dt: float,
    sigma2: float | jax.Array = 0.0,
) -> jax.Array:
    """Single Euler step of the simplicial (3-body) Kuramoto model.

    Extends standard Kuramoto with higher-order 3-body interactions that
    produce explosive (first-order) synchronization transitions.

    Parameters
    ----------
    phases : jax.Array
        (N,) oscillator phases in [0, 2pi).
    omegas : jax.Array
        (N,) natural frequencies.
    K : jax.Array
        (N, N) pairwise coupling matrix.
    dt : float
        integration timestep.
    sigma2 : float | jax.Array
        3-body coupling strength (0 = standard Kuramoto).

    Returns
    -------
    jax.Array
        (N,) updated phases, wrapped to [0, 2pi).
    """
    dphi = _simplicial_deriv(phases, omegas, K, sigma2)
    return (phases + dt * dphi) % TWO_PI

simplicial_rk4_step 

simplicial_rk4_step(
    phases: Array,
    omegas: Array,
    K: Array,
    dt: float,
    sigma2: float | Array = 0.0,
) -> jax.Array

Single RK4 step of the simplicial (3-body) Kuramoto model.

Parameters

phases : jax.Array (N,) oscillator phases in [0, 2pi). omegas : jax.Array (N,) natural frequencies. K : jax.Array (N, N) pairwise coupling matrix. dt : float integration timestep. sigma2 : float | jax.Array 3-body coupling strength (0 = standard Kuramoto).

Returns

jax.Array (N,) updated phases, wrapped to [0, 2pi).

Source code in src/scpn_phase_orchestrator/nn/functional.py 
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
def simplicial_rk4_step(
    phases: jax.Array,
    omegas: jax.Array,
    K: jax.Array,
    dt: float,
    sigma2: float | jax.Array = 0.0,
) -> jax.Array:
    """Single RK4 step of the simplicial (3-body) Kuramoto model.

    Parameters
    ----------
    phases : jax.Array
        (N,) oscillator phases in [0, 2pi).
    omegas : jax.Array
        (N,) natural frequencies.
    K : jax.Array
        (N, N) pairwise coupling matrix.
    dt : float
        integration timestep.
    sigma2 : float | jax.Array
        3-body coupling strength (0 = standard Kuramoto).

    Returns
    -------
    jax.Array
        (N,) updated phases, wrapped to [0, 2pi).
    """

    def deriv(p: jax.Array) -> jax.Array:
        """Evaluate the ODE right-hand side at a state."""
        return _simplicial_deriv(p, omegas, K, sigma2)

    k1 = deriv(phases)
    k2 = deriv(phases + 0.5 * dt * k1)
    k3 = deriv(phases + 0.5 * dt * k2)
    k4 = deriv(phases + dt * k3)
    new = phases + (dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4)
    return new % TWO_PI

simplicial_forward 

simplicial_forward(
    phases: Array,
    omegas: Array,
    K: Array,
    dt: float,
    n_steps: int,
    sigma2: float | Array = 0.0,
    method: str = "rk4",
) -> tuple[jax.Array, jax.Array]

Run N steps of simplicial Kuramoto, returning final state and trajectory.

Parameters

phases : jax.Array (N,) initial oscillator phases. omegas : jax.Array (N,) natural frequencies. K : jax.Array (N, N) pairwise coupling matrix. dt : float integration timestep. n_steps : int number of integration steps. sigma2 : float | jax.Array 3-body coupling strength (0 = standard Kuramoto). method : str "rk4" or "euler".

Returns

tuple[jax.Array, jax.Array] (final, trajectory) where trajectory is (n_steps, N).

Source code in src/scpn_phase_orchestrator/nn/functional.py 
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
def simplicial_forward(
    phases: jax.Array,
    omegas: jax.Array,
    K: jax.Array,
    dt: float,
    n_steps: int,
    sigma2: float | jax.Array = 0.0,
    method: str = "rk4",
) -> tuple[jax.Array, jax.Array]:
    """Run N steps of simplicial Kuramoto, returning final state and trajectory.

    Parameters
    ----------
    phases : jax.Array
        (N,) initial oscillator phases.
    omegas : jax.Array
        (N,) natural frequencies.
    K : jax.Array
        (N, N) pairwise coupling matrix.
    dt : float
        integration timestep.
    n_steps : int
        number of integration steps.
    sigma2 : float | jax.Array
        3-body coupling strength (0 = standard Kuramoto).
    method : str
        "rk4" or "euler".

    Returns
    -------
    tuple[jax.Array, jax.Array]
        (final, trajectory) where trajectory is (n_steps, N).
    """
    step_fn = simplicial_rk4_step if method == "rk4" else simplicial_step

    def body(carry: jax.Array, _: None) -> tuple[jax.Array, jax.Array]:
        """Return the loop body for the iteration."""
        p = step_fn(carry, omegas, K, dt, sigma2)
        return p, p

    final, trajectory = jax.lax.scan(body, phases, None, length=n_steps)
    return final, trajectory

stuart_landau_step 

stuart_landau_step(
    phases: Array,
    amplitudes: Array,
    omegas: Array,
    mu: Array,
    K: Array,
    K_r: Array,
    dt: float,
    epsilon: float = 1.0,
) -> tuple[jax.Array, jax.Array]

Single Euler step of the Stuart-Landau oscillator model.

Parameters

phases : jax.Array (N,) oscillator phases in [0, 2pi). amplitudes : jax.Array (N,) oscillator amplitudes (r >= 0). omegas : jax.Array (N,) natural frequencies. mu : jax.Array (N,) bifurcation parameters (supercritical if mu > 0). K : jax.Array (N, N) phase coupling matrix. K_r : jax.Array (N, N) amplitude coupling matrix. dt : float integration timestep. epsilon : float amplitude coupling strength.

Returns

tuple[jax.Array, jax.Array] (new_phases, new_amplitudes).

Source code in src/scpn_phase_orchestrator/nn/functional.py 
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
def stuart_landau_step(
    phases: jax.Array,
    amplitudes: jax.Array,
    omegas: jax.Array,
    mu: jax.Array,
    K: jax.Array,
    K_r: jax.Array,
    dt: float,
    epsilon: float = 1.0,
) -> tuple[jax.Array, jax.Array]:
    """Single Euler step of the Stuart-Landau oscillator model.

    Parameters
    ----------
    phases : jax.Array
        (N,) oscillator phases in [0, 2pi).
    amplitudes : jax.Array
        (N,) oscillator amplitudes (r >= 0).
    omegas : jax.Array
        (N,) natural frequencies.
    mu : jax.Array
        (N,) bifurcation parameters (supercritical if mu > 0).
    K : jax.Array
        (N, N) phase coupling matrix.
    K_r : jax.Array
        (N, N) amplitude coupling matrix.
    dt : float
        integration timestep.
    epsilon : float
        amplitude coupling strength.

    Returns
    -------
    tuple[jax.Array, jax.Array]
        (new_phases, new_amplitudes).
    """
    dtheta, dr = _stuart_landau_deriv(phases, amplitudes, omegas, mu, K, K_r, epsilon)
    new_phases = (phases + dt * dtheta) % TWO_PI
    new_amplitudes = jnp.maximum(amplitudes + dt * dr, 0.0)
    return new_phases, new_amplitudes

stuart_landau_rk4_step 

stuart_landau_rk4_step(
    phases: Array,
    amplitudes: Array,
    omegas: Array,
    mu: Array,
    K: Array,
    K_r: Array,
    dt: float,
    epsilon: float = 1.0,
) -> tuple[jax.Array, jax.Array]

Single RK4 step of the Stuart-Landau oscillator model.

Parameters

phases : jax.Array (N,) oscillator phases in [0, 2pi). amplitudes : jax.Array (N,) oscillator amplitudes (r >= 0). omegas : jax.Array (N,) natural frequencies. mu : jax.Array (N,) bifurcation parameters. K : jax.Array (N, N) phase coupling matrix. K_r : jax.Array (N, N) amplitude coupling matrix. dt : float integration timestep. epsilon : float amplitude coupling strength.

Returns

tuple[jax.Array, jax.Array] (new_phases, new_amplitudes).

Source code in src/scpn_phase_orchestrator/nn/functional.py 
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
def stuart_landau_rk4_step(
    phases: jax.Array,
    amplitudes: jax.Array,
    omegas: jax.Array,
    mu: jax.Array,
    K: jax.Array,
    K_r: jax.Array,
    dt: float,
    epsilon: float = 1.0,
) -> tuple[jax.Array, jax.Array]:
    """Single RK4 step of the Stuart-Landau oscillator model.

    Parameters
    ----------
    phases : jax.Array
        (N,) oscillator phases in [0, 2pi).
    amplitudes : jax.Array
        (N,) oscillator amplitudes (r >= 0).
    omegas : jax.Array
        (N,) natural frequencies.
    mu : jax.Array
        (N,) bifurcation parameters.
    K : jax.Array
        (N, N) phase coupling matrix.
    K_r : jax.Array
        (N, N) amplitude coupling matrix.
    dt : float
        integration timestep.
    epsilon : float
        amplitude coupling strength.

    Returns
    -------
    tuple[jax.Array, jax.Array]
        (new_phases, new_amplitudes).
    """

    def deriv(p: jax.Array, r: jax.Array) -> tuple[jax.Array, jax.Array]:
        """Evaluate the ODE right-hand side at a state."""
        return _stuart_landau_deriv(p, r, omegas, mu, K, K_r, epsilon)

    k1p, k1r = deriv(phases, amplitudes)
    k2p, k2r = deriv(phases + 0.5 * dt * k1p, amplitudes + 0.5 * dt * k1r)
    k3p, k3r = deriv(phases + 0.5 * dt * k2p, amplitudes + 0.5 * dt * k2r)
    k4p, k4r = deriv(phases + dt * k3p, amplitudes + dt * k3r)

    new_phases = (phases + (dt / 6.0) * (k1p + 2 * k2p + 2 * k3p + k4p)) % TWO_PI
    new_amps = amplitudes + (dt / 6.0) * (k1r + 2 * k2r + 2 * k3r + k4r)
    return new_phases, jnp.maximum(new_amps, 0.0)

stuart_landau_forward 

stuart_landau_forward(
    phases: Array,
    amplitudes: Array,
    omegas: Array,
    mu: Array,
    K: Array,
    K_r: Array,
    dt: float,
    n_steps: int,
    epsilon: float = 1.0,
    method: str = "rk4",
) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]

Run N Stuart-Landau steps, returning final state and trajectories.

Parameters

phases : jax.Array (N,) initial phases. amplitudes : jax.Array (N,) initial amplitudes. omegas : jax.Array (N,) natural frequencies. mu : jax.Array (N,) bifurcation parameters. K : jax.Array (N, N) phase coupling matrix. K_r : jax.Array (N, N) amplitude coupling matrix. dt : float integration timestep. n_steps : int number of steps. epsilon : float amplitude coupling strength. method : str "rk4" or "euler".

Returns

tuple[jax.Array, jax.Array, jax.Array, jax.Array] (final_phases, final_amplitudes, phase_traj, amp_traj) where trajectories are (n_steps, N).

Source code in src/scpn_phase_orchestrator/nn/functional.py 
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
def stuart_landau_forward(
    phases: jax.Array,
    amplitudes: jax.Array,
    omegas: jax.Array,
    mu: jax.Array,
    K: jax.Array,
    K_r: jax.Array,
    dt: float,
    n_steps: int,
    epsilon: float = 1.0,
    method: str = "rk4",
) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]:
    """Run N Stuart-Landau steps, returning final state and trajectories.

    Parameters
    ----------
    phases : jax.Array
        (N,) initial phases.
    amplitudes : jax.Array
        (N,) initial amplitudes.
    omegas : jax.Array
        (N,) natural frequencies.
    mu : jax.Array
        (N,) bifurcation parameters.
    K : jax.Array
        (N, N) phase coupling matrix.
    K_r : jax.Array
        (N, N) amplitude coupling matrix.
    dt : float
        integration timestep.
    n_steps : int
        number of steps.
    epsilon : float
        amplitude coupling strength.
    method : str
        "rk4" or "euler".

    Returns
    -------
    tuple[jax.Array, jax.Array, jax.Array, jax.Array]
        (final_phases, final_amplitudes, phase_traj, amp_traj) where trajectories are
        (n_steps, N).
    """
    step_fn = stuart_landau_rk4_step if method == "rk4" else stuart_landau_step

    def body(
        carry: tuple[jax.Array, jax.Array], _: None
    ) -> tuple[tuple[jax.Array, jax.Array], tuple[jax.Array, jax.Array]]:
        """Return the loop body for the iteration."""
        p, r = carry
        new_p, new_r = step_fn(p, r, omegas, mu, K, K_r, dt, epsilon)
        return (new_p, new_r), (new_p, new_r)

    (final_p, final_r), (traj_p, traj_r) = jax.lax.scan(
        body, (phases, amplitudes), None, length=n_steps
    )
    return final_p, final_r, traj_p, traj_r

order_parameter 

order_parameter(phases: Array) -> jax.Array

Kuramoto order parameter R = ||.

Differentiable scalar measuring global synchronization. R=1 means perfect sync, R~0 means incoherent.

Parameters

phases : jax.Array (N,) or (T, N) oscillator phases.

Returns

jax.Array Scalar R value (or (T,) if trajectory input).

Source code in src/scpn_phase_orchestrator/nn/functional.py 
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
def order_parameter(phases: jax.Array) -> jax.Array:
    """Kuramoto order parameter R = |<exp(i*phi)>|.

    Differentiable scalar measuring global synchronization.
    R=1 means perfect sync, R~0 means incoherent.

    Parameters
    ----------
    phases : jax.Array
        (N,) or (T, N) oscillator phases.

    Returns
    -------
    jax.Array
        Scalar R value (or (T,) if trajectory input).
    """
    z = jnp.exp(1j * phases)
    return jnp.abs(jnp.mean(z, axis=-1))

plv 

plv(trajectory: Array) -> jax.Array

Phase-Locking Value matrix from a phase trajectory.

PLV_ij = |_t|

Parameters

trajectory : jax.Array (T, N) phase trajectory.

Returns

jax.Array (N, N) PLV matrix, values in [0, 1].

Source code in src/scpn_phase_orchestrator/nn/functional.py 
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
def plv(trajectory: jax.Array) -> jax.Array:
    """Phase-Locking Value matrix from a phase trajectory.

    PLV_ij = |<exp(i*(phi_i(t) - phi_j(t)))>_t|

    Parameters
    ----------
    trajectory : jax.Array
        (T, N) phase trajectory.

    Returns
    -------
    jax.Array
        (N, N) PLV matrix, values in [0, 1].
    """
    # (T, N, 1) - (T, 1, N) -> (T, N, N) phase differences
    diff = trajectory[:, :, jnp.newaxis] - trajectory[:, jnp.newaxis, :]
    return jnp.abs(jnp.mean(jnp.exp(1j * diff), axis=0))

coupling_laplacian 

coupling_laplacian(K: Array) -> jax.Array

Compute the graph Laplacian from a coupling matrix.

L = D - K, where D_ii = sum_j K_ij.

Parameters

K : jax.Array (N, N) symmetric coupling matrix.

Returns

jax.Array (N, N) Laplacian matrix.

Source code in src/scpn_phase_orchestrator/nn/functional.py 
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
def coupling_laplacian(K: jax.Array) -> jax.Array:
    """Compute the graph Laplacian from a coupling matrix.

    L = D - K, where D_ii = sum_j K_ij.

    Parameters
    ----------
    K : jax.Array
        (N, N) symmetric coupling matrix.

    Returns
    -------
    jax.Array
        (N, N) Laplacian matrix.
    """
    K = _validate_square_coupling("K", K)
    D = jnp.diag(jnp.sum(K, axis=1))
    return D - K

saf_order_parameter 

saf_order_parameter(
    K: Array,
    omegas: Array,
    eps: float = 1e-08,
    solver: str = "auto",
    exact_size_limit: int = 256,
    cg_tol: float = 1e-05,
    cg_maxiter: int | None = None,
) -> jax.Array

Spectral Alignment Function: closed-form order parameter estimate.

r ≈ 1 - (1/2N) Σ_{j=2}^N λ_j⁻² ⟨v^j, ω⟩²

where λ_j are Laplacian eigenvalues and v^j are eigenvectors. Valid in the strongly-coupled regime. The exact path differentiates through Laplacian eigendecomposition. The conjugate-gradient path uses the equivalent identity Σ λ_j⁻² ⟨v^j, ω⟩² = ||L⁺ω||², avoids full eigendecomposition, and maps large dense problems to GPU-friendly matrix-vector operations.

Parameters

K : jax.Array (N, N) symmetric coupling matrix (non-negative). omegas : jax.Array (N,) natural frequencies. eps : float regularization for small eigenvalues. solver : str "auto", "eigh", or "cg". "auto" uses exact eigendecomposition up to exact_size_limit and conjugate gradient above it. exact_size_limit : int Largest N where "auto" keeps the exact eigensolver. cg_tol : float Relative tolerance for the conjugate-gradient solver. cg_maxiter : int | None Optional maximum conjugate-gradient iterations.

Returns

jax.Array Scalar estimated order parameter in [0, 1].

Raises

ValueError If the inputs are invalid or inconsistent.

Source code in src/scpn_phase_orchestrator/nn/functional.py 
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
def saf_order_parameter(
    K: jax.Array,
    omegas: jax.Array,
    eps: float = 1e-8,
    solver: str = "auto",
    exact_size_limit: int = 256,
    cg_tol: float = 1e-5,
    cg_maxiter: int | None = None,
) -> jax.Array:
    """Spectral Alignment Function: closed-form order parameter estimate.

    r ≈ 1 - (1/2N) Σ_{j=2}^N λ_j⁻² ⟨v^j, ω⟩²

    where λ_j are Laplacian eigenvalues and v^j are eigenvectors.
    Valid in the strongly-coupled regime. The exact path differentiates through
    Laplacian eigendecomposition. The conjugate-gradient path uses the
    equivalent identity Σ λ_j⁻² ⟨v^j, ω⟩² = ||L⁺ω||², avoids full
    eigendecomposition, and maps large dense problems to GPU-friendly
    matrix-vector operations.

    Parameters
    ----------
    K : jax.Array
        (N, N) symmetric coupling matrix (non-negative).
    omegas : jax.Array
        (N,) natural frequencies.
    eps : float
        regularization for small eigenvalues.
    solver : str
        "auto", "eigh", or "cg". "auto" uses exact eigendecomposition up to
        exact_size_limit and conjugate gradient above it.
    exact_size_limit : int
        Largest N where "auto" keeps the exact eigensolver.
    cg_tol : float
        Relative tolerance for the conjugate-gradient solver.
    cg_maxiter : int | None
        Optional maximum conjugate-gradient iterations.

    Returns
    -------
    jax.Array
        Scalar estimated order parameter in [0, 1].

    Raises
    ------
    ValueError
        If the inputs are invalid or inconsistent.
    """
    K = _validate_square_coupling("K", K)
    N = K.shape[0]
    omegas = _validate_frequency_vector(omegas, n=N)
    eps = _require_positive_real(eps, "eps")
    exact_size_limit = _require_positive_int(exact_size_limit, "exact_size_limit")
    cg_tol = _require_positive_real(cg_tol, "cg_tol")
    cg_maxiter = _require_positive_int_or_none(cg_maxiter, "cg_maxiter")

    if solver == "auto":
        solver = "eigh" if exact_size_limit >= N else "cg"
    if solver == "cg":
        return _saf_order_parameter_cg(K, omegas, eps, cg_tol, cg_maxiter)
    if solver != "eigh":
        raise ValueError("solver must be 'auto', 'eigh', or 'cg'")
    return _saf_order_parameter_eigh(K, omegas, eps)

saf_loss 

saf_loss(
    K: Array,
    omegas: Array,
    budget: float = 0.0,
    budget_weight: float = 0.1,
    solver: str = "auto",
    exact_size_limit: int = 256,
    cg_tol: float = 1e-05,
    cg_maxiter: int | None = None,
) -> jax.Array

Loss function for coupling topology optimization via SAF.

Minimizes -r_SAF (maximize synchronization) with optional L1 budget constraint on total coupling strength.

Parameters

K : jax.Array (N, N) symmetric coupling matrix. omegas : jax.Array (N,) natural frequencies. budget : float target total coupling strength (0 = no constraint). budget_weight : float penalty weight for budget violation. solver : str SAF solver passed to saf_order_parameter. exact_size_limit : int Largest N where "auto" keeps the exact eigensolver. cg_tol : float Relative tolerance for the conjugate-gradient solver. cg_maxiter : int | None Optional maximum conjugate-gradient iterations.

Returns

jax.Array Scalar loss (lower = better synchronization).

Source code in src/scpn_phase_orchestrator/nn/functional.py 
 972
 973
 974
 975
 976
 977
 978
 979
 980
 981
 982
 983
 984
 985
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
def saf_loss(
    K: jax.Array,
    omegas: jax.Array,
    budget: float = 0.0,
    budget_weight: float = 0.1,
    solver: str = "auto",
    exact_size_limit: int = 256,
    cg_tol: float = 1e-5,
    cg_maxiter: int | None = None,
) -> jax.Array:
    """Loss function for coupling topology optimization via SAF.

    Minimizes -r_SAF (maximize synchronization) with optional L1 budget
    constraint on total coupling strength.

    Parameters
    ----------
    K : jax.Array
        (N, N) symmetric coupling matrix.
    omegas : jax.Array
        (N,) natural frequencies.
    budget : float
        target total coupling strength (0 = no constraint).
    budget_weight : float
        penalty weight for budget violation.
    solver : str
        SAF solver passed to saf_order_parameter.
    exact_size_limit : int
        Largest N where "auto" keeps the exact eigensolver.
    cg_tol : float
        Relative tolerance for the conjugate-gradient solver.
    cg_maxiter : int | None
        Optional maximum conjugate-gradient iterations.

    Returns
    -------
    jax.Array
        Scalar loss (lower = better synchronization).
    """
    budget = _require_non_negative_real(budget, "budget")
    budget_weight = _require_non_negative_real(budget_weight, "budget_weight")

    r = saf_order_parameter(
        K,
        omegas,
        solver=solver,
        exact_size_limit=exact_size_limit,
        cg_tol=cg_tol,
        cg_maxiter=cg_maxiter,
    )
    loss = -r
    if budget > 0.0:
        total_coupling = jnp.sum(jnp.abs(K))
        loss = loss + budget_weight * jnp.maximum(total_coupling - budget, 0.0)
    return loss

KuramotoLayer

Equinox module wrapping Kuramoto dynamics as a learnable layer.

KuramotoLayer(
    n: int,              # number of oscillators
    n_steps: int = 50,   # integration steps per forward pass
    dt: float = 0.01,    # timestep
    K_scale: float = 0.1, # initialisation scale for K
    mask: jax.Array | None = None,  # sparse coupling mask
    key: jax.Array,      # PRNG key
)

Learnable parameters: K (coupling matrix), omegas (frequencies).

Method Signature Description
__call__ (phases) → final_phases Forward pass
forward_with_trajectory (phases) → (final, trajectory) With full trajectory
sync_score (phases) → R Order parameter after forward pass

kuramoto_layer

Equinox module wrapping Kuramoto dynamics as a differentiable layer.

The KuramotoLayer maps input features to oscillator phases, runs N steps of Kuramoto dynamics with a learnable coupling matrix K, and returns the synchronized phase representation. Fully differentiable via JAX autodiff.

Requires: jax>=0.4, equinox>=0.11

Classes

KuramotoLayer 

KuramotoLayer(
    n: int,
    n_steps: int = 50,
    dt: float = 0.01,
    K_scale: float = 0.1,
    mask: Array | None = None,
    *,
    key: Array,
)

Bases: Module

Differentiable Kuramoto oscillator layer.

Learnable parameters

K: (n, n) coupling matrix — controls which oscillators synchronize omegas: (n,) natural frequencies

Static config

n_steps: integration steps per forward pass dt: integration timestep

Source code in src/scpn_phase_orchestrator/nn/kuramoto_layer.py 
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
def __init__(
    self,
    n: int,
    n_steps: int = 50,
    dt: float = 0.01,
    K_scale: float = 0.1,
    mask: jax.Array | None = None,
    *,
    key: jax.Array,
) -> None:
    k1, k2 = jax.random.split(key)
    raw = K_scale * jax.random.normal(k1, (n, n))
    self.K = (raw + raw.T) / 2.0
    self.omegas = jax.random.normal(k2, (n,))
    self.mask = (
        None
        if mask is None
        else tuple(tuple(float(v) for v in row) for row in np.asarray(mask))
    )
    self.n_steps = n_steps
    self.dt = dt
    self.n = n
Attributes
coupling property 
coupling: Array

Symmetric coupling matrix used by the dynamics.

Kuramoto coupling is undirected, so the dynamics use the symmetric part (K + Kᵀ)/2. Because the loss depends only on this symmetric part, the gradient with respect to K is itself symmetric, so gradient training started from a symmetric K keeps K = Kᵀ instead of drifting into a physically meaningless directed matrix.

Returns

jax.Array Symmetric coupling matrix used by the dynamics.

Methods:
__call__ 
__call__(phases: Array) -> jax.Array

Run Kuramoto dynamics on input phases.

Parameters

phases : jax.Array (n,) initial phase angles in [0, 2pi).

Returns

jax.Array (n,) phase angles after n_steps of Kuramoto integration.

Source code in src/scpn_phase_orchestrator/nn/kuramoto_layer.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
@eqx.filter_jit
def __call__(self, phases: jax.Array) -> jax.Array:
    """Run Kuramoto dynamics on input phases.

    Parameters
    ----------
    phases : jax.Array
        (n,) initial phase angles in [0, 2pi).

    Returns
    -------
    jax.Array
        (n,) phase angles after n_steps of Kuramoto integration.
    """
    if self.mask is not None:
        final, _ = kuramoto_forward_masked(
            phases,
            self.omegas,
            self.coupling,
            jnp.asarray(self.mask),
            self.dt,
            self.n_steps,
        )
    else:
        final, _ = kuramoto_forward(
            phases,
            self.omegas,
            self.coupling,
            self.dt,
            self.n_steps,
        )
    return final
forward_with_trajectory 
forward_with_trajectory(
    phases: Array,
) -> tuple[jax.Array, jax.Array]

Run dynamics and return both final state and full trajectory.

Parameters

phases : jax.Array (n,) initial phase angles.

Returns

tuple[jax.Array, jax.Array] (final_phases, trajectory) where trajectory is (n_steps, n).

Source code in src/scpn_phase_orchestrator/nn/kuramoto_layer.py 
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
@eqx.filter_jit
def forward_with_trajectory(self, phases: jax.Array) -> tuple[jax.Array, jax.Array]:
    """Run dynamics and return both final state and full trajectory.

    Parameters
    ----------
    phases : jax.Array
        (n,) initial phase angles.

    Returns
    -------
    tuple[jax.Array, jax.Array]
        (final_phases, trajectory) where trajectory is (n_steps, n).
    """
    if self.mask is not None:
        return kuramoto_forward_masked(
            phases,
            self.omegas,
            self.coupling,
            jnp.asarray(self.mask),
            self.dt,
            self.n_steps,
        )
    return kuramoto_forward(
        phases, self.omegas, self.coupling, self.dt, self.n_steps
    )
sync_score 
sync_score(phases: Array) -> jax.Array

Run dynamics and return final synchronization (order parameter R).

Useful as a differentiable loss target: maximize R for sync, minimize R for desync.

Parameters

phases : jax.Array (n,) initial phases.

Returns

jax.Array Scalar R in [0, 1].

Source code in src/scpn_phase_orchestrator/nn/kuramoto_layer.py 
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
@eqx.filter_jit
def sync_score(self, phases: jax.Array) -> jax.Array:
    """Run dynamics and return final synchronization (order parameter R).

    Useful as a differentiable loss target: maximize R for sync,
    minimize R for desync.

    Parameters
    ----------
    phases : jax.Array
        (n,) initial phases.

    Returns
    -------
    jax.Array
        Scalar R in [0, 1].
    """
    final = self(phases)
    return order_parameter(final)

Functions:

SimplicialKuramotoLayer

Extends KuramotoLayer with learnable 3-body interaction strength σ₂.

SimplicialKuramotoLayer(
    n: int,
    n_steps: int = 50,
    dt: float = 0.01,
    K_scale: float = 0.1,
    sigma2_init: float = 0.0,  # initial 3-body strength
    key: jax.Array,
)

Learnable parameters: K, omegas, sigma2.

When sigma2=0, output matches KuramotoLayer (verified in tests).

simplicial_layer

Equinox module wrapping simplicial (3-body) Kuramoto dynamics.

Extends KuramotoLayer with a learnable 3-body coupling strength sigma2. When sigma2=0, reduces to standard pairwise Kuramoto. Nonzero sigma2 produces explosive (first-order) synchronization transitions (Gambuzza et al. 2023, Nature Physics).

First differentiable 3-body Kuramoto layer in open source.

Requires: jax>=0.4, equinox>=0.11

Classes

SimplicialKuramotoLayer 

SimplicialKuramotoLayer(
    n: int,
    n_steps: int = 50,
    dt: float = 0.01,
    K_scale: float = 0.1,
    sigma2_init: float = 0.0,
    *,
    key: Array,
)

Bases: Module

Differentiable simplicial Kuramoto layer with 3-body interactions.

Learnable parameters

K: (n, n) pairwise coupling matrix omegas: (n,) natural frequencies sigma2: scalar 3-body coupling strength

Static config

n_steps: integration steps per forward pass dt: integration timestep

Source code in src/scpn_phase_orchestrator/nn/simplicial_layer.py 
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
def __init__(
    self,
    n: int,
    n_steps: int = 50,
    dt: float = 0.01,
    K_scale: float = 0.1,
    sigma2_init: float = 0.0,
    *,
    key: jax.Array,
) -> None:
    k1, k2 = jax.random.split(key)
    raw = K_scale * jax.random.normal(k1, (n, n))
    self.K = (raw + raw.T) / 2.0
    self.omegas = jax.random.normal(k2, (n,))
    self.sigma2 = jnp.array(sigma2_init)
    self.n_steps = n_steps
    self.dt = dt
    self.n = n
Attributes
coupling property 
coupling: Array

Symmetric pairwise coupling matrix used by the dynamics (K + Kᵀ)/2.

Pairwise coupling is undirected, so the loss depends only on the symmetric part; the gradient w.r.t. K is therefore symmetric and training from a symmetric K keeps it symmetric rather than drifting into a directed matrix.

Returns

jax.Array Symmetric pairwise coupling matrix used by the dynamics (K + Kᵀ)/2.

Methods:
__call__ 
__call__(phases: Array) -> jax.Array

Run simplicial Kuramoto dynamics on input phases.

Parameters

phases : jax.Array (n,) initial phase angles in [0, 2pi).

Returns

jax.Array (n,) phase angles after n_steps of integration.

Source code in src/scpn_phase_orchestrator/nn/simplicial_layer.py 
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
@eqx.filter_jit
def __call__(self, phases: jax.Array) -> jax.Array:
    """Run simplicial Kuramoto dynamics on input phases.

    Parameters
    ----------
    phases : jax.Array
        (n,) initial phase angles in [0, 2pi).

    Returns
    -------
    jax.Array
        (n,) phase angles after n_steps of integration.
    """
    final, _ = simplicial_forward(
        phases,
        self.omegas,
        self.coupling,
        self.dt,
        self.n_steps,
        sigma2=self.sigma2,
    )
    return final
forward_with_trajectory 
forward_with_trajectory(
    phases: Array,
) -> tuple[jax.Array, jax.Array]

Run dynamics and return both final state and full trajectory.

Parameters

phases : jax.Array (n,) initial phase angles.

Returns

tuple[jax.Array, jax.Array] (final_phases, trajectory) where trajectory is (n_steps, n).

Source code in src/scpn_phase_orchestrator/nn/simplicial_layer.py 
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
@eqx.filter_jit
def forward_with_trajectory(
    self,
    phases: jax.Array,
) -> tuple[jax.Array, jax.Array]:
    """Run dynamics and return both final state and full trajectory.

    Parameters
    ----------
    phases : jax.Array
        (n,) initial phase angles.

    Returns
    -------
    tuple[jax.Array, jax.Array]
        (final_phases, trajectory) where trajectory is (n_steps, n).
    """
    return simplicial_forward(
        phases,
        self.omegas,
        self.coupling,
        self.dt,
        self.n_steps,
        sigma2=self.sigma2,
    )
sync_score 
sync_score(phases: Array) -> jax.Array

Run dynamics and return final synchronization (order parameter R).

Parameters

phases : jax.Array (n,) initial phases.

Returns

jax.Array Scalar R in [0, 1].

Source code in src/scpn_phase_orchestrator/nn/simplicial_layer.py 
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
@eqx.filter_jit
def sync_score(self, phases: jax.Array) -> jax.Array:
    """Run dynamics and return final synchronization (order parameter R).

    Parameters
    ----------
    phases : jax.Array
        (n,) initial phases.

    Returns
    -------
    jax.Array
        Scalar R in [0, 1].
    """
    final = self(phases)
    return order_parameter(final)

Functions:

StuartLandauLayer

Phase + amplitude dynamics with learnable bifurcation parameters.

StuartLandauLayer(
    n: int,
    n_steps: int = 50,
    dt: float = 0.01,
    K_scale: float = 0.1,
    epsilon: float = 1.0,
    key: jax.Array,
)

Learnable parameters: K, K_r (amplitude coupling), omegas, mu (bifurcation).

Method Returns Description
__call__ (phases, amps) Forward phase + amplitude
sync_score scalar R from final phases
mean_amplitude scalar Mean amplitude after forward

stuart_landau_layer

Equinox module wrapping Stuart-Landau dynamics as a differentiable layer.

Unlike the Kuramoto-only KuramotoLayer, this layer has both phase AND amplitude dynamics, enabling representation of feature presence/absence (amplitude) alongside binding relationships (phase).

Solves AKOrN's limitations: amplitude allows memory, no N>32 degradation, supercritical/subcritical bifurcation as a natural activation gate.

Requires: jax>=0.4, equinox>=0.11

Classes

StuartLandauLayer 

StuartLandauLayer(
    n: int,
    n_steps: int = 50,
    dt: float = 0.01,
    K_scale: float = 0.1,
    epsilon: float = 1.0,
    *,
    key: Array,
)

Bases: Module

Differentiable Stuart-Landau oscillator layer.

Learnable parameters

K: (n, n) phase coupling matrix K_r: (n, n) amplitude coupling matrix omegas: (n,) natural frequencies mu: (n,) bifurcation parameters (>0: supercritical, <0: subcritical)

Static config

n_steps: integration steps per forward pass dt: integration timestep epsilon: amplitude coupling strength

Source code in src/scpn_phase_orchestrator/nn/stuart_landau_layer.py 
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
def __init__(
    self,
    n: int,
    n_steps: int = 50,
    dt: float = 0.01,
    K_scale: float = 0.1,
    epsilon: float = 1.0,
    *,
    key: jax.Array,
) -> None:
    k1, k2, k3, k4 = jax.random.split(key, 4)
    raw = K_scale * jax.random.normal(k1, (n, n))
    self.K = (raw + raw.T) / 2.0
    raw_r = K_scale * jax.random.normal(k2, (n, n))
    self.K_r = (raw_r + raw_r.T) / 2.0
    self.omegas = jax.random.normal(k3, (n,))
    # mu > 0 by default: supercritical regime (oscillators have amplitude)
    self.mu = 0.5 + 0.1 * jax.random.normal(k4, (n,))
    self.n_steps = n_steps
    self.dt = dt
    self.epsilon = epsilon
    self.n = n
Attributes
coupling property 
coupling: Array

Symmetric phase-coupling matrix used by the dynamics (K + Kᵀ)/2.

Phase coupling is undirected, so the loss depends only on the symmetric part; the gradient w.r.t. K is therefore symmetric and training from a symmetric K keeps it symmetric instead of drifting directed.

Returns

jax.Array Symmetric phase-coupling matrix used by the dynamics (K + Kᵀ)/2.

coupling_r property 
coupling_r: Array

Symmetric amplitude-coupling matrix (K_r + K_rᵀ)/2 (see coupling).

Returns

jax.Array Symmetric amplitude-coupling matrix (K_r + K_rᵀ)/2 (see coupling).

Methods:
__call__ 
__call__(
    phases: Array, amplitudes: Array
) -> tuple[jax.Array, jax.Array]

Run Stuart-Landau dynamics on input state.

Parameters

phases : jax.Array (n,) initial phase angles in [0, 2pi). amplitudes : jax.Array (n,) initial amplitudes (r >= 0).

Returns

tuple[jax.Array, jax.Array] (final_phases, final_amplitudes).

Source code in src/scpn_phase_orchestrator/nn/stuart_landau_layer.py 
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
@eqx.filter_jit
def __call__(
    self, phases: jax.Array, amplitudes: jax.Array
) -> tuple[jax.Array, jax.Array]:
    """Run Stuart-Landau dynamics on input state.

    Parameters
    ----------
    phases : jax.Array
        (n,) initial phase angles in [0, 2pi).
    amplitudes : jax.Array
        (n,) initial amplitudes (r >= 0).

    Returns
    -------
    tuple[jax.Array, jax.Array]
        (final_phases, final_amplitudes).
    """
    fp, fr, _, _ = stuart_landau_forward(
        phases,
        amplitudes,
        self.omegas,
        self.mu,
        self.coupling,
        self.coupling_r,
        self.dt,
        self.n_steps,
        self.epsilon,
    )
    return fp, fr
forward_with_trajectory 
forward_with_trajectory(
    phases: Array, amplitudes: Array
) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]

Run dynamics and return full trajectories.

Returns
(final_phases, final_amplitudes, phase_trajectory, amplitude_trajectory)
Parameters

phases : jax.Array Oscillator phases in radians, shape (N,). amplitudes : jax.Array Oscillator amplitudes, shape (N,).

Source code in src/scpn_phase_orchestrator/nn/stuart_landau_layer.py 
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
@eqx.filter_jit
def forward_with_trajectory(
    self, phases: jax.Array, amplitudes: jax.Array
) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]:
    """Run dynamics and return full trajectories.

    Returns
    -------
        (final_phases, final_amplitudes, phase_trajectory, amplitude_trajectory)

    Parameters
    ----------
    phases : jax.Array
        Oscillator phases in radians, shape ``(N,)``.
    amplitudes : jax.Array
        Oscillator amplitudes, shape ``(N,)``.
    """
    return stuart_landau_forward(
        phases,
        amplitudes,
        self.omegas,
        self.mu,
        self.coupling,
        self.coupling_r,
        self.dt,
        self.n_steps,
        self.epsilon,
    )
sync_score 
sync_score(phases: Array, amplitudes: Array) -> jax.Array

Run dynamics and return final synchronization (order parameter R).

Parameters

phases : jax.Array (n,) initial phases. amplitudes : jax.Array (n,) initial amplitudes.

Returns

jax.Array Scalar R in [0, 1].

Source code in src/scpn_phase_orchestrator/nn/stuart_landau_layer.py 
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
@eqx.filter_jit
def sync_score(self, phases: jax.Array, amplitudes: jax.Array) -> jax.Array:
    """Run dynamics and return final synchronization (order parameter R).

    Parameters
    ----------
    phases : jax.Array
        (n,) initial phases.
    amplitudes : jax.Array
        (n,) initial amplitudes.

    Returns
    -------
    jax.Array
        Scalar R in [0, 1].
    """
    fp, _ = self(phases, amplitudes)
    return order_parameter(fp)
mean_amplitude 
mean_amplitude(
    phases: Array, amplitudes: Array
) -> jax.Array

Run dynamics and return mean final amplitude.

Useful as a differentiable activity measure: high mean amplitude means oscillators are active (supercritical), low means quiescent.

Parameters

phases : jax.Array (n,) initial phases. amplitudes : jax.Array (n,) initial amplitudes.

Returns

jax.Array Scalar mean amplitude.

Source code in src/scpn_phase_orchestrator/nn/stuart_landau_layer.py 
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
@eqx.filter_jit
def mean_amplitude(self, phases: jax.Array, amplitudes: jax.Array) -> jax.Array:
    """Run dynamics and return mean final amplitude.

    Useful as a differentiable activity measure: high mean amplitude
    means oscillators are active (supercritical), low means quiescent.

    Parameters
    ----------
    phases : jax.Array
        (n,) initial phases.
    amplitudes : jax.Array
        (n,) initial amplitudes.

    Returns
    -------
    jax.Array
        Scalar mean amplitude.
    """
    import jax.numpy as jnp

    _, fr = self(phases, amplitudes)
    return jnp.mean(fr)

Functions:

BOLD Signal Generator

Balloon-Windkessel hemodynamic model converting oscillator amplitudes to simulated fMRI BOLD signal. Differentiable for gradient-based fMRI fitting.

Function Description
balloon_windkessel_step One Euler step of BW hemodynamics
bold_signal V,Q → BOLD observation equation
bold_from_neural Full neural → BOLD conversion

State variables: signal (s), flow (f), volume (v), deoxyhemoglobin (q).

bold

Balloon-Windkessel hemodynamic model in JAX.

Converts neural activity (oscillator amplitude envelope) to simulated fMRI BOLD signal. Fully differentiable for gradient-based optimization of oscillator parameters to match empirical fMRI data.

Friston et al. 2000 (Balloon model), Stephan et al. 2007 (parameters). Requires: jax>=0.4

Functions:

balloon_windkessel_step 

balloon_windkessel_step(
    s: Array,
    f: Array,
    v: Array,
    q: Array,
    x: Array,
    dt: float,
    kappa: float = KAPPA,
    gamma: float = GAMMA,
    tau: float = TAU,
    alpha: float = ALPHA,
    e0: float = E0,
) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]

Single Euler step of the Balloon-Windkessel hemodynamic model.

Parameters

s : jax.Array (N,) vasodilatory signal. f : jax.Array (N,) blood inflow (normalised, resting=1). v : jax.Array (N,) blood volume (normalised, resting=1). q : jax.Array (N,) deoxyhaemoglobin content (normalised, resting=1). x : jax.Array (N,) neural input (amplitude envelope). dt : float integration timestep (seconds). kappa : float signal decay rate (default 0.65). gamma : float flow-dependent elimination (default 0.41). tau : float haemodynamic transit time (default 0.98). alpha : float Grubb's vessel stiffness exponent (default 0.32). e0 : float resting oxygen extraction fraction (default 0.34).

Returns

tuple[jax.Array, jax.Array, jax.Array, jax.Array] (new_s, new_f, new_v, new_q).

Source code in src/scpn_phase_orchestrator/nn/bold.py 
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
def balloon_windkessel_step(
    s: jax.Array,
    f: jax.Array,
    v: jax.Array,
    q: jax.Array,
    x: jax.Array,
    dt: float,
    kappa: float = KAPPA,
    gamma: float = GAMMA,
    tau: float = TAU,
    alpha: float = ALPHA,
    e0: float = E0,
) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]:
    """Single Euler step of the Balloon-Windkessel hemodynamic model.

    Parameters
    ----------
    s : jax.Array
        (N,) vasodilatory signal.
    f : jax.Array
        (N,) blood inflow (normalised, resting=1).
    v : jax.Array
        (N,) blood volume (normalised, resting=1).
    q : jax.Array
        (N,) deoxyhaemoglobin content (normalised, resting=1).
    x : jax.Array
        (N,) neural input (amplitude envelope).
    dt : float
        integration timestep (seconds).
    kappa : float
        signal decay rate (default 0.65).
    gamma : float
        flow-dependent elimination (default 0.41).
    tau : float
        haemodynamic transit time (default 0.98).
    alpha : float
        Grubb's vessel stiffness exponent (default 0.32).
    e0 : float
        resting oxygen extraction fraction (default 0.34).

    Returns
    -------
    tuple[jax.Array, jax.Array, jax.Array, jax.Array]
        (new_s, new_f, new_v, new_q).
    """
    E_f = 1.0 - (1.0 - e0) ** (1.0 / jnp.maximum(f, 0.01))

    ds = x - kappa * s - gamma * (f - 1.0)
    df = s
    dv = (1.0 / tau) * (f - v ** (1.0 / alpha))
    dq = (1.0 / tau) * (f * E_f / e0 - v ** (1.0 / alpha) * q / jnp.maximum(v, 0.01))

    new_s = s + dt * ds
    new_f = jnp.maximum(f + dt * df, 0.01)
    new_v = jnp.maximum(v + dt * dv, 0.01)
    new_q = jnp.maximum(q + dt * dq, 0.01)

    return new_s, new_f, new_v, new_q

bold_signal 

bold_signal(
    v: Array,
    q: Array,
    v0: float = V0,
    k1: float = K1,
    k2: float = K2,
    k3: float = K3,
) -> jax.Array

Compute BOLD signal from blood volume and deoxyhemoglobin.

Parameters

v : jax.Array (N,) or (T, N) blood volume. q : jax.Array (N,) or (T, N) deoxyhaemoglobin.

Returns

jax.Array BOLD signal, same shape as input.

Source code in src/scpn_phase_orchestrator/nn/bold.py 
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
def bold_signal(
    v: jax.Array,
    q: jax.Array,
    v0: float = V0,
    k1: float = K1,
    k2: float = K2,
    k3: float = K3,
) -> jax.Array:
    """Compute BOLD signal from blood volume and deoxyhemoglobin.

    Parameters
    ----------
    v : jax.Array
        (N,) or (T, N) blood volume.
    q : jax.Array
        (N,) or (T, N) deoxyhaemoglobin.

    Returns
    -------
    jax.Array
        BOLD signal, same shape as input.
    """
    return v0 * (k1 * (1.0 - q) + k2 * (1.0 - q / v) + k3 * (1.0 - v))

bold_from_neural 

bold_from_neural(
    neural: Array,
    dt: float,
    dt_bold: float = 0.5,
    kappa: float = KAPPA,
    gamma: float = GAMMA,
    tau: float = TAU,
    alpha: float = ALPHA,
    e0: float = E0,
) -> jax.Array

Generate BOLD signal from neural activity time series.

Runs the Balloon-Windkessel model on the neural input and returns the BOLD signal at a lower sampling rate (TR = dt_bold).

Parameters

neural : jax.Array (T, N) neural activity time series (e.g., amplitude envelope). dt : float simulation timestep (seconds). dt_bold : float BOLD sampling period (seconds, default 0.5s = 2Hz). kappa : float signal decay rate (default 0.65). gamma : float flow-dependent elimination (default 0.41). tau : float haemodynamic transit time (default 0.98). alpha : float Grubb's vessel stiffness exponent (default 0.32). e0 : float resting oxygen extraction fraction (default 0.34).

Returns

jax.Array (T_bold, N) BOLD signal, where T_bold = T * dt / dt_bold.

Source code in src/scpn_phase_orchestrator/nn/bold.py 
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
def bold_from_neural(
    neural: jax.Array,
    dt: float,
    dt_bold: float = 0.5,
    kappa: float = KAPPA,
    gamma: float = GAMMA,
    tau: float = TAU,
    alpha: float = ALPHA,
    e0: float = E0,
) -> jax.Array:
    """Generate BOLD signal from neural activity time series.

    Runs the Balloon-Windkessel model on the neural input and returns
    the BOLD signal at a lower sampling rate (TR = dt_bold).

    Parameters
    ----------
    neural : jax.Array
        (T, N) neural activity time series (e.g., amplitude envelope).
    dt : float
        simulation timestep (seconds).
    dt_bold : float
        BOLD sampling period (seconds, default 0.5s = 2Hz).
    kappa : float
        signal decay rate (default 0.65).
    gamma : float
        flow-dependent elimination (default 0.41).
    tau : float
        haemodynamic transit time (default 0.98).
    alpha : float
        Grubb's vessel stiffness exponent (default 0.32).
    e0 : float
        resting oxygen extraction fraction (default 0.34).

    Returns
    -------
    jax.Array
        (T_bold, N) BOLD signal, where T_bold = T * dt / dt_bold.
    """
    T, n_regions = neural.shape
    subsample = max(1, int(dt_bold / dt))

    def step(
        carry: tuple[jax.Array, jax.Array, jax.Array, jax.Array],
        x_t: jax.Array,
    ) -> tuple[tuple[jax.Array, jax.Array, jax.Array, jax.Array], jax.Array]:
        """One scan step: advance hemodynamic state and emit BOLD."""
        s, f, v, q = carry
        new_s, new_f, new_v, new_q = balloon_windkessel_step(
            s, f, v, q, x_t, dt, kappa, gamma, tau, alpha, e0
        )
        y = bold_signal(new_v, new_q)
        return (new_s, new_f, new_v, new_q), y

    s0 = jnp.zeros(n_regions)
    f0 = jnp.ones(n_regions)
    v0 = jnp.ones(n_regions)
    q0 = jnp.ones(n_regions)

    _, bold_full = jax.lax.scan(step, (s0, f0, v0, q0), neural)
    result: jax.Array = bold_full[::subsample]
    return result

Reservoir Computing

Kuramoto-based echo state network with linear readout.

Function Description
reservoir_features(phases) cos/sin feature extraction
reservoir_drive(phases, omegas, K, W_in, u, dt, n_steps) Driven reservoir dynamics
ridge_readout(features, targets, alpha=1e-4) Ridge regression readout weights
reservoir_predict(features, W_out) Prediction from trained readout

Universal approximation near edge-of-bifurcation (arXiv:2407.16172).

reservoir

Kuramoto-based reservoir computing in JAX.

Uses a Kuramoto oscillator network as a nonlinear reservoir. Input signals modulate natural frequencies; the reservoir's phase state is read out via a trained linear layer.

Theory: universal approximation near edge-of-bifurcation (arXiv:2407.16172, 2024). The Ott-Antonsen critical coupling K_c = 2*Delta defines the optimal operating point.

Requires: jax>=0.4

Functions:

reservoir_features 

reservoir_features(phases: Array) -> jax.Array

Extract features from oscillator phases for readout.

Features: [cos(theta_1), sin(theta_1), ..., cos(theta_N), sin(theta_N), R] Total: 2*N + 1 features.

Parameters

phases : jax.Array (N,) oscillator phases.

Returns

jax.Array (2*N + 1,) feature vector.

Source code in src/scpn_phase_orchestrator/nn/reservoir.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def reservoir_features(phases: jax.Array) -> jax.Array:
    """Extract features from oscillator phases for readout.

    Features: [cos(theta_1), sin(theta_1), ..., cos(theta_N), sin(theta_N), R]
    Total: 2*N + 1 features.

    Parameters
    ----------
    phases : jax.Array
        (N,) oscillator phases.

    Returns
    -------
    jax.Array
        (2*N + 1,) feature vector.
    """
    R = order_parameter(phases)
    return jnp.concatenate([jnp.cos(phases), jnp.sin(phases), R[jnp.newaxis]])

reservoir_drive 

reservoir_drive(
    phases: Array,
    omegas: Array,
    K: Array,
    W_in: Array,
    u: Array,
    dt: float,
    n_steps: int,
) -> jax.Array

Drive reservoir with input signal and collect features at each step.

Input is injected into natural frequencies: omega_i(t) = omega_i + W_in @ u(t).

Parameters

phases : jax.Array (N,) initial oscillator phases. omegas : jax.Array (N,) base natural frequencies. K : jax.Array (N, N) fixed coupling matrix. W_in : jax.Array (N, D_in) input weight matrix. u : jax.Array (T, D_in) input signal sequence. dt : float integration timestep. n_steps : int Kuramoto steps per input sample.

Returns

jax.Array (T, 2*N + 1) feature matrix for readout training.

Source code in src/scpn_phase_orchestrator/nn/reservoir.py 
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
def reservoir_drive(
    phases: jax.Array,
    omegas: jax.Array,
    K: jax.Array,
    W_in: jax.Array,
    u: jax.Array,
    dt: float,
    n_steps: int,
) -> jax.Array:
    """Drive reservoir with input signal and collect features at each step.

    Input is injected into natural frequencies: omega_i(t) = omega_i + W_in @ u(t).

    Parameters
    ----------
    phases : jax.Array
        (N,) initial oscillator phases.
    omegas : jax.Array
        (N,) base natural frequencies.
    K : jax.Array
        (N, N) fixed coupling matrix.
    W_in : jax.Array
        (N, D_in) input weight matrix.
    u : jax.Array
        (T, D_in) input signal sequence.
    dt : float
        integration timestep.
    n_steps : int
        Kuramoto steps per input sample.

    Returns
    -------
    jax.Array
        (T, 2*N + 1) feature matrix for readout training.
    """

    def process_sample(carry: jax.Array, u_t: jax.Array) -> tuple[jax.Array, jax.Array]:
        """Drive reservoir one input step and extract features."""
        p = carry
        driven_omegas = omegas + W_in @ u_t
        p, _ = kuramoto_forward(p, driven_omegas, K, dt, n_steps)
        feat = reservoir_features(p)
        return p, feat

    _, feat_matrix = jax.lax.scan(process_sample, phases, u)
    result: jax.Array = feat_matrix
    return result

ridge_readout 

ridge_readout(
    features: Array, targets: Array, alpha: float = 0.0001
) -> jax.Array

Train linear readout via ridge regression.

W_out = (F^T F + alpha I)^{-1} F^T Y

Parameters

features : jax.Array (T, D_feat) reservoir feature matrix. targets : jax.Array (T, D_out) target outputs. alpha : float L2 regularization strength.

Returns

jax.Array (D_feat, D_out) readout weight matrix.

Source code in src/scpn_phase_orchestrator/nn/reservoir.py 
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
def ridge_readout(
    features: jax.Array,
    targets: jax.Array,
    alpha: float = 1e-4,
) -> jax.Array:
    """Train linear readout via ridge regression.

    W_out = (F^T F + alpha I)^{-1} F^T Y

    Parameters
    ----------
    features : jax.Array
        (T, D_feat) reservoir feature matrix.
    targets : jax.Array
        (T, D_out) target outputs.
    alpha : float
        L2 regularization strength.

    Returns
    -------
    jax.Array
        (D_feat, D_out) readout weight matrix.
    """
    FtF = features.T @ features + alpha * jnp.eye(features.shape[1])
    FtY = features.T @ targets
    result: jax.Array = jnp.linalg.solve(FtF, FtY)
    return result

reservoir_predict 

reservoir_predict(
    features: Array, W_out: Array
) -> jax.Array

Apply trained readout to reservoir features.

Parameters

features : jax.Array (T, D_feat) feature matrix. W_out : jax.Array (D_feat, D_out) readout weights.

Returns

jax.Array (T, D_out) predictions.

Source code in src/scpn_phase_orchestrator/nn/reservoir.py 
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
def reservoir_predict(
    features: jax.Array,
    W_out: jax.Array,
) -> jax.Array:
    """Apply trained readout to reservoir features.

    Parameters
    ----------
    features : jax.Array
        (T, D_feat) feature matrix.
    W_out : jax.Array
        (D_feat, D_out) readout weights.

    Returns
    -------
    jax.Array
        (T, D_out) predictions.
    """
    result: jax.Array = features @ W_out
    return result

Differentiable Chimera Metrics

JAX-native local order-parameter and chimera-index helpers for gradient-aware topology searches.

chimera

JAX-based chimera state detection for coupled oscillator networks.

Chimera states are spatiotemporal patterns where synchronised and incoherent domains coexist (Kuramoto & Battogtokh 2002). This module provides differentiable detection, enabling gradient-based search for chimera-producing coupling matrices.

Requires: jax>=0.4

Functions:

local_order_parameter 

local_order_parameter(phases: Array, K: Array) -> jax.Array

Local Kuramoto order parameter R_i for each oscillator.

R_i = |mean(exp(i·Δθ_j)) for neighbours j of i|

Neighbours defined by nonzero entries in K. Vectorised — no Python loops.

Parameters

phases : jax.Array (N,) oscillator phases. K : jax.Array (N, N) coupling matrix (nonzero = neighbour).

Returns

jax.Array (N,) local order parameters in [0, 1].

Source code in src/scpn_phase_orchestrator/nn/chimera.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
def local_order_parameter(
    phases: jax.Array,
    K: jax.Array,
) -> jax.Array:
    """Local Kuramoto order parameter R_i for each oscillator.

    R_i = |mean(exp(i·Δθ_j)) for neighbours j of i|

    Neighbours defined by nonzero entries in K. Vectorised — no Python loops.

    Parameters
    ----------
    phases : jax.Array
        (N,) oscillator phases.
    K : jax.Array
        (N, N) coupling matrix (nonzero = neighbour).

    Returns
    -------
    jax.Array
        (N,) local order parameters in [0, 1].
    """
    mask = (K != 0).astype(jnp.float32)
    diff = phases[jnp.newaxis, :] - phases[:, jnp.newaxis]  # (N, N)
    # Complex phasors weighted by adjacency
    cos_diff = jnp.cos(diff) * mask
    sin_diff = jnp.sin(diff) * mask
    n_neighbours = jnp.sum(mask, axis=1).clip(min=1.0)
    mean_cos = jnp.sum(cos_diff, axis=1) / n_neighbours
    mean_sin = jnp.sum(sin_diff, axis=1) / n_neighbours
    return jnp.sqrt(mean_cos**2 + mean_sin**2)

chimera_index 

chimera_index(phases: Array, K: Array) -> jax.Array

Scalar chimera index: variance of local order parameters.

High variance = coexistence of coherent (R≈1) and incoherent (R≈0) domains. Zero variance = uniform state (either all sync or all desync). Differentiable.

Parameters

phases : jax.Array (N,) oscillator phases. K : jax.Array (N, N) coupling matrix.

Returns

jax.Array Scalar chimera index (higher = more chimera-like).

Source code in src/scpn_phase_orchestrator/nn/chimera.py 
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
def chimera_index(
    phases: jax.Array,
    K: jax.Array,
) -> jax.Array:
    """Scalar chimera index: variance of local order parameters.

    High variance = coexistence of coherent (R≈1) and incoherent (R≈0)
    domains. Zero variance = uniform state (either all sync or all desync).
    Differentiable.

    Parameters
    ----------
    phases : jax.Array
        (N,) oscillator phases.
    K : jax.Array
        (N, N) coupling matrix.

    Returns
    -------
    jax.Array
        Scalar chimera index (higher = more chimera-like).
    """
    R_local = local_order_parameter(phases, K)
    return jnp.var(R_local)

detect_chimera 

detect_chimera(
    phases: Array,
    K: Array,
    coherent_threshold: float = 0.8,
    incoherent_threshold: float = 0.3,
) -> tuple[jax.Array, jax.Array]

Classify oscillators as coherent or incoherent.

Parameters

phases : jax.Array (N,) oscillator phases. K : jax.Array (N, N) coupling matrix. coherent_threshold : float R_i above this → coherent. incoherent_threshold : float R_i below this → incoherent.

Returns

tuple[jax.Array, jax.Array] (coherent_mask, incoherent_mask): (N,) boolean arrays.

Source code in src/scpn_phase_orchestrator/nn/chimera.py 
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
def detect_chimera(
    phases: jax.Array,
    K: jax.Array,
    coherent_threshold: float = 0.8,
    incoherent_threshold: float = 0.3,
) -> tuple[jax.Array, jax.Array]:
    """Classify oscillators as coherent or incoherent.

    Parameters
    ----------
    phases : jax.Array
        (N,) oscillator phases.
    K : jax.Array
        (N, N) coupling matrix.
    coherent_threshold : float
        R_i above this → coherent.
    incoherent_threshold : float
        R_i below this → incoherent.

    Returns
    -------
    tuple[jax.Array, jax.Array]
        (coherent_mask, incoherent_mask): (N,) boolean arrays.
    """
    R_local = local_order_parameter(phases, K)
    coherent = R_local >= coherent_threshold
    incoherent = R_local <= incoherent_threshold
    return coherent, incoherent

Differentiable Spectral Metrics

JAX-native graph Laplacian metrics used by topology and synchronisability experiments.

spectral

Differentiable spectral metrics for coupling matrix analysis.

All functions are differentiable via jnp.linalg.eigh, enabling gradient-based topology optimisation: find the sparsest K that maintains synchronisability above a target threshold.

Requires: jax>=0.4

Functions:

laplacian_spectrum 

laplacian_spectrum(K: Array) -> jax.Array

Sorted eigenvalues of the graph Laplacian L = D - K.

Parameters

K : jax.Array (N, N) symmetric coupling matrix (non-negative weights).

Returns

jax.Array (N,) eigenvalues in ascending order. First is ~0 (connected graph).

Source code in src/scpn_phase_orchestrator/nn/spectral.py 
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
def laplacian_spectrum(K: jax.Array) -> jax.Array:
    """Sorted eigenvalues of the graph Laplacian L = D - K.

    Parameters
    ----------
    K : jax.Array
        (N, N) symmetric coupling matrix (non-negative weights).

    Returns
    -------
    jax.Array
        (N,) eigenvalues in ascending order. First is ~0 (connected graph).
    """
    D = jnp.diag(jnp.sum(K, axis=1))
    L = D - K
    eigenvalues: jax.Array = jnp.linalg.eigh(L)[0]
    return eigenvalues

algebraic_connectivity 

algebraic_connectivity(K: Array) -> jax.Array

Second-smallest Laplacian eigenvalue (Fiedler value).

Measures how well-connected the network is. Zero iff disconnected. Differentiable — gradient flows through eigh.

Parameters

K : jax.Array (N, N) symmetric coupling matrix.

Returns

jax.Array Scalar lambda_2.

Source code in src/scpn_phase_orchestrator/nn/spectral.py 
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
def algebraic_connectivity(K: jax.Array) -> jax.Array:
    """Second-smallest Laplacian eigenvalue (Fiedler value).

    Measures how well-connected the network is. Zero iff disconnected.
    Differentiable — gradient flows through eigh.

    Parameters
    ----------
    K : jax.Array
        (N, N) symmetric coupling matrix.

    Returns
    -------
    jax.Array
        Scalar lambda_2.
    """
    return laplacian_spectrum(K)[1]

eigenratio 

eigenratio(K: Array) -> jax.Array

Ratio lambda_N / lambda_2 (synchronisability metric).

Lower eigenratio = more synchronisable (Barahona & Pecora 2002). The MSF (master stability function) approach shows that coupled oscillators synchronise when all transverse eigenvalues fall within the MSF stability interval.

Parameters

K : jax.Array (N, N) symmetric coupling matrix.

Returns

jax.Array Scalar lambda_N / lambda_2.

Source code in src/scpn_phase_orchestrator/nn/spectral.py 
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
def eigenratio(K: jax.Array) -> jax.Array:
    """Ratio lambda_N / lambda_2 (synchronisability metric).

    Lower eigenratio = more synchronisable (Barahona & Pecora 2002).
    The MSF (master stability function) approach shows that coupled
    oscillators synchronise when all transverse eigenvalues fall
    within the MSF stability interval.

    Parameters
    ----------
    K : jax.Array
        (N, N) symmetric coupling matrix.

    Returns
    -------
    jax.Array
        Scalar lambda_N / lambda_2.
    """
    eigs = laplacian_spectrum(K)
    lambda_2 = eigs[1]
    lambda_N = eigs[-1]
    return lambda_N / jnp.maximum(lambda_2, 1e-10)

sync_threshold 

sync_threshold(K: Array, omegas: Array) -> jax.Array

Critical coupling strength estimate (Dorfler & Bullo 2014).

K_c ≈ max|ω_i - ω_j| / lambda_2

Below K_c, the network cannot synchronise. Above, it can.

Parameters

K : jax.Array (N, N) symmetric coupling matrix. omegas : jax.Array (N,) natural frequencies.

Returns

jax.Array Scalar estimated critical coupling.

Source code in src/scpn_phase_orchestrator/nn/spectral.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def sync_threshold(
    K: jax.Array,
    omegas: jax.Array,
) -> jax.Array:
    """Critical coupling strength estimate (Dorfler & Bullo 2014).

    K_c ≈ max|ω_i - ω_j| / lambda_2

    Below K_c, the network cannot synchronise. Above, it can.

    Parameters
    ----------
    K : jax.Array
        (N, N) symmetric coupling matrix.
    omegas : jax.Array
        (N,) natural frequencies.

    Returns
    -------
    jax.Array
        Scalar estimated critical coupling.
    """
    lambda_2 = algebraic_connectivity(K)
    omega_spread = jnp.max(omegas) - jnp.min(omegas)
    return omega_spread / jnp.maximum(lambda_2, 1e-10)

Theta Neuron Dynamics

Differentiable Ermentrout-Kopell theta-neuron dynamics for excitable systems.

theta_neuron

Theta neuron (Ermentrout-Kopell canonical model) for coupled excitable systems.

dθ_i/dt = (1 - cos(θ_i)) + (1 + cos(θ_i)) · (η_i + I_syn_i)

where I_syn_i = Σ_j K_ij · (1 - cos(θ_j)) is synaptic input.

The theta neuron is the canonical model for Type I neuronal excitability (Ermentrout & Kopell 1986). Unlike Kuramoto oscillators which are always oscillating, theta neurons can be excitable (η < 0) — they fire only when driven by sufficient synaptic input.

Requires: jax>=0.4

Classes

ThetaNeuronLayer 

ThetaNeuronLayer(
    n: int,
    n_steps: int = 50,
    dt: float = 0.01,
    K_scale: float = 0.1,
    eta_mean: float = -0.5,
    *,
    key: Array,
)

Bases: Module

Differentiable theta neuron layer.

Learnable parameters

K: (n, n) synaptic coupling matrix eta: (n,) excitability parameters

Static config

n_steps, dt

Source code in src/scpn_phase_orchestrator/nn/theta_neuron.py 
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
def __init__(
    self,
    n: int,
    n_steps: int = 50,
    dt: float = 0.01,
    K_scale: float = 0.1,
    eta_mean: float = -0.5,
    *,
    key: jax.Array,
) -> None:
    k1, k2 = jax.random.split(key)
    raw = K_scale * jax.random.normal(k1, (n, n))
    self.K = (raw + raw.T) / 2.0
    # Default η<0 (excitable regime)
    self.eta = eta_mean + 0.1 * jax.random.normal(k2, (n,))
    self.n_steps = n_steps
    self.dt = dt
    self.n = n
Methods:
__call__ 
__call__(phases: Array) -> jax.Array

Run theta neuron dynamics on input phases.

Source code in src/scpn_phase_orchestrator/nn/theta_neuron.py 
185
186
187
188
189
190
191
192
193
194
195
@eqx.filter_jit
def __call__(self, phases: jax.Array) -> jax.Array:
    """Run theta neuron dynamics on input phases."""
    final, _ = theta_neuron_forward(
        phases,
        self.eta,
        self.K,
        self.dt,
        self.n_steps,
    )
    return final
forward_with_trajectory 
forward_with_trajectory(
    phases: Array,
) -> tuple[jax.Array, jax.Array]

Run dynamics and return full trajectory.

Parameters

phases : jax.Array Oscillator phases in radians, shape (N,).

Returns

tuple[jax.Array, jax.Array] The final phases and the full trajectory.

Source code in src/scpn_phase_orchestrator/nn/theta_neuron.py 
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
@eqx.filter_jit
def forward_with_trajectory(
    self,
    phases: jax.Array,
) -> tuple[jax.Array, jax.Array]:
    """Run dynamics and return full trajectory.

    Parameters
    ----------
    phases : jax.Array
        Oscillator phases in radians, shape ``(N,)``.

    Returns
    -------
    tuple[jax.Array, jax.Array]
        The final phases and the full trajectory.
    """
    return theta_neuron_forward(
        phases,
        self.eta,
        self.K,
        self.dt,
        self.n_steps,
    )

Functions:

theta_neuron_step 

theta_neuron_step(
    phases: Array, eta: Array, K: Array, dt: float
) -> jax.Array

Single Euler step of the theta neuron model.

Parameters

phases : jax.Array (N,) neuron phases in [0, 2pi). eta : jax.Array (N,) excitability parameters (η>0: oscillatory, η<0: excitable). K : jax.Array (N, N) synaptic coupling matrix. dt : float integration timestep.

Returns

jax.Array (N,) updated phases.

Source code in src/scpn_phase_orchestrator/nn/theta_neuron.py 
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
def theta_neuron_step(
    phases: jax.Array,
    eta: jax.Array,
    K: jax.Array,
    dt: float,
) -> jax.Array:
    """Single Euler step of the theta neuron model.

    Parameters
    ----------
    phases : jax.Array
        (N,) neuron phases in [0, 2pi).
    eta : jax.Array
        (N,) excitability parameters (η>0: oscillatory, η<0: excitable).
    K : jax.Array
        (N, N) synaptic coupling matrix.
    dt : float
        integration timestep.

    Returns
    -------
    jax.Array
        (N,) updated phases.
    """
    dphi = _theta_deriv(phases, eta, K)
    return (phases + dt * dphi) % TWO_PI

theta_neuron_rk4_step 

theta_neuron_rk4_step(
    phases: Array, eta: Array, K: Array, dt: float
) -> jax.Array

Single RK4 step of the theta neuron model.

Parameters

phases : jax.Array Oscillator phases in radians, shape (N,). eta : jax.Array Per-neuron excitability parameters, shape (N,). K : jax.Array Coupling matrix K, shape (N, N). dt : float Integration step size.

Returns

jax.Array The phases after one theta-neuron RK4 step.

Source code in src/scpn_phase_orchestrator/nn/theta_neuron.py 
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
def theta_neuron_rk4_step(
    phases: jax.Array,
    eta: jax.Array,
    K: jax.Array,
    dt: float,
) -> jax.Array:
    """Single RK4 step of the theta neuron model.

    Parameters
    ----------
    phases : jax.Array
        Oscillator phases in radians, shape ``(N,)``.
    eta : jax.Array
        Per-neuron excitability parameters, shape ``(N,)``.
    K : jax.Array
        Coupling matrix ``K``, shape ``(N, N)``.
    dt : float
        Integration step size.

    Returns
    -------
    jax.Array
        The phases after one theta-neuron RK4 step.
    """

    def deriv(p: jax.Array) -> jax.Array:
        """Theta neuron derivative at arbitrary phase."""
        return _theta_deriv(p, eta, K)

    k1 = deriv(phases)
    k2 = deriv(phases + 0.5 * dt * k1)
    k3 = deriv(phases + 0.5 * dt * k2)
    k4 = deriv(phases + dt * k3)
    return (phases + (dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4)) % TWO_PI

theta_neuron_forward 

theta_neuron_forward(
    phases: Array,
    eta: Array,
    K: Array,
    dt: float,
    n_steps: int,
    method: str = "rk4",
) -> tuple[jax.Array, jax.Array]

Run N steps of theta neuron dynamics.

Parameters

phases : jax.Array (N,) initial phases. eta : jax.Array (N,) excitability parameters. K : jax.Array (N, N) synaptic coupling. dt : float timestep. n_steps : int integration steps. method : str "rk4" or "euler".

Returns

tuple[jax.Array, jax.Array] (final, trajectory) where trajectory is (n_steps, N).

Source code in src/scpn_phase_orchestrator/nn/theta_neuron.py 
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
def theta_neuron_forward(
    phases: jax.Array,
    eta: jax.Array,
    K: jax.Array,
    dt: float,
    n_steps: int,
    method: str = "rk4",
) -> tuple[jax.Array, jax.Array]:
    """Run N steps of theta neuron dynamics.

    Parameters
    ----------
    phases : jax.Array
        (N,) initial phases.
    eta : jax.Array
        (N,) excitability parameters.
    K : jax.Array
        (N, N) synaptic coupling.
    dt : float
        timestep.
    n_steps : int
        integration steps.
    method : str
        "rk4" or "euler".

    Returns
    -------
    tuple[jax.Array, jax.Array]
        (final, trajectory) where trajectory is (n_steps, N).
    """
    step_fn = theta_neuron_rk4_step if method == "rk4" else theta_neuron_step

    def body(carry: jax.Array, _: None) -> tuple[jax.Array, jax.Array]:
        """Single scan iteration: step and record phase."""
        p = step_fn(carry, eta, K, dt)
        return p, p

    final, trajectory = jax.lax.scan(body, phases, None, length=n_steps)
    return final, trajectory

UDE-Kuramoto (Universal Differential Equation)

Physics backbone (sin(Δθ) coupling) plus a learned neural residual. The MLP handles model mismatch that the analytical Kuramoto model cannot capture.

CouplingResidual (eqx.Module)

CouplingResidual(hidden: int = 16, key: jax.Array)

Small MLP: Linear(1, hidden) → tanh → Linear(hidden, 1).

UDEKuramotoLayer (eqx.Module)

UDEKuramotoLayer(n, n_steps=50, dt=0.01, K_scale=0.1, hidden=16, key=key)

Learnable: K, omegas, residual (CouplingResidual MLP).

ude

UDE-Kuramoto: physics backbone + learned neural residual.

dθ_i/dt = ω_i + Σ_j K_ij · [sin(θ_j - θ_i) + NN_φ(θ_j - θ_i)]

The known Kuramoto structure provides the mechanistic backbone. A small neural network NN_φ handles model mismatch: higher harmonics, asymmetric coupling, amplitude-dependent effects. Trained end-to-end via JAX autodiff.

Rackauckas et al. 2020 (UDE framework); Frontiers Comp. Neuro. 2025. First Python UDE implementation for oscillator networks.

Requires: jax>=0.4, equinox>=0.11

Classes

CouplingResidual 

CouplingResidual(hidden: int = 16, *, key: Array)

Bases: Module

Small MLP that learns the residual coupling function.

Maps phase difference Δθ → correction to sin(Δθ).

Source code in src/scpn_phase_orchestrator/nn/ude.py 
42
43
44
45
46
47
48
def __init__(self, hidden: int = 16, *, key: jax.Array) -> None:
    k1, k2, k3 = jax.random.split(key, 3)
    self.layers = [
        eqx.nn.Linear(1, hidden, key=k1),
        eqx.nn.Linear(hidden, hidden, key=k2),
        eqx.nn.Linear(hidden, 1, key=k3),
    ]
Methods:
__call__ 
__call__(delta_theta: Array) -> jax.Array

Evaluate residual for a single phase difference scalar.

The output is squashed to [-1, 1] with tanh: a physical coupling function of a phase difference is bounded (the sin backbone has magnitude ≤ 1), so the learned correction must be too. Without the bound the linear output head extrapolates without limit on phase differences unseen during training, and forward integration outside the training window diverges to NaN.

Source code in src/scpn_phase_orchestrator/nn/ude.py 
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
def __call__(self, delta_theta: jax.Array) -> jax.Array:
    """Evaluate residual for a single phase difference scalar.

    The output is squashed to ``[-1, 1]`` with ``tanh``: a physical coupling
    function of a phase difference is bounded (the ``sin`` backbone has
    magnitude ≤ 1), so the learned correction must be too. Without the bound
    the linear output head extrapolates without limit on phase differences
    unseen during training, and forward integration outside the training
    window diverges to NaN.
    """
    x = delta_theta[jnp.newaxis]  # (1,)
    x = jnp.tanh(self.layers[0](x))
    x = jnp.tanh(self.layers[1](x))
    x = jnp.tanh(self.layers[2](x))
    result: jax.Array = x[0]
    return result

UDEKuramotoLayer 

UDEKuramotoLayer(
    n: int,
    n_steps: int = 50,
    dt: float = 0.01,
    K_scale: float = 0.1,
    hidden: int = 16,
    *,
    key: Array,
)

Bases: Module

UDE-Kuramoto layer: physics backbone + learned residual.

Learnable parameters

K: (n, n) coupling matrix omegas: (n,) natural frequencies residual: CouplingResidual MLP

Source code in src/scpn_phase_orchestrator/nn/ude.py 
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
def __init__(
    self,
    n: int,
    n_steps: int = 50,
    dt: float = 0.01,
    K_scale: float = 0.1,
    hidden: int = 16,
    *,
    key: jax.Array,
) -> None:
    k1, k2, k3 = jax.random.split(key, 3)
    raw = K_scale * jax.random.normal(k1, (n, n))
    self.K = (raw + raw.T) / 2.0
    self.omegas = jax.random.normal(k2, (n,))
    self.residual = CouplingResidual(hidden=hidden, key=k3)
    self.n_steps = n_steps
    self.dt = dt
    self.n = n
Methods:
__call__ 
__call__(phases: Array) -> jax.Array

Integrate the UDE-Kuramoto forward map and return the final phases.

Source code in src/scpn_phase_orchestrator/nn/ude.py 
194
195
196
197
198
199
200
@eqx.filter_jit
def __call__(self, phases: jax.Array) -> jax.Array:
    """Integrate the UDE-Kuramoto forward map and return the final phases."""
    final, _ = ude_kuramoto_forward(
        phases, self.omegas, self.K, self.residual, self.dt, self.n_steps
    )
    return final
forward_with_trajectory 
forward_with_trajectory(
    phases: Array, *, backend: str = "euler"
) -> tuple[jax.Array, jax.Array]

Run dynamics and return (final_phases, trajectory).

The backend is validated here, outside the compiled region, so an invalid value fails fast with a plain Python error rather than a tracing-time exception; each backend body is a separately compiled helper.

Parameters

phases : jax.Array Oscillator phases in radians, shape (N,). backend : str Integration backend. "euler" (default) is the reproducible explicit jax.lax.scan map whose fixed grid keeps trajectory hashes stable; "diffrax" routes through :func:scpn_phase_orchestrator.nn.neural_ode.solve_ude_adjoint, an adaptive solver under a checkpointed continuous adjoint that samples the same n_steps grid, giving O(1)-memory training gradients. Requires the diffrax dependency.

Returns

tuple[jax.Array, jax.Array] The final phases (shape (N,)) and the trajectory (shape (n_steps, N)).

Raises

ValueError If backend is neither "euler" nor "diffrax".

Source code in src/scpn_phase_orchestrator/nn/ude.py 
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
def forward_with_trajectory(
    self, phases: jax.Array, *, backend: str = "euler"
) -> tuple[jax.Array, jax.Array]:
    """Run dynamics and return (final_phases, trajectory).

    The backend is validated here, outside the compiled region, so an
    invalid value fails fast with a plain Python error rather than a
    tracing-time exception; each backend body is a separately compiled
    helper.

    Parameters
    ----------
    phases : jax.Array
        Oscillator phases in radians, shape ``(N,)``.
    backend : str
        Integration backend. ``"euler"`` (default) is the reproducible
        explicit ``jax.lax.scan`` map whose fixed grid keeps trajectory
        hashes stable; ``"diffrax"`` routes through
        :func:`scpn_phase_orchestrator.nn.neural_ode.solve_ude_adjoint`,
        an adaptive solver under a checkpointed continuous adjoint that
        samples the same ``n_steps`` grid, giving ``O(1)``-memory training
        gradients. Requires the ``diffrax`` dependency.

    Returns
    -------
    tuple[jax.Array, jax.Array]
        The final phases (shape ``(N,)``) and the trajectory (shape
        ``(n_steps, N)``).

    Raises
    ------
    ValueError
        If ``backend`` is neither ``"euler"`` nor ``"diffrax"``.
    """
    if backend == "euler":
        return self._forward_trajectory_euler(phases)
    if backend == "diffrax":
        return self._forward_trajectory_diffrax(phases)
    raise ValueError("backend must be 'euler' or 'diffrax'")
sync_score 
sync_score(phases: Array) -> jax.Array

Kuramoto order parameter R after running the layer forward.

Parameters

phases : jax.Array Oscillator phases in radians, shape (N,).

Returns

jax.Array The Kuramoto order parameter R.

Source code in src/scpn_phase_orchestrator/nn/ude.py 
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
@eqx.filter_jit
def sync_score(self, phases: jax.Array) -> jax.Array:
    """Kuramoto order parameter R after running the layer forward.

    Parameters
    ----------
    phases : jax.Array
        Oscillator phases in radians, shape ``(N,)``.

    Returns
    -------
    jax.Array
        The Kuramoto order parameter ``R``.
    """
    return order_parameter(self(phases))

Functions:

ude_kuramoto_step 

ude_kuramoto_step(
    phases: Array,
    omegas: Array,
    K: Array,
    residual_fn: CouplingResidual,
    dt: float,
) -> jax.Array

Single Euler step of UDE-Kuramoto.

Parameters

phases : jax.Array (N,) oscillator phases. omegas : jax.Array (N,) natural frequencies. K : jax.Array (N, N) coupling matrix. residual_fn : CouplingResidual learned coupling correction. dt : float integration timestep.

Returns

jax.Array (N,) updated phases.

Source code in src/scpn_phase_orchestrator/nn/ude.py 
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
def ude_kuramoto_step(
    phases: jax.Array,
    omegas: jax.Array,
    K: jax.Array,
    residual_fn: CouplingResidual,
    dt: float,
) -> jax.Array:
    """Single Euler step of UDE-Kuramoto.

    Parameters
    ----------
    phases : jax.Array
        (N,) oscillator phases.
    omegas : jax.Array
        (N,) natural frequencies.
    K : jax.Array
        (N, N) coupling matrix.
    residual_fn : CouplingResidual
        learned coupling correction.
    dt : float
        integration timestep.

    Returns
    -------
    jax.Array
        (N,) updated phases.
    """
    dphi = _ude_deriv(phases, omegas, K, residual_fn)
    return (phases + dt * dphi) % TWO_PI

ude_kuramoto_forward 

ude_kuramoto_forward(
    phases: Array,
    omegas: Array,
    K: Array,
    residual_fn: CouplingResidual,
    dt: float,
    n_steps: int,
) -> tuple[jax.Array, jax.Array]

Run N steps of UDE-Kuramoto, returning final state and trajectory.

Parameters

phases : jax.Array (N,) initial phases. omegas : jax.Array (N,) natural frequencies. K : jax.Array (N, N) coupling matrix. residual_fn : CouplingResidual learned coupling correction. dt : float integration timestep. n_steps : int number of steps.

Returns

tuple[jax.Array, jax.Array] (final_phases, trajectory).

Source code in src/scpn_phase_orchestrator/nn/ude.py 
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
def ude_kuramoto_forward(
    phases: jax.Array,
    omegas: jax.Array,
    K: jax.Array,
    residual_fn: CouplingResidual,
    dt: float,
    n_steps: int,
) -> tuple[jax.Array, jax.Array]:
    """Run N steps of UDE-Kuramoto, returning final state and trajectory.

    Parameters
    ----------
    phases : jax.Array
        (N,) initial phases.
    omegas : jax.Array
        (N,) natural frequencies.
    K : jax.Array
        (N, N) coupling matrix.
    residual_fn : CouplingResidual
        learned coupling correction.
    dt : float
        integration timestep.
    n_steps : int
        number of steps.

    Returns
    -------
    tuple[jax.Array, jax.Array]
        (final_phases, trajectory).
    """

    def body(carry: jax.Array, _: None) -> tuple[jax.Array, jax.Array]:
        """Single scan iteration: UDE-Kuramoto step and record phase."""
        p = ude_kuramoto_step(carry, omegas, K, residual_fn, dt)
        return p, p

    final, trajectory = jax.lax.scan(body, phases, None, length=n_steps)
    return final, trajectory

Neural-ODE continuous adjoint (diffrax)

The explicit Euler map stores every step, so reverse-mode gradients cost O(n_steps) memory. solve_ude_adjoint integrates the same UDE-Kuramoto vector field with an adaptive solver (diffrax.Tsit5 by default) under a configurable adjoint — RecursiveCheckpointAdjoint (logarithmic checkpointing) or BacksolveAdjoint (O(1) memory). Integration runs on the unwrapped phase (the coupling is 2π-periodic, so the field is wrap-invariant while an adaptive solver must not see the % 2π discontinuities); wrapping is applied once, to the returned states. The solver never mutates the global jax_enable_x64 flag, so the dtype of every intermediate follows the input arrays.

solve_ude_adjoint(phases, omegas, K, residual_fn, *, t1, dt0=0.01,
                  solver=None, adjoint=None, rtol=1e-6, atol=1e-6,
                  max_steps=4096, saveat_ts=None, wrap=True)

Requires the diffrax dependency (the nn, jax, or full extra).

neural_ode

Continuous-time adjoint integration of the UDE-Kuramoto vector field.

The explicit jax.lax.scan Euler map in :mod:scpn_phase_orchestrator.nn.ude stores every intermediate state, so reverse-mode gradients cost O(n_steps) memory. This module integrates the same vector field

dθ_i/dt = ω_i + Σ_j K_ij · [sin(θ_j − θ_i) + NN_φ(θ_j − θ_i)]

with an adaptive higher-order solver (diffrax.Tsit5 by default) under a configurable adjoint. diffrax.RecursiveCheckpointAdjoint gives logarithmic checkpointing; diffrax.BacksolveAdjoint reconstructs the forward trajectory backwards for O(1) memory. Both differentiate through the coupling matrix K and the learned residual, so this is the production gradient path the finite-difference estimator in :mod:scpn_phase_orchestrator.upde.adjoint approximates.

The integration runs on the unwrapped phase: the coupling depends only on phase differences and sin is 2π-periodic, so the vector field is invariant to wrapping, while an adaptive solver must not see the % 2π discontinuities that the Euler map introduces at each step. Wrapping is applied once, to the returned states.

The dtype of every intermediate follows the input arrays — the solver never mutates the global jax_enable_x64 flag, so callers keep the float32 default of the rest of nn unless they opt into x64 themselves.

Requires: jax>=0.4, equinox>=0.11, diffrax>=0.5.

Classes

Functions:

solve_ude_adjoint 

solve_ude_adjoint(
    phases: Array,
    omegas: Array,
    K: Array,
    residual_fn: CouplingResidual,
    *,
    t1: float,
    dt0: float = 0.01,
    solver: AbstractSolver[Any] | None = None,
    adjoint: AbstractAdjoint | None = None,
    rtol: float = 1e-06,
    atol: float = 1e-06,
    max_steps: int = 4096,
    saveat_ts: Array | None = None,
    wrap: bool = True,
    throw: bool = True,
) -> jax.Array

Integrate UDE-Kuramoto with an adaptive solver and continuous adjoint.

Parameters

phases : jax.Array Initial oscillator phases in radians, shape (N,). omegas : jax.Array Natural frequencies in rad/s, shape (N,). K : jax.Array Coupling matrix, shape (N, N). residual_fn : CouplingResidual Learned per-pair coupling correction. t1 : float Final integration time; the interval is [0, t1]. Must be positive. dt0 : float Initial step size handed to the adaptive controller. Must be positive. solver : diffrax.AbstractSolver or None The ODE solver. Defaults to :class:diffrax.Tsit5 (5th-order adaptive). adjoint : diffrax.AbstractAdjoint or None The reverse-mode strategy. Defaults to :class:diffrax.RecursiveCheckpointAdjoint; pass :class:diffrax.BacksolveAdjoint for O(1) memory. rtol : float Relative tolerance for the PID step-size controller. Must be positive. atol : float Absolute tolerance for the PID step-size controller. Must be positive. max_steps : int Upper bound on solver steps. Must be positive. saveat_ts : jax.Array or None Times at which to save the trajectory. None saves only the final state and returns shape (N,); a length-T array returns shape (T, N). wrap : bool When True (default) the returned phases are wrapped into [0, 2π); when False the unwrapped phases are returned. throw : bool Stiffness guard. When True (default) a solve that exhausts max_steps — the symptom of a stiff or blowing-up field — raises instead of silently returning non-finite phases, so a diverging integration can never masquerade as a valid result. Set False to recover the non-finite solution for inspection (e.g. to locate the offending oscillator) rather than raising; raise max_steps or soften the field when this trips.

Returns

jax.Array The final phases (shape (N,)) when saveat_ts is None, else the saved trajectory (shape (T, N)).

Raises

ValueError If t1, dt0, rtol, atol or max_steps is not positive, or if phases is not one-dimensional.

Source code in src/scpn_phase_orchestrator/nn/neural_ode.py 
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
def solve_ude_adjoint(
    phases: jax.Array,
    omegas: jax.Array,
    K: jax.Array,
    residual_fn: CouplingResidual,
    *,
    t1: float,
    dt0: float = 0.01,
    solver: diffrax.AbstractSolver[Any] | None = None,
    adjoint: diffrax.AbstractAdjoint | None = None,
    rtol: float = 1e-6,
    atol: float = 1e-6,
    max_steps: int = 4096,
    saveat_ts: jax.Array | None = None,
    wrap: bool = True,
    throw: bool = True,
) -> jax.Array:
    """Integrate UDE-Kuramoto with an adaptive solver and continuous adjoint.

    Parameters
    ----------
    phases : jax.Array
        Initial oscillator phases in radians, shape ``(N,)``.
    omegas : jax.Array
        Natural frequencies in rad/s, shape ``(N,)``.
    K : jax.Array
        Coupling matrix, shape ``(N, N)``.
    residual_fn : CouplingResidual
        Learned per-pair coupling correction.
    t1 : float
        Final integration time; the interval is ``[0, t1]``. Must be positive.
    dt0 : float
        Initial step size handed to the adaptive controller. Must be positive.
    solver : diffrax.AbstractSolver or None
        The ODE solver. Defaults to :class:`diffrax.Tsit5` (5th-order adaptive).
    adjoint : diffrax.AbstractAdjoint or None
        The reverse-mode strategy. Defaults to
        :class:`diffrax.RecursiveCheckpointAdjoint`; pass
        :class:`diffrax.BacksolveAdjoint` for ``O(1)`` memory.
    rtol : float
        Relative tolerance for the PID step-size controller. Must be positive.
    atol : float
        Absolute tolerance for the PID step-size controller. Must be positive.
    max_steps : int
        Upper bound on solver steps. Must be positive.
    saveat_ts : jax.Array or None
        Times at which to save the trajectory. ``None`` saves only the final
        state and returns shape ``(N,)``; a length-``T`` array returns shape
        ``(T, N)``.
    wrap : bool
        When ``True`` (default) the returned phases are wrapped into
        ``[0, 2π)``; when ``False`` the unwrapped phases are returned.
    throw : bool
        Stiffness guard. When ``True`` (default) a solve that exhausts
        ``max_steps`` — the symptom of a stiff or blowing-up field — raises
        instead of silently returning non-finite phases, so a diverging
        integration can never masquerade as a valid result. Set ``False`` to
        recover the non-finite solution for inspection (e.g. to locate the
        offending oscillator) rather than raising; raise ``max_steps`` or soften
        the field when this trips.

    Returns
    -------
    jax.Array
        The final phases (shape ``(N,)``) when ``saveat_ts`` is ``None``, else
        the saved trajectory (shape ``(T, N)``).

    Raises
    ------
    ValueError
        If ``t1``, ``dt0``, ``rtol``, ``atol`` or ``max_steps`` is not
        positive, or if ``phases`` is not one-dimensional.
    """
    if phases.ndim != 1:
        raise ValueError("phases must be a one-dimensional array")
    if t1 <= 0.0:
        raise ValueError("t1 must be positive")
    if dt0 <= 0.0:
        raise ValueError("dt0 must be positive")
    if rtol <= 0.0 or atol <= 0.0:
        raise ValueError("rtol and atol must be positive")
    if max_steps <= 0:
        raise ValueError("max_steps must be positive")

    active_solver = diffrax.Tsit5() if solver is None else solver
    active_adjoint = (
        diffrax.RecursiveCheckpointAdjoint() if adjoint is None else adjoint
    )
    if saveat_ts is None:
        saveat = diffrax.SaveAt(t1=True)
    else:
        saveat = diffrax.SaveAt(ts=saveat_ts)

    solution = diffrax.diffeqsolve(
        diffrax.ODETerm(_ude_vector_field),
        active_solver,
        t0=0.0,
        t1=t1,
        dt0=dt0,
        y0=phases,
        args=(omegas, K, residual_fn),
        saveat=saveat,
        stepsize_controller=diffrax.PIDController(rtol=rtol, atol=atol),
        adjoint=active_adjoint,
        max_steps=max_steps,
        throw=throw,
    )

    saved = solution.ys
    states = saved if saveat_ts is not None else saved[-1]
    if wrap:
        wrapped: jax.Array = states % TWO_PI
        return wrapped
    result: jax.Array = states
    return result

Inverse Kuramoto

Gradient-based inference of K and ω from observed phase trajectories.

Function Description
infer_coupling(observed, dt, n_epochs, lr, ...) Full gradient descent inference
analytical_inverse(observed, dt, alpha) Closed-form least-squares
hybrid_inverse(observed, dt, ...) Analytical + gradient refinement
inverse_loss(K, omegas, observed, dt, l1) Differentiable loss
coupling_correlation(K_true, K_inferred) Pearson r for validation

inverse

Infer coupling matrix K and natural frequencies ω from observed phases.

Three methods, in order of preference:

  1. analytical_inverse (Pikovsky 2008) — O(N³) linear regression on sin(Δθ) basis functions. Exact for noiseless Kuramoto, >0.95 correlation, completes in seconds. Use this by default.

  2. hybrid_inverse — analytical init + gradient refinement. Handles model mismatch (noise, higher harmonics) by starting from the analytical solution and running a few Adam epochs.

  3. infer_coupling — pure gradient descent through ODE solver. Kept for backward compatibility. Slow (minutes), lower accuracy.

Requires: jax>=0.4

Functions:

inverse_loss 

inverse_loss(
    K: Array,
    omegas: Array,
    observed: Array,
    dt: float,
    l1_weight: float = 0.0,
) -> jax.Array

Loss for inverse Kuramoto: prediction error + optional L1 sparsity.

Runs the forward model from observed[0] and compares the predicted trajectory against the observed trajectory.

Parameters

K : jax.Array (N, N) coupling matrix to optimize. omegas : jax.Array (N,) natural frequencies to optimize. observed : jax.Array (T, N) observed phase trajectory. dt : float integration timestep. l1_weight : float L1 penalty on K for sparsity (0 = no penalty).

Returns

jax.Array Scalar loss.

Source code in src/scpn_phase_orchestrator/nn/inverse.py 
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
def inverse_loss(
    K: jax.Array,
    omegas: jax.Array,
    observed: jax.Array,
    dt: float,
    l1_weight: float = 0.0,
) -> jax.Array:
    """Loss for inverse Kuramoto: prediction error + optional L1 sparsity.

    Runs the forward model from observed[0] and compares the predicted
    trajectory against the observed trajectory.

    Parameters
    ----------
    K : jax.Array
        (N, N) coupling matrix to optimize.
    omegas : jax.Array
        (N,) natural frequencies to optimize.
    observed : jax.Array
        (T, N) observed phase trajectory.
    dt : float
        integration timestep.
    l1_weight : float
        L1 penalty on K for sparsity (0 = no penalty).

    Returns
    -------
    jax.Array
        Scalar loss.
    """
    n_steps = observed.shape[0] - 1
    initial = observed[0]

    _, predicted = kuramoto_forward(initial, omegas, K, dt, n_steps)

    diff = observed[1:] - predicted
    phase_error = jnp.mean(1.0 - jnp.cos(diff))

    loss = phase_error
    if l1_weight > 0.0:
        loss = loss + l1_weight * jnp.sum(jnp.abs(K))
    return loss

analytical_inverse 

analytical_inverse(
    observed: Array, dt: float, alpha: float = 0.0
) -> tuple[jax.Array, jax.Array]

Recover K and ω from observed phases via linear regression.

Exploits the Kuramoto structure directly (Pikovsky 2008): dθ_i/dt = ω_i + Σ_j K_ij sin(θ_j - θ_i)

Finite-difference dθ/dt, build sin(Δθ) basis, solve via lstsq. O(N³) per oscillator, no ODE backprop, no gradient vanishing.

Parameters

observed : jax.Array (T, N) phase trajectory, T >= 3. dt : float integration timestep. alpha : float Tikhonov (ridge) regularisation strength. 0 = no reg.

Returns

tuple[jax.Array, jax.Array] (K, omegas): inferred (N, N) coupling and (N,) frequencies.

Source code in src/scpn_phase_orchestrator/nn/inverse.py 
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
def analytical_inverse(
    observed: jax.Array,
    dt: float,
    alpha: float = 0.0,
) -> tuple[jax.Array, jax.Array]:
    """Recover K and ω from observed phases via linear regression.

    Exploits the Kuramoto structure directly (Pikovsky 2008):
      dθ_i/dt = ω_i + Σ_j K_ij sin(θ_j - θ_i)

    Finite-difference dθ/dt, build sin(Δθ) basis, solve via lstsq.
    O(N³) per oscillator, no ODE backprop, no gradient vanishing.

    Parameters
    ----------
    observed : jax.Array
        (T, N) phase trajectory, T >= 3.
    dt : float
        integration timestep.
    alpha : float
        Tikhonov (ridge) regularisation strength. 0 = no reg.

    Returns
    -------
    tuple[jax.Array, jax.Array]
        (K, omegas): inferred (N, N) coupling and (N,) frequencies.
    """
    T, N = observed.shape
    # Phase-aware central finite differences: unwrap Δθ via atan2
    # to handle 2π boundary crossings correctly
    raw_diff = observed[2:] - observed[:-2]
    dtheta_dt = jnp.arctan2(jnp.sin(raw_diff), jnp.cos(raw_diff)) / (2.0 * dt)
    phases_mid = observed[1:-1]  # (T_mid, N)

    # Build 3D basis: B_all[i, t, j] = sin(θ_j(t) - θ_i(t))
    # phases_mid[:, :, None] - phases_mid[:, None, :] → (T_mid, N, N)
    # then transpose to (N, T_mid, N) for per-oscillator solve
    diff_3d = phases_mid[:, jnp.newaxis, :] - phases_mid[:, :, jnp.newaxis]
    sin_basis = jnp.sin(diff_3d).transpose(1, 0, 2)  # (N, T_mid, N)
    targets = dtheta_dt.T  # (N, T_mid)

    # Augment each per-oscillator design with an intercept column so ω_i is
    # estimated jointly with its coupling row rather than as a post-hoc residual.
    # Without the intercept the lstsq absorbs the ω-driven phase drift into the
    # sin(Δθ) basis, inflating K for weakly coupled or uncoupled data (the
    # ω/coupling confounding). The intercept removes that bias: uncoupled data
    # recovers K ≈ 0 while coupled recovery is unchanged.
    n_mid = sin_basis.shape[1]
    design = jnp.concatenate([sin_basis, jnp.ones((N, n_mid, 1))], axis=2)

    if alpha > 0:
        # Ridge-penalise the coupling block only; the intercept (ω) is unpenalised.
        reg = alpha * jnp.diag(jnp.concatenate([jnp.ones(N), jnp.zeros(1)]))

        def _solve(design_i: jax.Array, target: jax.Array) -> jax.Array:
            """Solve the inverse problem and return the result."""
            solution_i: jax.Array = jnp.linalg.solve(
                design_i.T @ design_i + reg, design_i.T @ target
            )
            return solution_i

        solution = jax.vmap(_solve)(design, targets)
    else:

        def _solve(design_i: jax.Array, target: jax.Array) -> jax.Array:
            """Solve the inverse problem and return the result."""
            row, _, _, _ = jnp.linalg.lstsq(design_i, target)
            return row

        solution = jax.vmap(_solve)(design, targets)

    K = _symmetrise_K(solution[:, :N])
    omegas = solution[:, N]
    return K, omegas

hybrid_inverse 

hybrid_inverse(
    observed: Array,
    dt: float,
    alpha: float = 0.0,
    n_refine: int = 50,
    lr: float = 0.005,
    window_size: int = 10,
) -> tuple[jax.Array, jax.Array, list[float]]

Analytical inverse + gradient refinement for noisy data.

Runs analytical_inverse() for the initial estimate, then refines with a few Adam epochs using multiple shooting. Handles model mismatch (noise, higher harmonics, amplitude effects).

Parameters

observed : jax.Array (T, N) phase trajectory. dt : float integration timestep. alpha : float Tikhonov regularisation for analytical step. n_refine : int Adam refinement epochs (0 = analytical only). lr : float learning rate for refinement. window_size : int shooting window size for refinement.

Returns

tuple[jax.Array, jax.Array, list[float]] (K, omegas, losses): inferred params + refinement loss history.

Source code in src/scpn_phase_orchestrator/nn/inverse.py 
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
def hybrid_inverse(
    observed: jax.Array,
    dt: float,
    alpha: float = 0.0,
    n_refine: int = 50,
    lr: float = 0.005,
    window_size: int = 10,
) -> tuple[jax.Array, jax.Array, list[float]]:
    """Analytical inverse + gradient refinement for noisy data.

    Runs analytical_inverse() for the initial estimate, then refines
    with a few Adam epochs using multiple shooting. Handles model
    mismatch (noise, higher harmonics, amplitude effects).

    Parameters
    ----------
    observed : jax.Array
        (T, N) phase trajectory.
    dt : float
        integration timestep.
    alpha : float
        Tikhonov regularisation for analytical step.
    n_refine : int
        Adam refinement epochs (0 = analytical only).
    lr : float
        learning rate for refinement.
    window_size : int
        shooting window size for refinement.

    Returns
    -------
    tuple[jax.Array, jax.Array, list[float]]
        (K, omegas, losses): inferred params + refinement loss history.
    """
    K, omegas = analytical_inverse(observed, dt, alpha=alpha)

    if n_refine <= 0:
        return K, omegas, []

    starts, targets = _build_windows(observed, window_size)

    def loss_fn(k: jax.Array, o: jax.Array) -> jax.Array:
        """Shooting loss for refinement step."""
        return _shooting_loss(k, o, starts, targets, dt, window_size, 0.0)

    loss_and_grad = jax.value_and_grad(loss_fn, argnums=(0, 1))

    m_K = jnp.zeros_like(K)
    v_K = jnp.zeros_like(K)
    m_o = jnp.zeros_like(omegas)
    v_o = jnp.zeros_like(omegas)
    beta1, beta2, eps = 0.9, 0.999, 1e-8
    losses: list[float] = []

    for epoch in range(n_refine):
        loss_val, (grad_K, grad_o) = loss_and_grad(K, omegas)
        g_norm = jnp.sqrt(jnp.sum(grad_K**2) + jnp.sum(grad_o**2) + 1e-10)
        scale = jnp.minimum(1.0, 1.0 / g_norm)
        grad_K = grad_K * scale
        grad_o = grad_o * scale

        t = epoch + 1
        m_K = beta1 * m_K + (1 - beta1) * grad_K
        v_K = beta2 * v_K + (1 - beta2) * grad_K**2
        m_o = beta1 * m_o + (1 - beta1) * grad_o
        v_o = beta2 * v_o + (1 - beta2) * grad_o**2
        bc1 = 1 - beta1**t
        bc2 = 1 - beta2**t
        K = K - lr * (m_K / bc1) / (jnp.sqrt(v_K / bc2) + eps)
        omegas = omegas - lr * (m_o / bc1) / (jnp.sqrt(v_o / bc2) + eps)
        K = _symmetrise_K(K)
        losses.append(float(loss_val))

    return K, omegas, losses

infer_coupling 

infer_coupling(
    observed: Array,
    dt: float,
    n_epochs: int = 200,
    lr: float = 0.01,
    l1_weight: float = 0.001,
    seed: int = 0,
    window_size: int = 0,
    grad_clip: float = 1.0,
) -> tuple[jax.Array, jax.Array, list[float]]

Infer coupling matrix K and frequencies ω from observed phases.

Uses Adam optimiser with gradient clipping and optional multiple shooting for gradient-stable training through ODE solvers.

Parameters

observed : jax.Array (T, N) observed phase trajectory. dt : float integration timestep used to generate the data. n_epochs : int optimisation epochs. lr : float learning rate (for Adam). l1_weight : float L1 sparsity penalty on K. seed : int random seed for initialisation. window_size : int if >0, use multiple shooting with this window size. Recommended: 10-20 steps. 0 = single-shot (original behaviour). grad_clip : float maximum gradient norm (0 = no clipping).

Returns

tuple[jax.Array, jax.Array, list[float]] (K, omegas, losses) where: K: (N, N) inferred coupling matrix omegas: (N,) inferred natural frequencies losses: list of loss values per epoch.

Source code in src/scpn_phase_orchestrator/nn/inverse.py 
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
def infer_coupling(
    observed: jax.Array,
    dt: float,
    n_epochs: int = 200,
    lr: float = 0.01,
    l1_weight: float = 0.001,
    seed: int = 0,
    window_size: int = 0,
    grad_clip: float = 1.0,
) -> tuple[jax.Array, jax.Array, list[float]]:
    """Infer coupling matrix K and frequencies ω from observed phases.

    Uses Adam optimiser with gradient clipping and optional multiple
    shooting for gradient-stable training through ODE solvers.

    Parameters
    ----------
    observed : jax.Array
        (T, N) observed phase trajectory.
    dt : float
        integration timestep used to generate the data.
    n_epochs : int
        optimisation epochs.
    lr : float
        learning rate (for Adam).
    l1_weight : float
        L1 sparsity penalty on K.
    seed : int
        random seed for initialisation.
    window_size : int
        if >0, use multiple shooting with this window size. Recommended: 10-20 steps. 0
        = single-shot (original behaviour).
    grad_clip : float
        maximum gradient norm (0 = no clipping).

    Returns
    -------
    tuple[jax.Array, jax.Array, list[float]]
        (K, omegas, losses) where: K: (N, N) inferred coupling matrix omegas: (N,)
        inferred natural frequencies losses: list of loss values per epoch.
    """
    N = observed.shape[1]
    key = jax.random.PRNGKey(seed)
    k1, _ = jax.random.split(key)

    K = jax.random.normal(k1, (N, N)) * 0.05
    K = _symmetrise_K(K)
    omegas = jnp.zeros(N)

    # Adam state
    m_K = jnp.zeros_like(K)
    v_K = jnp.zeros_like(K)
    m_o = jnp.zeros_like(omegas)
    v_o = jnp.zeros_like(omegas)
    beta1, beta2, eps = 0.9, 0.999, 1e-8

    if window_size > 0:
        starts, targets = _build_windows(observed, window_size)

        def loss_fn(k: jax.Array, o: jax.Array) -> jax.Array:
            """Multiple-shooting loss with L1 penalty."""
            return _shooting_loss(k, o, starts, targets, dt, window_size, l1_weight)
    else:

        def loss_fn(k: jax.Array, o: jax.Array) -> jax.Array:
            """Single-shot inverse loss with L1 penalty."""
            return inverse_loss(k, o, observed, dt, l1_weight)

    loss_and_grad = jax.value_and_grad(loss_fn, argnums=(0, 1))
    losses: list[float] = []

    for epoch in range(n_epochs):
        loss_val, (grad_K, grad_o) = loss_and_grad(K, omegas)

        # Gradient clipping
        if grad_clip > 0:
            g_norm = jnp.sqrt(jnp.sum(grad_K**2) + jnp.sum(grad_o**2) + 1e-10)
            scale = jnp.minimum(1.0, grad_clip / g_norm)
            grad_K = grad_K * scale
            grad_o = grad_o * scale

        # Adam update
        t = epoch + 1
        m_K = beta1 * m_K + (1 - beta1) * grad_K
        v_K = beta2 * v_K + (1 - beta2) * grad_K**2
        m_o = beta1 * m_o + (1 - beta1) * grad_o
        v_o = beta2 * v_o + (1 - beta2) * grad_o**2

        bc1 = 1 - beta1**t
        bc2 = 1 - beta2**t
        K = K - lr * (m_K / bc1) / (jnp.sqrt(v_K / bc2) + eps)
        omegas = omegas - lr * (m_o / bc1) / (jnp.sqrt(v_o / bc2) + eps)

        K = _symmetrise_K(K)
        losses.append(float(loss_val))

    return K, omegas, losses

coupling_correlation 

coupling_correlation(
    K_true: Array, K_inferred: Array
) -> jax.Array

Pearson correlation between true and inferred coupling matrices.

Parameters

K_true : jax.Array (N, N) ground truth coupling. K_inferred : jax.Array (N, N) inferred coupling.

Returns

jax.Array Scalar correlation in [-1, 1].

Source code in src/scpn_phase_orchestrator/nn/inverse.py 
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
def coupling_correlation(K_true: jax.Array, K_inferred: jax.Array) -> jax.Array:
    """Pearson correlation between true and inferred coupling matrices.

    Parameters
    ----------
    K_true : jax.Array
        (N, N) ground truth coupling.
    K_inferred : jax.Array
        (N, N) inferred coupling.

    Returns
    -------
    jax.Array
        Scalar correlation in [-1, 1].
    """
    # Flatten upper triangle (exclude diagonal)
    N = K_true.shape[0]
    idx = jnp.triu_indices(N, k=1)
    a = K_true[idx]
    b = K_inferred[idx]
    a_centered = a - jnp.mean(a)
    b_centered = b - jnp.mean(b)
    num = jnp.sum(a_centered * b_centered)
    denom = jnp.sqrt(jnp.sum(a_centered**2) * jnp.sum(b_centered**2) + 1e-10)
    result: jax.Array = num / denom
    return result

Oscillator Ising Machine (OIM)

Combinatorial optimisation via phase clustering. Maps graph colouring, max-cut, and QUBO to Kuramoto dynamics.

Function Description
oim_solve(adj, n_colors, key, ...) Full solver with annealing + restarts
oim_forward(phases, adj, n_colors, dt, n_steps) Forward integration
extract_coloring(phases, n_colors) Hard colour assignment
coloring_violations(colors, adj) Count constraint violations
coloring_energy(phases, adj, n_colors) Continuous energy

First open-source OIM simulator.

oim

Kuramoto-based combinatorial optimization via phase clustering.

Maps NP-hard problems (graph coloring, max-cut, QUBO) to coupled oscillator dynamics. Oscillators settle into k distinct phase clusters, each cluster corresponding to a color/partition.

The coupling function is modified from standard sin(Δθ) to produce equidistant phase clusters (Nature Scientific Reports 2017, Böhm & Schumacher 2020). GPU-accelerated via JAX.

First open-source oscillator Ising machine simulator.

Requires: jax>=0.4

Functions:

oim_step 

oim_step(
    phases: Array,
    adjacency: Array,
    n_colors: int,
    dt: float,
    coupling_strength: float = 1.0,
) -> jax.Array

Single step of OIM coloring dynamics.

Parameters

phases : jax.Array (N,) oscillator phases. adjacency : jax.Array (N, N) graph adjacency matrix (1 = edge, 0 = no edge). n_colors : int number of colors (phase clusters). dt : float integration timestep. coupling_strength : float overall coupling scale.

Returns

jax.Array (N,) updated phases.

Source code in src/scpn_phase_orchestrator/nn/oim.py 
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
def oim_step(
    phases: jax.Array,
    adjacency: jax.Array,
    n_colors: int,
    dt: float,
    coupling_strength: float = 1.0,
) -> jax.Array:
    """Single step of OIM coloring dynamics.

    Parameters
    ----------
    phases : jax.Array
        (N,) oscillator phases.
    adjacency : jax.Array
        (N, N) graph adjacency matrix (1 = edge, 0 = no edge).
    n_colors : int
        number of colors (phase clusters).
    dt : float
        integration timestep.
    coupling_strength : float
        overall coupling scale.

    Returns
    -------
    jax.Array
        (N,) updated phases.
    """
    dphi = _oim_deriv(phases, adjacency, n_colors, coupling_strength)
    return (phases + dt * dphi) % TWO_PI

oim_forward 

oim_forward(
    phases: Array,
    adjacency: Array,
    n_colors: int,
    dt: float,
    n_steps: int,
    coupling_strength: float = 1.0,
) -> tuple[jax.Array, jax.Array]

Run OIM dynamics for n_steps, returning final phases and trajectory.

Parameters

phases : jax.Array (N,) initial random phases. adjacency : jax.Array (N, N) graph adjacency matrix. n_colors : int number of colors. dt : float timestep. n_steps : int number of integration steps. coupling_strength : float overall coupling scale.

Returns

tuple[jax.Array, jax.Array] (final_phases, trajectory) where trajectory is (n_steps, N).

Source code in src/scpn_phase_orchestrator/nn/oim.py 
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
def oim_forward(
    phases: jax.Array,
    adjacency: jax.Array,
    n_colors: int,
    dt: float,
    n_steps: int,
    coupling_strength: float = 1.0,
) -> tuple[jax.Array, jax.Array]:
    """Run OIM dynamics for n_steps, returning final phases and trajectory.

    Parameters
    ----------
    phases : jax.Array
        (N,) initial random phases.
    adjacency : jax.Array
        (N, N) graph adjacency matrix.
    n_colors : int
        number of colors.
    dt : float
        timestep.
    n_steps : int
        number of integration steps.
    coupling_strength : float
        overall coupling scale.

    Returns
    -------
    tuple[jax.Array, jax.Array]
        (final_phases, trajectory) where trajectory is (n_steps, N).
    """

    def body(carry: jax.Array, _: None) -> tuple[jax.Array, jax.Array]:
        """Return the loop body for the iteration."""
        p = oim_step(carry, adjacency, n_colors, dt, coupling_strength)
        return p, p

    final, trajectory = jax.lax.scan(body, phases, None, length=n_steps)
    return final, trajectory

extract_coloring 

extract_coloring(phases: Array, n_colors: int) -> jax.Array

Extract integer color assignment from oscillator phases.

Maps each phase to the nearest cluster center at 2πk/n_colors.

Parameters

phases : jax.Array (N,) oscillator phases in [0, 2π). n_colors : int number of colors.

Returns

jax.Array (N,) integer colour labels in {0, 1, ..., n_colors-1}.

Source code in src/scpn_phase_orchestrator/nn/oim.py 
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
def extract_coloring(phases: jax.Array, n_colors: int) -> jax.Array:
    """Extract integer color assignment from oscillator phases.

    Maps each phase to the nearest cluster center at 2πk/n_colors.

    Parameters
    ----------
    phases : jax.Array
        (N,) oscillator phases in [0, 2π).
    n_colors : int
        number of colors.

    Returns
    -------
    jax.Array
        (N,) integer colour labels in {0, 1, ..., n_colors-1}.
    """
    # Cluster centers at 2πk/n_colors
    bucket_size = TWO_PI / n_colors
    result: jax.Array = jnp.floor(phases / bucket_size).astype(jnp.int32) % n_colors
    return result

extract_coloring_soft 

extract_coloring_soft(
    phases: Array, n_colors: int
) -> jax.Array

Extract color assignment using circular distance to cluster centres.

More accurate than floor bucketing when phases sit near bucket boundaries. Assigns each oscillator to the nearest of the n_colors equidistant cluster centres.

Parameters

phases : jax.Array (N,) oscillator phases in [0, 2π). n_colors : int number of colors.

Returns

jax.Array (N,) integer colour labels in {0, 1, ..., n_colors-1}.

Source code in src/scpn_phase_orchestrator/nn/oim.py 
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
def extract_coloring_soft(phases: jax.Array, n_colors: int) -> jax.Array:
    """Extract color assignment using circular distance to cluster centres.

    More accurate than floor bucketing when phases sit near bucket
    boundaries. Assigns each oscillator to the nearest of the n_colors
    equidistant cluster centres.

    Parameters
    ----------
    phases : jax.Array
        (N,) oscillator phases in [0, 2π).
    n_colors : int
        number of colors.

    Returns
    -------
    jax.Array
        (N,) integer colour labels in {0, 1, ..., n_colors-1}.
    """
    centres = jnp.linspace(0, TWO_PI, n_colors, endpoint=False)
    # Circular distance: |angle_diff| wrapped to [-π, π]
    diff = phases[:, jnp.newaxis] - centres[jnp.newaxis, :]
    circ_dist = jnp.abs(jnp.arctan2(jnp.sin(diff), jnp.cos(diff)))
    result: jax.Array = jnp.argmin(circ_dist, axis=1).astype(jnp.int32)
    return result

oim_solve 

oim_solve(
    adjacency: Array,
    n_colors: int,
    *,
    key: Array,
    dt: float = 0.05,
    k_min: float = 0.1,
    k_max: float = 10.0,
    n_anneal: int = 1000,
    n_refine: int = 500,
    n_restarts: int = 10,
) -> tuple[jax.Array, jax.Array, float]

Solve graph coloring via OIM with annealing and multi-start.

Fully vectorised: restarts run in parallel via vmap, annealing and refinement use jax.lax.scan (no Python loops). 70x faster than the sequential version on GPU.

Parameters

adjacency : jax.Array (N, N) graph adjacency matrix. n_colors : int number of colors. key : jax.Array PRNG key. dt : float Requested integration timestep, used as an upper bound. The effective step is reduced when k_max * coupling_n * max_degree * dt would exceed the explicit-Euler stability radius, so the dynamics settle into the ground state instead of overshooting it. k_min : float initial coupling strength (low = exploration). k_max : float final coupling strength (high = exploitation). n_anneal : int ramp-up steps. n_refine : int hold steps after annealing. n_restarts : int number of random restarts.

Returns

tuple[jax.Array, jax.Array, float] (best_colors, best_phases, best_energy).

Source code in src/scpn_phase_orchestrator/nn/oim.py 
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
def oim_solve(
    adjacency: jax.Array,
    n_colors: int,
    *,
    key: jax.Array,
    dt: float = 0.05,
    k_min: float = 0.1,
    k_max: float = 10.0,
    n_anneal: int = 1000,
    n_refine: int = 500,
    n_restarts: int = 10,
) -> tuple[jax.Array, jax.Array, float]:
    """Solve graph coloring via OIM with annealing and multi-start.

    Fully vectorised: restarts run in parallel via vmap, annealing and
    refinement use jax.lax.scan (no Python loops). 70x faster than the
    sequential version on GPU.

    Parameters
    ----------
    adjacency : jax.Array
        (N, N) graph adjacency matrix.
    n_colors : int
        number of colors.
    key : jax.Array
        PRNG key.
    dt : float
        Requested integration timestep, used as an upper bound. The effective
        step is reduced when ``k_max * coupling_n * max_degree * dt`` would
        exceed the explicit-Euler stability radius, so the dynamics settle into
        the ground state instead of overshooting it.
    k_min : float
        initial coupling strength (low = exploration).
    k_max : float
        final coupling strength (high = exploitation).
    n_anneal : int
        ramp-up steps.
    n_refine : int
        hold steps after annealing.
    n_restarts : int
        number of random restarts.

    Returns
    -------
    tuple[jax.Array, jax.Array, float]
        (best_colors, best_phases, best_energy).
    """
    N = adjacency.shape[0]
    # For 2-coloring, sin(2*Δθ) equilibrium is at π/2 (between
    # cluster centres), so use sin(Δθ) coupling (anti-phase at π).
    coupling_n = 1 if n_colors == 2 else n_colors

    # Forward-Euler stability. The linearised gain of the sin(coupling_n·Δθ)
    # coupling is k · coupling_n · (weighted node degree); an explicit step
    # beyond the stability radius overshoots and stalls the dynamics well above
    # the ground state (dt=0.05, k_max=10 on K_{3,3} stalls at E≈-0.67 instead
    # of the optimum -9). Cap the effective step so the peak-coupling gain stays
    # within a safe fraction of that radius, adapting to the graph's degree.
    max_degree = jnp.maximum(jnp.max(jnp.sum(jnp.abs(adjacency), axis=1)), 1.0)
    stable_dt = 0.5 / (k_max * coupling_n * max_degree)
    dt_eff = jnp.minimum(dt, stable_dt)

    # Precompute annealing schedules
    frac = jnp.linspace(0.0, 1.0, n_anneal)
    k_schedule = k_min + (k_max - k_min) * frac
    noise_schedule = 0.1 * (1.0 - frac)

    def _single_restart(restart_key: jax.Array) -> tuple[jax.Array, jax.Array]:
        """Run one anneal+refine restart, return (phases, violations)."""
        init_key, noise_base_key = jax.random.split(restart_key)
        phases0 = jax.random.uniform(init_key, (N,), maxval=TWO_PI)
        # Pre-split noise keys for all anneal steps
        noise_keys = jax.random.split(noise_base_key, n_anneal)

        # Annealing via scan
        def anneal_body(phases: jax.Array, xs: Any) -> tuple[jax.Array, None]:
            """Return the annealing-phase loop body."""
            k, ns, nk = xs
            dphi = _oim_deriv(phases, adjacency, coupling_n, k)
            phases = (phases + dt_eff * dphi) % TWO_PI
            noise = ns * jax.random.normal(nk, (N,))
            phases = (phases + noise) % TWO_PI
            return phases, None

        phases, _ = jax.lax.scan(
            anneal_body,
            phases0,
            (k_schedule, noise_schedule, noise_keys),
        )

        # Refinement via scan (fixed coupling, no noise)
        def refine_body(phases: jax.Array, _: Any) -> tuple[jax.Array, None]:
            """Return the refinement-phase loop body."""
            dphi = _oim_deriv(phases, adjacency, coupling_n, k_max)
            return (phases + dt_eff * dphi) % TWO_PI, None

        phases, _ = jax.lax.scan(refine_body, phases, None, length=n_refine)

        colors = extract_coloring_soft(phases, n_colors)
        v = coloring_violations(colors, adjacency)
        return phases, v

    # vmap across all restarts
    restart_keys = jax.random.split(key, n_restarts)
    all_phases, all_violations = jax.vmap(_single_restart)(restart_keys)

    # Select best restart (lowest violations)
    best_idx = jnp.argmin(all_violations)
    best_phases = all_phases[best_idx]
    best_colors = extract_coloring_soft(best_phases, n_colors)
    best_energy = float(coloring_energy(best_phases, adjacency, coupling_n))

    return best_colors, best_phases, best_energy

coloring_violations 

coloring_violations(
    colors: Array, adjacency: Array
) -> jax.Array

Count edges where both endpoints have the same color.

Parameters

colors : jax.Array (N,) integer colour labels. adjacency : jax.Array (N, N) adjacency matrix.

Returns

jax.Array Scalar: number of violated edges.

Source code in src/scpn_phase_orchestrator/nn/oim.py 
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
def coloring_violations(
    colors: jax.Array,
    adjacency: jax.Array,
) -> jax.Array:
    """Count edges where both endpoints have the same color.

    Parameters
    ----------
    colors : jax.Array
        (N,) integer colour labels.
    adjacency : jax.Array
        (N, N) adjacency matrix.

    Returns
    -------
    jax.Array
        Scalar: number of violated edges.
    """
    same_color = (colors[jnp.newaxis, :] == colors[:, jnp.newaxis]).astype(jnp.float32)
    # Count upper triangle only (each edge once)
    violations = jnp.sum(jnp.triu(adjacency * same_color, k=1))
    result: jax.Array = violations
    return result

coloring_energy 

coloring_energy(
    phases: Array, adjacency: Array, n_colors: int
) -> jax.Array

Continuous energy function for the coloring problem.

E = Σ_{(i,j)∈E} cos(n_colors · (θ_i - θ_j))

Minimized when connected nodes are in different phase clusters (cos(n·Δθ) = -1 when Δθ = π/n, i.e., maximally separated). Differentiable for gradient-based optimization.

Parameters

phases : jax.Array (N,) oscillator phases. adjacency : jax.Array (N, N) adjacency matrix. n_colors : int number of colors.

Returns

jax.Array Scalar energy (lower = better colouring).

Source code in src/scpn_phase_orchestrator/nn/oim.py 
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
def coloring_energy(
    phases: jax.Array,
    adjacency: jax.Array,
    n_colors: int,
) -> jax.Array:
    """Continuous energy function for the coloring problem.

    E = Σ_{(i,j)∈E} cos(n_colors · (θ_i - θ_j))

    Minimized when connected nodes are in different phase clusters
    (cos(n·Δθ) = -1 when Δθ = π/n, i.e., maximally separated).
    Differentiable for gradient-based optimization.

    Parameters
    ----------
    phases : jax.Array
        (N,) oscillator phases.
    adjacency : jax.Array
        (N, N) adjacency matrix.
    n_colors : int
        number of colors.

    Returns
    -------
    jax.Array
        Scalar energy (lower = better colouring).
    """
    diff = phases[jnp.newaxis, :] - phases[:, jnp.newaxis]
    cost_matrix = jnp.cos(n_colors * diff)
    energy: jax.Array = jnp.sum(jnp.triu(adjacency * cost_matrix, k=1))
    return energy

Training Utilities

Loss functions

Function Description
sync_loss(model, phases, target_R=1.0) (1 - R)² loss
trajectory_loss(model, phases, observed) MSE on phase trajectory
coupling_sparsity_loss(K, target_density) L1 sparsity penalty

Training loop

from scpn_phase_orchestrator.nn.training import train

model, losses = train(
    model=layer,
    loss_fn=lambda m: sync_loss(m, phases),
    optimizer=optax.adam(1e-3),
    n_epochs=500,
    callback=lambda ep, m, l: print(f"Epoch {ep}: loss={float(l):.4f}"),
)

Data generation

Function Returns
generate_kuramoto_data(N, T, dt, K_scale, key) (K, omegas, phases_init, trajectory)
generate_chimera_data(N, T, dt, coupling, range, key) (K, omegas, trajectory)

training

Loss functions, training loops, and data generation for nn/ layers.

Integrates with optax for optimisation. All loss functions are differentiable via JAX autodiff and compatible with equinox modules.

Requires: jax>=0.4, equinox>=0.11, optax>=0.2

Functions:

sync_loss 

sync_loss(
    model: Module, phases: Array, target_R: float = 1.0
) -> jax.Array

Drive oscillator layer toward a target synchronisation level.

Parameters

model : eqx.Module equinox layer with call(phases) → final_phases. phases : jax.Array (N,) initial phases. target_R : float target order parameter R in [0, 1].

Returns

jax.Array Scalar loss (R - target_R)^2.

Source code in src/scpn_phase_orchestrator/nn/training.py 
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
def sync_loss(
    model: eqx.Module,
    phases: jax.Array,
    target_R: float = 1.0,
) -> jax.Array:
    """Drive oscillator layer toward a target synchronisation level.

    Parameters
    ----------
    model : eqx.Module
        equinox layer with __call__(phases) → final_phases.
    phases : jax.Array
        (N,) initial phases.
    target_R : float
        target order parameter R in [0, 1].

    Returns
    -------
    jax.Array
        Scalar loss (R - target_R)^2.
    """
    # type ignore: Equinox modules expose __call__ dynamically by subclass.
    final = model(phases)  # type: ignore[operator]
    R = order_parameter(final)
    return (R - target_R) ** 2

trajectory_loss 

trajectory_loss(
    model: Module,
    phases: Array,
    observed: Array,
    *,
    backend: str = "euler",
) -> jax.Array

Fit model trajectory to observed phase data.

Uses circular distance (via cos) to handle 2pi wrapping.

Parameters

model : eqx.Module equinox layer with forward_with_trajectory(phases) → (final, traj). phases : jax.Array (N,) initial phases. observed : jax.Array (T, N) observed phase trajectory. backend : str Integration backend forwarded to the layer. "euler" (default) calls forward_with_trajectory(phases) unchanged, so any trajectory-capable layer works and existing hashes are preserved. "diffrax" requests the checkpointed continuous-adjoint path and therefore requires a backend-aware layer such as :class:~scpn_phase_orchestrator.nn.ude.UDEKuramotoLayer.

Returns

jax.Array Scalar mean circular distance.

Source code in src/scpn_phase_orchestrator/nn/training.py 
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
def trajectory_loss(
    model: eqx.Module,
    phases: jax.Array,
    observed: jax.Array,
    *,
    backend: str = "euler",
) -> jax.Array:
    """Fit model trajectory to observed phase data.

    Uses circular distance (via cos) to handle 2pi wrapping.

    Parameters
    ----------
    model : eqx.Module
        equinox layer with forward_with_trajectory(phases) → (final, traj).
    phases : jax.Array
        (N,) initial phases.
    observed : jax.Array
        (T, N) observed phase trajectory.
    backend : str
        Integration backend forwarded to the layer. ``"euler"`` (default)
        calls ``forward_with_trajectory(phases)`` unchanged, so any
        trajectory-capable layer works and existing hashes are preserved.
        ``"diffrax"`` requests the checkpointed continuous-adjoint path and
        therefore requires a backend-aware layer such as
        :class:`~scpn_phase_orchestrator.nn.ude.UDEKuramotoLayer`.

    Returns
    -------
    jax.Array
        Scalar mean circular distance.
    """
    # type ignore: training accepts the trajectory-capable Equinox protocol.
    if backend == "euler":
        _, predicted = model.forward_with_trajectory(phases)  # type: ignore[attr-defined]
    else:
        # type ignore: the backend-aware layer forwards the integration backend.
        _, predicted = model.forward_with_trajectory(  # type: ignore[attr-defined]
            phases, backend=backend
        )
    T = min(predicted.shape[0], observed.shape[0])
    pred = predicted[:T]
    obs = observed[:T]
    return jnp.mean(1.0 - jnp.cos(pred - obs))

coupling_sparsity_loss 

coupling_sparsity_loss(
    K: Array, target_density: float = 0.1
) -> jax.Array

L1 penalty driving K toward target density.

Parameters

K : jax.Array (N, N) coupling matrix. target_density : float fraction of nonzero entries desired.

Returns

jax.Array Scalar penalty: |mean(|K|) - target_density * mean(|K|_initial)|.

Source code in src/scpn_phase_orchestrator/nn/training.py 
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
def coupling_sparsity_loss(
    K: jax.Array,
    target_density: float = 0.1,
) -> jax.Array:
    """L1 penalty driving K toward target density.

    Parameters
    ----------
    K : jax.Array
        (N, N) coupling matrix.
    target_density : float
        fraction of nonzero entries desired.

    Returns
    -------
    jax.Array
        Scalar penalty: |mean(|K|) - target_density * mean(|K|_initial)|.
    """
    return jnp.mean(jnp.abs(K)) - target_density * jnp.mean(jnp.abs(K))

train_step 

train_step(
    model: Module,
    loss_fn: Callable[[Module], Array],
    opt_state: Any,
    optimizer: GradientTransformation,
) -> tuple[eqx.Module, Any, jax.Array]

Single optimisation step using optax.

Parameters

model : eqx.Module equinox module to optimise. loss_fn : Callable[[eqx.Module], jax.Array] callable(model) → scalar loss. opt_state : Any optax optimiser state. optimizer : optax.GradientTransformation optax optimiser (e.g. optax.adam(1e-3)).

Returns

tuple[eqx.Module, Any, jax.Array] (updated_model, updated_opt_state, loss_value).

Source code in src/scpn_phase_orchestrator/nn/training.py 
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
def train_step(
    model: eqx.Module,
    loss_fn: Callable[[eqx.Module], jax.Array],
    opt_state: Any,
    optimizer: optax.GradientTransformation,
) -> tuple[eqx.Module, Any, jax.Array]:
    """Single optimisation step using optax.

    Parameters
    ----------
    model : eqx.Module
        equinox module to optimise.
    loss_fn : Callable[[eqx.Module], jax.Array]
        callable(model) → scalar loss.
    opt_state : Any
        optax optimiser state.
    optimizer : optax.GradientTransformation
        optax optimiser (e.g. optax.adam(1e-3)).

    Returns
    -------
    tuple[eqx.Module, Any, jax.Array]
        (updated_model, updated_opt_state, loss_value).
    """
    loss, grads = eqx.filter_value_and_grad(loss_fn)(model)
    updates, opt_state = optimizer.update(grads, opt_state, model)
    model = eqx.apply_updates(model, updates)
    return model, opt_state, loss

train 

train(
    model: Module,
    loss_fn: Callable[[Module], Array],
    optimizer: GradientTransformation,
    n_epochs: int,
    *,
    callback: Callable[[int, Module, Array], None]
    | None = None,
) -> tuple[eqx.Module, list[float]]

Full training loop.

Parameters

model : eqx.Module equinox module to train. loss_fn : Callable[[eqx.Module], jax.Array] callable(model) → scalar loss. optimizer : optax.GradientTransformation optax optimiser. n_epochs : int number of training steps. callback : Callable[[int, eqx.Module, jax.Array], None] | None optional fn(epoch, model, loss) called each step.

Returns

tuple[eqx.Module, list[float]] (trained_model, loss_history).

Source code in src/scpn_phase_orchestrator/nn/training.py 
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
def train(
    model: eqx.Module,
    loss_fn: Callable[[eqx.Module], jax.Array],
    optimizer: optax.GradientTransformation,
    n_epochs: int,
    *,
    callback: Callable[[int, eqx.Module, jax.Array], None] | None = None,
) -> tuple[eqx.Module, list[float]]:
    """Full training loop.

    Parameters
    ----------
    model : eqx.Module
        equinox module to train.
    loss_fn : Callable[[eqx.Module], jax.Array]
        callable(model) → scalar loss.
    optimizer : optax.GradientTransformation
        optax optimiser.
    n_epochs : int
        number of training steps.
    callback : Callable[[int, eqx.Module, jax.Array], None] | None
        optional fn(epoch, model, loss) called each step.

    Returns
    -------
    tuple[eqx.Module, list[float]]
        (trained_model, loss_history).
    """
    opt_state = optimizer.init(eqx.filter(model, eqx.is_array))
    losses: list[float] = []

    step_fn = eqx.filter_jit(
        lambda m, s: train_step(m, loss_fn, s, optimizer),
    )

    for epoch in range(n_epochs):
        model, opt_state, loss = step_fn(model, opt_state)
        loss_val = float(loss)
        losses.append(loss_val)
        if callback is not None:
            callback(epoch, model, loss)

    return model, losses

generate_kuramoto_data 

generate_kuramoto_data(
    N: int,
    T: int,
    dt: float = 0.01,
    K_scale: float = 0.3,
    *,
    key: Array,
) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]

Generate synthetic Kuramoto trajectory with known ground truth.

Parameters

N : int number of oscillators. T : int number of timesteps. dt : float integration timestep. K_scale : float coupling matrix scale. key : jax.Array PRNG key.

Returns

tuple[jax.Array, jax.Array, jax.Array, jax.Array] (K_true, omegas_true, phases0, trajectory) where trajectory is (T, N).

Source code in src/scpn_phase_orchestrator/nn/training.py 
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
def generate_kuramoto_data(
    N: int,
    T: int,
    dt: float = 0.01,
    K_scale: float = 0.3,
    *,
    key: jax.Array,
) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]:
    """Generate synthetic Kuramoto trajectory with known ground truth.

    Parameters
    ----------
    N : int
        number of oscillators.
    T : int
        number of timesteps.
    dt : float
        integration timestep.
    K_scale : float
        coupling matrix scale.
    key : jax.Array
        PRNG key.

    Returns
    -------
    tuple[jax.Array, jax.Array, jax.Array, jax.Array]
        (K_true, omegas_true, phases0, trajectory) where trajectory is (T, N).
    """
    k1, k2, k3 = jax.random.split(key, 3)
    raw = jax.random.normal(k1, (N, N)) * K_scale
    K_true = (raw + raw.T) / 2.0
    K_true = K_true.at[jnp.diag_indices(N)].set(0.0)
    omegas_true = jax.random.normal(k2, (N,)) * 0.3
    phases0 = jax.random.uniform(k3, (N,), maxval=2.0 * jnp.pi)
    _, trajectory = kuramoto_forward(phases0, omegas_true, K_true, dt, T)
    return K_true, omegas_true, phases0, trajectory

generate_chimera_data 

generate_chimera_data(
    N: int,
    T: int,
    dt: float = 0.01,
    coupling_strength: float = 0.5,
    coupling_range: int = 4,
    *,
    key: Array,
) -> tuple[jax.Array, jax.Array, jax.Array]

Generate chimera-producing Kuramoto dynamics on a ring.

Uses non-local coupling on a 1D ring (Kuramoto & Battogtokh 2002) that produces coexistence of synchronised and incoherent domains.

Parameters

N : int number of oscillators on the ring. T : int number of timesteps. dt : float timestep. coupling_strength : float overall coupling scale. coupling_range : int number of neighbours on each side. key : jax.Array PRNG key.

Returns

tuple[jax.Array, jax.Array, jax.Array] (K, phases0, trajectory) where K is (N, N), trajectory is (T, N).

Source code in src/scpn_phase_orchestrator/nn/training.py 
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
def generate_chimera_data(
    N: int,
    T: int,
    dt: float = 0.01,
    coupling_strength: float = 0.5,
    coupling_range: int = 4,
    *,
    key: jax.Array,
) -> tuple[jax.Array, jax.Array, jax.Array]:
    """Generate chimera-producing Kuramoto dynamics on a ring.

    Uses non-local coupling on a 1D ring (Kuramoto & Battogtokh 2002)
    that produces coexistence of synchronised and incoherent domains.

    Parameters
    ----------
    N : int
        number of oscillators on the ring.
    T : int
        number of timesteps.
    dt : float
        timestep.
    coupling_strength : float
        overall coupling scale.
    coupling_range : int
        number of neighbours on each side.
    key : jax.Array
        PRNG key.

    Returns
    -------
    tuple[jax.Array, jax.Array, jax.Array]
        (K, phases0, trajectory) where K is (N, N), trajectory is (T, N).
    """
    k1 = key
    # Non-local ring coupling: each oscillator couples to ±coupling_range neighbours
    K = jnp.zeros((N, N))
    for offset in range(1, coupling_range + 1):
        idx_fwd = jnp.arange(N)
        idx_back = jnp.arange(N)
        K = K.at[idx_fwd, (idx_fwd + offset) % N].set(coupling_strength / N)
        K = K.at[idx_back, (idx_back - offset) % N].set(coupling_strength / N)

    omegas = jnp.zeros(N)
    # Start with partially coherent state (chimera seed)
    phases0 = jnp.where(
        jnp.arange(N) < N // 2,
        0.1 * jax.random.normal(k1, (N,)),
        jax.random.uniform(k1, (N,), maxval=2.0 * jnp.pi),
    )
    _, trajectory = kuramoto_forward(phases0, omegas, K, dt, T)
    return K, phases0, trajectory

Phase autoencoder

nn.phase_autoencoder learns the asymptotic phase, isochrons and phase-sensitivity function of a limit-cycle oscillator from state time series alone (Yawata, Fukami, Taira & Nakao 2024, Chaos 34, 063111). The encoder maps the state to a three-component latent whose first two components lie on the unit circle so that θ = atan2(Y₂, Y₁) is the asymptotic phase; the latent evolves by an exactly-linear normal-form flow with learnable frequency ω and decay λ, trained against a four-term reconstruction/phase/deviation/centring loss. The trained weights are extracted to the pure-NumPy oscillators.phase_reduction evaluator so the phase and the phase response curve are available on the control path without JAX.

phase_autoencoder

A phase autoencoder for model-free phase reduction of limit-cycle oscillators.

Classical phase reduction needs the vector field. The phase autoencoder of Yawata, Fukami, Taira & Nakao (2024) learns the asymptotic phase, the isochrons and the phase-sensitivity function from state time series alone. An encoder maps the oscillator state x to a three-component latent Y = (Y₁, Y₂, Y₃) whose first two components are constrained to the unit circle Y₁² + Y₂² = 1 so that θ = atan2(Y₂, Y₁) is the asymptotic phase; the latent then evolves by the exactly-linear normal-form flow

Y₁,ₜ₊τ = Y₁ cos(ωτ) − Y₂ sin(ωτ),
Y₂,ₜ₊τ = Y₁ sin(ωτ) + Y₂ cos(ωτ),
Y₃,ₜ₊τ = e^{λτ} Y₃,        λ < 0,

with learnable frequency ω and decay λ; a decoder reconstructs x. The four-term training objective ties reconstruction, the uniform phase rotation, the amplitude decay and a centring term that prevents the trivial ω = 0 solution.

The trained encoder/decoder weights and (ω, λ) are extracted to a :class:PhaseReductionWeights record consumed by the pure-NumPy, dependency-light evaluator in oscillators.phase_reduction so the asymptotic phase and the phase-sensitivity function are available on the control path without JAX.

This module follows the published latent constraint and four-term loss; the encoder and decoder are plain ReLU multilayer perceptrons (no batch normalisation, which is a training-stability detail outside the phase-reduction mathematics and would complicate the frozen-weights evaluator).

References

  • Yawata, Fukami, Taira & Nakao 2024, Chaos 34, 063111 (arXiv:2403.06992) — phase autoencoder for limit-cycle oscillators.

Classes

PhaseReductionWeights dataclass 

PhaseReductionWeights(
    encoder_weights: tuple[FloatArray, ...],
    encoder_biases: tuple[FloatArray, ...],
    decoder_weights: tuple[FloatArray, ...],
    decoder_biases: tuple[FloatArray, ...],
    omega: float,
    decay: float,
    state_dim: int,
)

Frozen encoder/decoder weights and (ω, λ) of a phase autoencoder.

Parameters

encoder_weights, encoder_biases : tuple[numpy.ndarray, ...] Per-layer encoder weight matrices (out, in) and bias vectors. decoder_weights, decoder_biases : tuple[numpy.ndarray, ...] Per-layer decoder weight matrices and bias vectors. omega : float The learned angular frequency ω. decay : float The learned amplitude decay λ < 0. state_dim : int The oscillator state dimension n.

PhaseAutoencoder 

PhaseAutoencoder(
    state_dim: int, *, hidden: int = 64, key: Array
)

Bases: Module

Encoder/decoder with a phase-circle latent and normal-form dynamics.

Parameters

state_dim : int The oscillator state dimension n. hidden : int Hidden width of the encoder and decoder multilayer perceptrons. key : jax.Array PRNG key for weight initialisation.

Source code in src/scpn_phase_orchestrator/nn/phase_autoencoder.py 
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
def __init__(self, state_dim: int, *, hidden: int = 64, key: jax.Array) -> None:
    keys = jax.random.split(key, 7)
    self.encoder = [
        eqx.nn.Linear(state_dim, hidden, key=keys[0]),
        eqx.nn.Linear(hidden, hidden, key=keys[1]),
        eqx.nn.Linear(hidden, _LATENT_DIM, key=keys[2]),
    ]
    self.decoder = [
        eqx.nn.Linear(_LATENT_DIM, hidden, key=keys[3]),
        eqx.nn.Linear(hidden, hidden, key=keys[4]),
        eqx.nn.Linear(hidden, state_dim, key=keys[5]),
    ]
    # ω initialised near 1; λ = −softplus(raw_decay) is strictly negative.
    self.raw_omega = jnp.asarray(1.0)
    self.raw_decay = jnp.asarray(0.0)
    self.state_dim = state_dim
Attributes
omega property 
omega: Array

The learned angular frequency ω of the latent rotation.

Returns

jax.Array The learned angular frequency ω.

decay property 
decay: Array

The learned amplitude decay λ = −softplus(raw_decay) < 0.

Returns

jax.Array The learned amplitude decay λ, strictly negative.

Methods:
encode_raw 
encode_raw(state: Array) -> jax.Array

Encode a single state to the unnormalised latent (Ỹ₁, Ỹ₂, Ỹ₃).

Parameters

state : jax.Array The oscillator state of shape (state_dim,).

Returns

jax.Array The unnormalised latent of shape (3,).

Source code in src/scpn_phase_orchestrator/nn/phase_autoencoder.py 
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
def encode_raw(self, state: jax.Array) -> jax.Array:
    """Encode a single state to the unnormalised latent ``(Ỹ₁, Ỹ₂, Ỹ₃)``.

    Parameters
    ----------
    state : jax.Array
        The oscillator state of shape ``(state_dim,)``.

    Returns
    -------
    jax.Array
        The unnormalised latent of shape ``(3,)``.
    """
    activation = state
    for layer in self.encoder[:-1]:
        activation = jax.nn.relu(layer(activation))
    return self.encoder[-1](activation)
encode 
encode(state: Array) -> jax.Array

Encode a single state to the phase-circle latent (Y₁, Y₂, Y₃).

Parameters

state : jax.Array The oscillator state of shape (state_dim,).

Returns

jax.Array The latent of shape (3,) with Y₁² + Y₂² = 1.

Source code in src/scpn_phase_orchestrator/nn/phase_autoencoder.py 
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
def encode(self, state: jax.Array) -> jax.Array:
    """Encode a single state to the phase-circle latent ``(Y₁, Y₂, Y₃)``.

    Parameters
    ----------
    state : jax.Array
        The oscillator state of shape ``(state_dim,)``.

    Returns
    -------
    jax.Array
        The latent of shape ``(3,)`` with ``Y₁² + Y₂² = 1``.
    """
    raw = self.encode_raw(state)
    radius = jnp.sqrt(raw[0] ** 2 + raw[1] ** 2) + 1.0e-9
    return jnp.stack((raw[0] / radius, raw[1] / radius, raw[2]))
decode 
decode(latent: Array) -> jax.Array

Decode a latent (Y₁, Y₂, Y₃) back to the oscillator state.

Parameters

latent : jax.Array The latent of shape (3,).

Returns

jax.Array The reconstructed oscillator state of shape (state_dim,).

Source code in src/scpn_phase_orchestrator/nn/phase_autoencoder.py 
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
def decode(self, latent: jax.Array) -> jax.Array:
    """Decode a latent ``(Y₁, Y₂, Y₃)`` back to the oscillator state.

    Parameters
    ----------
    latent : jax.Array
        The latent of shape ``(3,)``.

    Returns
    -------
    jax.Array
        The reconstructed oscillator state of shape ``(state_dim,)``.
    """
    activation = latent
    for layer in self.decoder[:-1]:
        activation = jax.nn.relu(layer(activation))
    return self.decoder[-1](activation)
advance 
advance(latent: Array, dt: float) -> jax.Array

Advance a latent by dt under the exactly-linear normal-form flow.

Parameters

latent : jax.Array The latent of shape (3,). dt : float The time increment.

Returns

jax.Array The advanced latent of shape (3,).

Source code in src/scpn_phase_orchestrator/nn/phase_autoencoder.py 
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
def advance(self, latent: jax.Array, dt: float) -> jax.Array:
    """Advance a latent by ``dt`` under the exactly-linear normal-form flow.

    Parameters
    ----------
    latent : jax.Array
        The latent of shape ``(3,)``.
    dt : float
        The time increment.

    Returns
    -------
    jax.Array
        The advanced latent of shape ``(3,)``.
    """
    angle = self.omega * dt
    cos = jnp.cos(angle)
    sin = jnp.sin(angle)
    rotated_1 = latent[0] * cos - latent[1] * sin
    rotated_2 = latent[0] * sin + latent[1] * cos
    decayed = jnp.exp(self.decay * dt) * latent[2]
    return jnp.stack((rotated_1, rotated_2, decayed))
asymptotic_phase 
asymptotic_phase(state: Array) -> jax.Array

Return the asymptotic phase θ = atan2(Y₂, Y₁) of a state.

Parameters

state : jax.Array The oscillator state of shape (state_dim,).

Returns

jax.Array The asymptotic phase in (−π, π].

Source code in src/scpn_phase_orchestrator/nn/phase_autoencoder.py 
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
def asymptotic_phase(self, state: jax.Array) -> jax.Array:
    """Return the asymptotic phase ``θ = atan2(Y₂, Y₁)`` of a state.

    Parameters
    ----------
    state : jax.Array
        The oscillator state of shape ``(state_dim,)``.

    Returns
    -------
    jax.Array
        The asymptotic phase in ``(−π, π]``.
    """
    latent = self.encode(state)
    return jnp.arctan2(latent[1], latent[0])

Functions:

extract_phase_reduction_weights 

extract_phase_reduction_weights(
    model: PhaseAutoencoder,
) -> PhaseReductionWeights

Extract a trained model's weights into a frozen NumPy record.

Parameters

model : PhaseAutoencoder The trained phase autoencoder.

Returns

PhaseReductionWeights The frozen encoder/decoder weights and (ω, λ) for the pure-NumPy oscillators.phase_reduction evaluator.

Source code in src/scpn_phase_orchestrator/nn/phase_autoencoder.py 
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
def extract_phase_reduction_weights(
    model: PhaseAutoencoder,
) -> PhaseReductionWeights:
    """Extract a trained model's weights into a frozen NumPy record.

    Parameters
    ----------
    model : PhaseAutoencoder
        The trained phase autoencoder.

    Returns
    -------
    PhaseReductionWeights
        The frozen encoder/decoder weights and ``(ω, λ)`` for the pure-NumPy
        ``oscillators.phase_reduction`` evaluator.
    """
    encoder_weights, encoder_biases = _layer_arrays(model.encoder)
    decoder_weights, decoder_biases = _layer_arrays(model.decoder)
    return PhaseReductionWeights(
        encoder_weights=encoder_weights,
        encoder_biases=encoder_biases,
        decoder_weights=decoder_weights,
        decoder_biases=decoder_biases,
        omega=float(model.omega),
        decay=float(model.decay),
        state_dim=int(model.state_dim),
    )

phase_autoencoder_loss 

phase_autoencoder_loss(
    model: PhaseAutoencoder,
    windows: Array,
    *,
    dt: float,
    weight_recon: float = 1.0,
    weight_phase: float = 0.5,
    weight_deviation: float = 0.5,
    weight_aux: float = 2.0,
) -> jax.Array

Return the four-term phase-autoencoder training loss (Yawata et al. 2024).

Parameters

model : PhaseAutoencoder The model under training. windows : jax.Array Trajectory windows of shape (batch, K + 1, state_dim) — K + 1 consecutive states sampled at spacing dt. dt : float The sampling interval between consecutive states in a window. weight_recon, weight_phase, weight_deviation, weight_aux : float The loss-term weights.

Returns

jax.Array The scalar total loss.

Source code in src/scpn_phase_orchestrator/nn/phase_autoencoder.py 
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
def phase_autoencoder_loss(
    model: PhaseAutoencoder,
    windows: jax.Array,
    *,
    dt: float,
    weight_recon: float = 1.0,
    weight_phase: float = 0.5,
    weight_deviation: float = 0.5,
    weight_aux: float = 2.0,
) -> jax.Array:
    """Return the four-term phase-autoencoder training loss (Yawata et al. 2024).

    Parameters
    ----------
    model : PhaseAutoencoder
        The model under training.
    windows : jax.Array
        Trajectory windows of shape ``(batch, K + 1, state_dim)`` — ``K + 1``
        consecutive states sampled at spacing ``dt``.
    dt : float
        The sampling interval between consecutive states in a window.
    weight_recon, weight_phase, weight_deviation, weight_aux : float
        The loss-term weights.

    Returns
    -------
    jax.Array
        The scalar total loss.
    """
    encode = jax.vmap(model.encode)
    decode = jax.vmap(model.decode)

    flat = windows.reshape(-1, model.state_dim)
    reconstruction = jnp.mean((decode(encode(flat)) - flat) ** 2)

    initial = windows[:, 0, :]
    latent = jax.vmap(model.encode)(initial)
    horizon = windows.shape[1] - 1
    phase_loss = jnp.asarray(0.0)
    deviation_loss = jnp.asarray(0.0)
    for step in range(1, horizon + 1):
        latent = jax.vmap(lambda y: model.advance(y, dt))(latent)
        observed = jax.vmap(model.encode)(windows[:, step, :])
        scale = 1.0 / step
        phase_loss = phase_loss + scale * jnp.mean(
            (observed[:, :2] - latent[:, :2]) ** 2
        )
        deviation_loss = deviation_loss + scale * jnp.mean(
            (observed[:, 2] - latent[:, 2]) ** 2
        )

    encoded_initial = jax.vmap(model.encode_raw)(initial)
    auxiliary = (
        jnp.mean(encoded_initial[:, 0]) ** 2 + jnp.mean(encoded_initial[:, 1]) ** 2
    )

    return (
        weight_recon * reconstruction
        + weight_phase * phase_loss
        + weight_deviation * deviation_loss
        + weight_aux * auxiliary
    )

train_phase_autoencoder 

train_phase_autoencoder(
    windows: Array,
    *,
    dt: float,
    state_dim: int,
    hidden: int = 64,
    epochs: int = 200,
    learning_rate: float = 0.001,
    seed: int = 0,
    loss_kwargs: dict[str, float] | None = None,
) -> tuple[PhaseAutoencoder, jax.Array]

Train a phase autoencoder on trajectory windows.

Parameters

windows : jax.Array Trajectory windows of shape (batch, K + 1, state_dim). dt : float The sampling interval within a window. state_dim : int The oscillator state dimension n. hidden : int Hidden width of the encoder/decoder. epochs : int Number of full-batch Adam steps. learning_rate : float The Adam learning rate. seed : int PRNG seed for weight initialisation. loss_kwargs : dict[str, float] | None Optional overrides for the loss-term weights.

Returns

tuple[PhaseAutoencoder, jax.Array] The trained model and the final loss value.

Source code in src/scpn_phase_orchestrator/nn/phase_autoencoder.py 
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
def train_phase_autoencoder(
    windows: jax.Array,
    *,
    dt: float,
    state_dim: int,
    hidden: int = 64,
    epochs: int = 200,
    learning_rate: float = 1.0e-3,
    seed: int = 0,
    loss_kwargs: dict[str, float] | None = None,
) -> tuple[PhaseAutoencoder, jax.Array]:
    """Train a phase autoencoder on trajectory windows.

    Parameters
    ----------
    windows : jax.Array
        Trajectory windows of shape ``(batch, K + 1, state_dim)``.
    dt : float
        The sampling interval within a window.
    state_dim : int
        The oscillator state dimension ``n``.
    hidden : int
        Hidden width of the encoder/decoder.
    epochs : int
        Number of full-batch Adam steps.
    learning_rate : float
        The Adam learning rate.
    seed : int
        PRNG seed for weight initialisation.
    loss_kwargs : dict[str, float] | None
        Optional overrides for the loss-term weights.

    Returns
    -------
    tuple[PhaseAutoencoder, jax.Array]
        The trained model and the final loss value.
    """
    import optax

    overrides = loss_kwargs or {}
    model = PhaseAutoencoder(state_dim, hidden=hidden, key=jax.random.key(seed))
    optimizer = optax.adam(learning_rate)
    opt_state = optimizer.init(eqx.filter(model, eqx.is_inexact_array))

    def _loss(candidate: PhaseAutoencoder) -> jax.Array:
        """Return the reconstruction loss for a batch."""
        return phase_autoencoder_loss(candidate, windows, dt=dt, **overrides)

    step = _make_train_step(optimizer, _loss)
    loss_value = jnp.asarray(jnp.inf)
    for _ in range(epochs):
        model, opt_state, loss_value = step(model, opt_state)
    return model, loss_value

Differentiable Supervisor

nn.supervisor provides the differentiable neural policy surface for closed-loop Kuramoto control. It is intentionally separate from supervisor.policy.SupervisorPolicy: the neural policy remains a JAX/equinox module trained over simulator or replay rollouts, while live actuation still flows through ControlAction, mapper limits, and safety gates.

The built-in objective maximizes good-partition synchrony while penalizing bad-partition synchrony, control energy, and abrupt action changes. The module also includes a squashed-Gaussian action sampler and clipped PPO loss/train step for on-policy RL experiments. This is a production-quality differentiable training surface, not a claim that large-scale RL benchmarks or a preprint have already been completed.

import equinox as eqx
import jax
import jax.numpy as jnp
import optax

from scpn_phase_orchestrator.nn import (
    DifferentiableSupervisorConfig,
    DifferentiableSupervisorPolicy,
    KuramotoSupervisorScenario,
    supervisor_train_step,
)

scenario = KuramotoSupervisorScenario(
    phases=jnp.array([0.0, 0.1, 2.7, 3.1]),
    omegas=jnp.array([0.04, 0.03, -0.03, -0.04]),
    base_K=jnp.full((4, 4), 0.03) - jnp.eye(4) * 0.03,
    good_mask=jnp.array([1.0, 1.0, 0.0, 0.0]),
    bad_mask=jnp.array([0.0, 0.0, 1.0, 1.0]),
    dt=0.02,
    inner_steps=4,
    horizon=3,
)
policy = DifferentiableSupervisorPolicy(
    DifferentiableSupervisorConfig(n_oscillators=4),
    key=jax.random.PRNGKey(0),
)
optimizer = optax.adam(1e-3)
opt_state = optimizer.init(eqx.filter(policy, eqx.is_array))
policy, opt_state, loss = supervisor_train_step(
    policy,
    scenario,
    opt_state,
    optimizer,
)

supervisor

Differentiable supervisor policies for closed-loop Kuramoto control.

This package is the JAX/equinox counterpart to supervisor.policy.SupervisorPolicy. It keeps the learning surface fully differentiable and array-native, then exposes a small adapter for the existing ControlAction actuation path. Safety projection, rate limits, and live interlocks remain outside the gradient path.

Classes

DifferentiableSupervisorConfig dataclass 

DifferentiableSupervisorConfig(
    n_oscillators: int,
    hidden_width: int = 32,
    hidden_depth: int = 2,
    n_layer_controls: int = 2,
    max_global_delta_K: float = 0.05,
    max_global_delta_zeta: float = 0.1,
    max_layer_delta_K: float = 0.03,
    control_energy_weight: float = 0.01,
    bad_sync_weight: float = 0.25,
    smoothness_weight: float = 0.001,
)

Static configuration for DifferentiableSupervisorPolicy.

Parameters

n_oscillators : object Number of oscillators in the controlled Kuramoto system. hidden_width : object Width of each MLP hidden layer. hidden_depth : object Number of hidden layers in the MLP. n_layer_controls : object Number of mask-scoped K controls. The default maps to good and bad partitions in KuramotoSupervisorScenario. max_global_delta_K : object Absolute bound for global coupling increments. max_global_delta_zeta : object Absolute bound for global damping/drive command. max_layer_delta_K : object Absolute bound for partition-local coupling deltas. control_energy_weight : object Quadratic penalty on control action magnitude. bad_sync_weight : object Penalty for synchronising the bad partition. smoothness_weight : object Quadratic penalty on action changes over rollout.

KuramotoSupervisorScenario

Bases: NamedTuple

Closed-loop differentiable Kuramoto control problem.

good_mask and bad_mask are non-negative oscillator membership weights. Binary masks are typical, but soft memberships are supported for differentiable curriculum construction.

SupervisorAction

Bases: NamedTuple

Continuous differentiable control emitted by the neural supervisor.

SupervisorActionProjection

Bases: NamedTuple

Non-actuating safety projection record for a neural supervisor proposal.

SupervisorBaselineReport

Bases: NamedTuple

Aggregate audit report for supervisor baseline comparison records.

Methods:
to_audit_record 
to_audit_record() -> dict[str, Any]

Return a JSON-serialisable baseline aggregation report.

Returns

dict[str, Any] Return a JSON-serialisable baseline aggregation report.

Source code in src/scpn_phase_orchestrator/nn/supervisor/_types.py 
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
def to_audit_record(self) -> dict[str, Any]:
    """Return a JSON-serialisable baseline aggregation report.

    Returns
    -------
    dict[str, Any]
        Return a JSON-serialisable baseline aggregation report.
    """
    return _json_object(
        {
            "proposal_type": "differentiable_supervisor_baseline_report",
            "actuation_permitted": self.actuation_permitted,
            "report_label": self.report_label,
            "summary": dict(self.summary),
            "comparisons": [dict(record) for record in self.comparisons],
        },
        "supervisor_baseline_report",
    )

SupervisorCorpusReplayProposals

Bases: NamedTuple

Replay-only proposal set generated from a supervisor scenario corpus.

Methods:
to_audit_record 
to_audit_record() -> dict[str, Any]

Return a JSON-serialisable corpus proposal record.

Returns

dict[str, Any] Return a JSON-serialisable corpus proposal record.

Source code in src/scpn_phase_orchestrator/nn/supervisor/_types.py 
182
183
184
185
186
187
188
189
190
191
192
193
194
195
def to_audit_record(self) -> dict[str, Any]:
    """Return a JSON-serialisable corpus proposal record.

    Returns
    -------
    dict[str, Any]
        Return a JSON-serialisable corpus proposal record.
    """
    return {
        "proposal_type": "differentiable_supervisor_corpus_replay",
        "actuation_permitted": self.actuation_permitted,
        "scenario_count": len(self.proposals),
        "proposals": [proposal.to_audit_record() for proposal in self.proposals],
    }

SupervisorExperimentManifest

Bases: NamedTuple

Reproducibility manifest for supervisor baseline experiment artefacts.

Methods:
to_audit_record 
to_audit_record() -> dict[str, Any]

Return a JSON-serialisable reproducibility manifest.

Returns

dict[str, Any] Return a JSON-serialisable reproducibility manifest.

Source code in src/scpn_phase_orchestrator/nn/supervisor/_types.py 
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
def to_audit_record(self) -> dict[str, Any]:
    """Return a JSON-serialisable reproducibility manifest.

    Returns
    -------
    dict[str, Any]
        Return a JSON-serialisable reproducibility manifest.
    """
    return _json_object(
        {
            "proposal_type": "differentiable_supervisor_experiment_manifest",
            "actuation_permitted": self.actuation_permitted,
            "command": self.command,
            "git_sha": self.git_sha,
            "dependency_lock": dict(self.dependency_lock),
            "device_info": dict(self.device_info),
            "seed_list": list(self.seed_list),
            "artifacts": dict(self.artifacts),
            "baseline_report": dict(self.baseline_report),
        },
        "supervisor_experiment_manifest",
    )

SupervisorHandTunedBaselineComparison

Bases: NamedTuple

Audit-only comparison against the rule-based SupervisorPolicy.

Methods:
to_audit_record 
to_audit_record() -> dict[str, Any]

Return a JSON-serialisable hand-tuned-baseline comparison record.

Returns

dict[str, Any] Return a JSON-serialisable hand-tuned-baseline comparison record.

Source code in src/scpn_phase_orchestrator/nn/supervisor/_types.py 
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
def to_audit_record(self) -> dict[str, Any]:
    """Return a JSON-serialisable hand-tuned-baseline comparison record.

    Returns
    -------
    dict[str, Any]
        Return a JSON-serialisable hand-tuned-baseline comparison record.
    """
    return _json_object(
        {
            "proposal_type": "differentiable_supervisor_hand_tuned_baseline",
            "actuation_permitted": self.actuation_permitted,
            "comparison_label": self.comparison_label,
            "scenario_summary": dict(self.scenario_summary),
            "baseline": dict(self.baseline),
            "supervisor": dict(self.supervisor),
            "metrics": dict(self.metrics),
        },
        "supervisor_hand_tuned_baseline_comparison",
    )

SupervisorLearnerProposalComparison

Bases: NamedTuple

Audit-only comparison against learner-shaped autotune proposal records.

Methods:
to_audit_record 
to_audit_record() -> dict[str, Any]

Return a JSON-serialisable learner-proposal comparison record.

Returns

dict[str, Any] Return a JSON-serialisable learner-proposal comparison record.

Source code in src/scpn_phase_orchestrator/nn/supervisor/_types.py 
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
def to_audit_record(self) -> dict[str, Any]:
    """Return a JSON-serialisable learner-proposal comparison record.

    Returns
    -------
    dict[str, Any]
        Return a JSON-serialisable learner-proposal comparison record.
    """
    return _json_object(
        {
            "proposal_type": "differentiable_supervisor_learner_comparison",
            "actuation_permitted": self.actuation_permitted,
            "comparison_label": self.comparison_label,
            "supervisor": dict(self.supervisor),
            "learner_proposals": [
                dict(proposal) for proposal in self.learner_proposals
            ],
            "metrics": dict(self.metrics),
        },
        "supervisor_learner_proposal_comparison",
    )

SupervisorLossAux

Bases: NamedTuple

Diagnostics returned by closed_loop_supervisor_loss.

SupervisorPPOAux

Bases: NamedTuple

Diagnostics returned by ppo_supervisor_loss.

SupervisorPPOBatch

Bases: NamedTuple

On-policy PPO batch for the differentiable supervisor.

Arrays carry a leading batch dimension. actions must contain bounded continuous actions produced by pack_supervisor_action.

SupervisorPPOCheckpoint

Bases: NamedTuple

Loaded PPO checkpoint state for deterministic supervisor training resume.

SupervisorPPOCorpusRollout

Bases: NamedTuple

Corpus-wide replay rollout with per-episode scenario provenance.

SupervisorPPORollout

Bases: NamedTuple

Replay-only rollout outputs for PPO-style supervisor training.

SupervisorPPOTrainResult

Bases: NamedTuple

PPO training result with checkpoint-resume bookkeeping.

SupervisorRandomBaselineComparison

Bases: NamedTuple

Audit-only comparison against a seeded bounded-random action baseline.

Methods:
to_audit_record 
to_audit_record() -> dict[str, Any]

Return a JSON-serialisable random-baseline comparison record.

Returns

dict[str, Any] Return a JSON-serialisable random-baseline comparison record.

Source code in src/scpn_phase_orchestrator/nn/supervisor/_types.py 
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
def to_audit_record(self) -> dict[str, Any]:
    """Return a JSON-serialisable random-baseline comparison record.

    Returns
    -------
    dict[str, Any]
        Return a JSON-serialisable random-baseline comparison record.
    """
    return _json_object(
        {
            "proposal_type": "differentiable_supervisor_random_baseline",
            "actuation_permitted": self.actuation_permitted,
            "comparison_label": self.comparison_label,
            "scenario_summary": dict(self.scenario_summary),
            "baseline": dict(self.baseline),
            "supervisor": dict(self.supervisor),
            "metrics": dict(self.metrics),
        },
        "supervisor_random_baseline_comparison",
    )

SupervisorReplayComparison

Bases: NamedTuple

Audit-only comparison between neural supervisor and replay policy search.

Methods:
to_audit_record 
to_audit_record() -> dict[str, Any]

Return a JSON-serialisable non-actuating comparison record.

Returns

dict[str, Any] Return a JSON-serialisable non-actuating comparison record.

Source code in src/scpn_phase_orchestrator/nn/supervisor/_types.py 
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
def to_audit_record(self) -> dict[str, Any]:
    """Return a JSON-serialisable non-actuating comparison record.

    Returns
    -------
    dict[str, Any]
        Return a JSON-serialisable non-actuating comparison record.
    """
    return _json_object(
        {
            "proposal_type": "differentiable_supervisor_replay_comparison",
            "actuation_permitted": self.actuation_permitted,
            "comparison_label": self.comparison_label,
            "supervisor": dict(self.supervisor),
            "replay_policy_search": dict(self.replay_policy_search),
            "metrics": dict(self.metrics),
        },
        "supervisor_replay_comparison",
    )

SupervisorReplayProposal

Bases: NamedTuple

Replay-only neural supervisor proposal record for audit review.

Methods:
to_audit_record 
to_audit_record() -> dict[str, Any]

Return a JSON-serialisable replay proposal record.

Returns

dict[str, Any] Return a JSON-serialisable replay proposal record.

Source code in src/scpn_phase_orchestrator/nn/supervisor/_types.py 
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
def to_audit_record(self) -> dict[str, Any]:
    """Return a JSON-serialisable replay proposal record.

    Returns
    -------
    dict[str, Any]
        Return a JSON-serialisable replay proposal record.
    """
    return {
        "proposal_type": "differentiable_supervisor_replay_proposal",
        "actuation_permitted": self.actuation_permitted,
        "scenario_summary": dict(self.scenario_summary),
        "scenario_metadata": dict(self.scenario_metadata),
        "metrics": dict(self.metrics),
        "action": _supervisor_action_to_record(self.action),
        "projection": dict(self.projection.audit_record),
    }

SupervisorScenarioCorpus

Bases: NamedTuple

Validated replay scenario corpus for supervisor training.

SupervisorStaticBaselineComparison

Bases: NamedTuple

Audit-only comparison against a static zero-action supervisor baseline.

Methods:
to_audit_record 
to_audit_record() -> dict[str, Any]

Return a JSON-serialisable static-baseline comparison record.

Returns

dict[str, Any] Return a JSON-serialisable static-baseline comparison record.

Source code in src/scpn_phase_orchestrator/nn/supervisor/_types.py 
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
def to_audit_record(self) -> dict[str, Any]:
    """Return a JSON-serialisable static-baseline comparison record.

    Returns
    -------
    dict[str, Any]
        Return a JSON-serialisable static-baseline comparison record.
    """
    return _json_object(
        {
            "proposal_type": "differentiable_supervisor_static_baseline",
            "actuation_permitted": self.actuation_permitted,
            "comparison_label": self.comparison_label,
            "scenario_summary": dict(self.scenario_summary),
            "baseline": dict(self.baseline),
            "supervisor": dict(self.supervisor),
            "metrics": dict(self.metrics),
        },
        "supervisor_static_baseline_comparison",
    )

DifferentiableSupervisorPolicy 

DifferentiableSupervisorPolicy(
    config: DifferentiableSupervisorConfig, *, key: Array
)

Bases: Module

Equinox neural supervisor for differentiable Kuramoto control.

The policy consumes a compact feature vector derived from phases, masks, and coupling statistics. Its output is a bounded continuous control action that can be differentiated through a full jax.lax.scan rollout.

Source code in src/scpn_phase_orchestrator/nn/supervisor/policy.py 
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
def __init__(
    self,
    config: DifferentiableSupervisorConfig,
    *,
    key: jax.Array,
) -> None:
    self.config = config
    out_size = 3 + config.n_layer_controls
    self.log_std = jnp.full((2 + config.n_layer_controls,), -0.5)
    self.network = eqx.nn.MLP(
        in_size=8,
        out_size=out_size,
        width_size=config.hidden_width,
        depth=config.hidden_depth,
        activation=jax.nn.tanh,
        key=key,
    )
Methods:
__call__ 
__call__(
    scenario: KuramotoSupervisorScenario,
) -> SupervisorAction

Return a bounded continuous control action for scenario.

Source code in src/scpn_phase_orchestrator/nn/supervisor/policy.py 
70
71
72
73
74
75
76
77
def __call__(self, scenario: KuramotoSupervisorScenario) -> SupervisorAction:
    """Return a bounded continuous control action for ``scenario``."""
    action_mean, value = _policy_mean_and_value(self, scenario)
    return unpack_supervisor_action(
        jnp.tanh(action_mean) * _action_bounds(self.config),
        value_estimate=value,
        config=self.config,
    )

Functions:

masked_order_parameter 

masked_order_parameter(
    phases: Array, weights: Array
) -> jax.Array

Weighted Kuramoto order parameter for a partition of oscillators.

Parameters

phases : jax.Array Oscillator phases in radians, shape (N,). weights : jax.Array Per-oscillator partition weights.

Returns

jax.Array The weighted Kuramoto order parameter.

Source code in src/scpn_phase_orchestrator/nn/supervisor/_shared.py 
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
def masked_order_parameter(phases: jax.Array, weights: jax.Array) -> jax.Array:
    """Weighted Kuramoto order parameter for a partition of oscillators.

    Parameters
    ----------
    phases : jax.Array
        Oscillator phases in radians, shape ``(N,)``.
    weights : jax.Array
        Per-oscillator partition weights.

    Returns
    -------
    jax.Array
        The weighted Kuramoto order parameter.
    """
    safe_weights = jnp.clip(weights, min=0.0)
    total = jnp.maximum(jnp.sum(safe_weights), 1.0e-12)
    z = jnp.sum(safe_weights * jnp.exp(1j * phases)) / total
    return jnp.abs(z)

supervisor_action_to_candidate 

supervisor_action_to_candidate(
    action: SupervisorAction,
    *,
    base: KnobPolicyCandidate | None = None,
) -> KnobPolicyCandidate

Map a supervisor action onto a candidate, relative to a base candidate.

Parameters

action : SupervisorAction The supervisor control action: a global coupling delta, a global damping delta, and per-layer coupling deltas. base : KnobPolicyCandidate | None The candidate the deltas are applied to. The alpha, Psi and cross_channel_gains knobs are carried through from it unchanged. Defaults to the zero candidate.

Returns

KnobPolicyCandidate A candidate with K and zeta advanced by the global deltas and channel_weights advanced by the per-layer coupling deltas.

Source code in src/scpn_phase_orchestrator/nn/supervisor/candidate_bridge.py 
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
def supervisor_action_to_candidate(
    action: SupervisorAction,
    *,
    base: KnobPolicyCandidate | None = None,
) -> KnobPolicyCandidate:
    """Map a supervisor action onto a candidate, relative to a base candidate.

    Parameters
    ----------
    action : SupervisorAction
        The supervisor control action: a global coupling delta, a global damping
        delta, and per-layer coupling deltas.
    base : KnobPolicyCandidate | None
        The candidate the deltas are applied to. The ``alpha``, ``Psi`` and
        ``cross_channel_gains`` knobs are carried through from it unchanged.
        Defaults to the zero candidate.

    Returns
    -------
    KnobPolicyCandidate
        A candidate with ``K`` and ``zeta`` advanced by the global deltas and
        ``channel_weights`` advanced by the per-layer coupling deltas.
    """
    reference = base if base is not None else KnobPolicyCandidate()
    delta_coupling = float(action.delta_K_global)
    delta_damping = float(action.delta_zeta_global)
    layer_deltas = np.asarray(action.delta_K_layers, dtype=float).ravel()
    base_weights = reference.channel_weights
    channel_weights = tuple(
        (base_weights[index] if index < len(base_weights) else 0.0)
        + float(layer_deltas[index])
        for index in range(len(layer_deltas))
    )
    return KnobPolicyCandidate(
        K=_add_global(reference.K, delta_coupling),
        alpha=reference.alpha,
        zeta=_add_global(reference.zeta, delta_damping),
        Psi=reference.Psi,
        channel_weights=channel_weights,
        cross_channel_gains=reference.cross_channel_gains,
    )

supervisor_policy_to_candidate 

supervisor_policy_to_candidate(
    policy: DifferentiableSupervisorPolicy,
    scenario: KuramotoSupervisorScenario,
    *,
    base: KnobPolicyCandidate | None = None,
) -> KnobPolicyCandidate

Run a supervisor policy deterministically and map its action to a candidate.

Parameters

policy : DifferentiableSupervisorPolicy The learned supervisor policy. It is evaluated for its deterministic mean action; no stochastic sample is drawn. scenario : KuramotoSupervisorScenario The scenario the policy is evaluated on. base : KnobPolicyCandidate | None The base candidate the deltas are applied to. Defaults to the zero candidate.

Returns

KnobPolicyCandidate The candidate equivalent of the policy's recommendation.

Source code in src/scpn_phase_orchestrator/nn/supervisor/candidate_bridge.py 
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
def supervisor_policy_to_candidate(
    policy: DifferentiableSupervisorPolicy,
    scenario: KuramotoSupervisorScenario,
    *,
    base: KnobPolicyCandidate | None = None,
) -> KnobPolicyCandidate:
    """Run a supervisor policy deterministically and map its action to a candidate.

    Parameters
    ----------
    policy : DifferentiableSupervisorPolicy
        The learned supervisor policy. It is evaluated for its deterministic mean
        action; no stochastic sample is drawn.
    scenario : KuramotoSupervisorScenario
        The scenario the policy is evaluated on.
    base : KnobPolicyCandidate | None
        The base candidate the deltas are applied to. Defaults to the zero
        candidate.

    Returns
    -------
    KnobPolicyCandidate
        The candidate equivalent of the policy's recommendation.
    """
    return supervisor_action_to_candidate(policy(scenario), base=base)

load_supervisor_ppo_checkpoint 

load_supervisor_ppo_checkpoint(
    checkpoint_dir: str | Path,
    *,
    template_policy: DifferentiableSupervisorPolicy,
    template_opt_state: Any,
) -> SupervisorPPOCheckpoint

Load a PPO supervisor checkpoint against explicit policy/state templates.

Parameters

checkpoint_dir : str | Path Directory for training checkpoints, or None. template_policy : DifferentiableSupervisorPolicy Template policy used to reconstruct the checkpoint. template_opt_state : Any Template optimiser state used to reconstruct the checkpoint.

Returns

SupervisorPPOCheckpoint The loaded PPO checkpoint.

Raises

FileNotFoundError If the checkpoint cannot be found.

Source code in src/scpn_phase_orchestrator/nn/supervisor/checkpoint.py 
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
def load_supervisor_ppo_checkpoint(
    checkpoint_dir: str | Path,
    *,
    template_policy: DifferentiableSupervisorPolicy,
    template_opt_state: Any,
) -> SupervisorPPOCheckpoint:
    """Load a PPO supervisor checkpoint against explicit policy/state templates.

    Parameters
    ----------
    checkpoint_dir : str | Path
        Directory for training checkpoints, or ``None``.
    template_policy : DifferentiableSupervisorPolicy
        Template policy used to reconstruct the checkpoint.
    template_opt_state : Any
        Template optimiser state used to reconstruct the checkpoint.

    Returns
    -------
    SupervisorPPOCheckpoint
        The loaded PPO checkpoint.

    Raises
    ------
    FileNotFoundError
        If the checkpoint cannot be found.
    """
    checkpoint_path = Path(checkpoint_dir)
    metadata_path = checkpoint_path / "metadata.json"
    state_path = checkpoint_path / "state.eqx"
    metadata = _load_checkpoint_metadata(metadata_path)
    if not state_path.exists():
        raise FileNotFoundError(f"missing checkpoint payload: {state_path}")

    key_shape = _metadata_shape(metadata, "key_shape")
    loss_history_shape = _metadata_shape(metadata, "loss_history_shape")
    key_dtype = _metadata_dtype(metadata, "key_dtype")
    loss_history_dtype = _metadata_dtype(metadata, "loss_history_dtype")
    template_payload = _SupervisorPPOCheckpointPayload(
        policy=template_policy,
        opt_state=template_opt_state,
        key=jnp.zeros(key_shape, dtype=key_dtype),
        loss_history=jnp.zeros(loss_history_shape, dtype=loss_history_dtype),
    )
    loaded = eqx.tree_deserialise_leaves(state_path, template_payload)
    return SupervisorPPOCheckpoint(
        policy=loaded.policy,
        opt_state=loaded.opt_state,
        key=loaded.key,
        loss_history=loaded.loss_history,
        n_updates=_non_negative_int(metadata["n_updates"], "n_updates"),
        metadata=_json_object(metadata.get("metadata"), "checkpoint metadata"),
    )

save_supervisor_ppo_checkpoint 

save_supervisor_ppo_checkpoint(
    checkpoint_dir: str | Path,
    *,
    policy: DifferentiableSupervisorPolicy,
    opt_state: Any,
    key: Array,
    n_updates: int,
    loss_history: Array,
    metadata: dict[str, Any] | None = None,
    overwrite: bool = False,
) -> Path

Persist PPO supervisor trainer state for deterministic resume.

Parameters

checkpoint_dir : str | Path Directory for training checkpoints, or None. policy : DifferentiableSupervisorPolicy The differentiable supervisor policy. opt_state : Any The optax optimiser state. key : jax.Array JAX PRNG key. n_updates : int Number of optimiser updates performed. loss_history : jax.Array Recorded per-step loss history. metadata : dict[str, Any] | None Associated metadata mapping, or None. overwrite : bool Whether to overwrite an existing checkpoint.

Returns

Path The path of the written checkpoint.

Raises

NotADirectoryError If the checkpoint directory path is not a directory. FileExistsError If the checkpoint already exists and overwrite is false.

Source code in src/scpn_phase_orchestrator/nn/supervisor/checkpoint.py 
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
def save_supervisor_ppo_checkpoint(
    checkpoint_dir: str | Path,
    *,
    policy: DifferentiableSupervisorPolicy,
    opt_state: Any,
    key: jax.Array,
    n_updates: int,
    loss_history: jax.Array,
    metadata: dict[str, Any] | None = None,
    overwrite: bool = False,
) -> Path:
    """Persist PPO supervisor trainer state for deterministic resume.

    Parameters
    ----------
    checkpoint_dir : str | Path
        Directory for training checkpoints, or ``None``.
    policy : DifferentiableSupervisorPolicy
        The differentiable supervisor policy.
    opt_state : Any
        The optax optimiser state.
    key : jax.Array
        JAX PRNG key.
    n_updates : int
        Number of optimiser updates performed.
    loss_history : jax.Array
        Recorded per-step loss history.
    metadata : dict[str, Any] | None
        Associated metadata mapping, or ``None``.
    overwrite : bool
        Whether to overwrite an existing checkpoint.

    Returns
    -------
    Path
        The path of the written checkpoint.

    Raises
    ------
    NotADirectoryError
        If the checkpoint directory path is not a directory.
    FileExistsError
        If the checkpoint already exists and ``overwrite`` is false.
    """
    checkpoint_path = Path(checkpoint_dir)
    if checkpoint_path.exists() and not checkpoint_path.is_dir():
        raise NotADirectoryError(f"{checkpoint_path} is not a checkpoint directory")
    checkpoint_path.mkdir(parents=True, exist_ok=True)

    state_path = checkpoint_path / "state.eqx"
    metadata_path = checkpoint_path / "metadata.json"
    if not overwrite and (state_path.exists() or metadata_path.exists()):
        raise FileExistsError(
            f"{checkpoint_path} already contains a supervisor PPO checkpoint"
        )

    n_updates = _non_negative_int(n_updates, "n_updates")
    key = jnp.asarray(key)
    loss_history = jnp.asarray(loss_history)
    user_metadata = _json_object(metadata, "metadata")
    payload = _SupervisorPPOCheckpointPayload(
        policy=policy,
        opt_state=opt_state,
        key=key,
        loss_history=loss_history,
    )
    checkpoint_metadata = {
        "format": _SUPERVISOR_PPO_CHECKPOINT_FORMAT,
        "schema_version": _SUPERVISOR_PPO_CHECKPOINT_SCHEMA_VERSION,
        "n_updates": n_updates,
        "key_shape": list(key.shape),
        "key_dtype": str(key.dtype),
        "loss_history_shape": list(loss_history.shape),
        "loss_history_dtype": str(loss_history.dtype),
        "metadata": user_metadata,
    }

    state_tmp = checkpoint_path / "state.eqx.tmp"
    metadata_tmp = checkpoint_path / "metadata.json.tmp"
    eqx.tree_serialise_leaves(state_tmp, payload)
    metadata_tmp.write_text(
        json.dumps(checkpoint_metadata, sort_keys=True, indent=2, allow_nan=False)
        + "\n",
        encoding="utf-8",
    )
    state_tmp.replace(state_path)
    metadata_tmp.replace(metadata_path)
    return checkpoint_path

compare_supervisor_hand_tuned_baseline 

compare_supervisor_hand_tuned_baseline(
    policy: DifferentiableSupervisorPolicy,
    scenario: KuramotoSupervisorScenario,
    *,
    boundary_state: BoundaryState | None = None,
    comparison_label: str = "hand_tuned_supervisor_policy",
) -> SupervisorHandTunedBaselineComparison

Compare a neural supervisor proposal against rule-based SupervisorPolicy.

Parameters

policy : DifferentiableSupervisorPolicy The differentiable supervisor policy. scenario : KuramotoSupervisorScenario The Kuramoto supervisor scenario. boundary_state : BoundaryState | None The boundary-observer state, or None. comparison_label : str Label identifying the comparison.

Returns

SupervisorHandTunedBaselineComparison The supervisor-vs-rule-based comparison.

Raises

ValueError If the inputs are invalid or inconsistent.

Source code in src/scpn_phase_orchestrator/nn/supervisor/comparison.py 
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
def compare_supervisor_hand_tuned_baseline(
    policy: DifferentiableSupervisorPolicy,
    scenario: KuramotoSupervisorScenario,
    *,
    boundary_state: BoundaryState | None = None,
    comparison_label: str = "hand_tuned_supervisor_policy",
) -> SupervisorHandTunedBaselineComparison:
    """Compare a neural supervisor proposal against rule-based ``SupervisorPolicy``.

    Parameters
    ----------
    policy : DifferentiableSupervisorPolicy
        The differentiable supervisor policy.
    scenario : KuramotoSupervisorScenario
        The Kuramoto supervisor scenario.
    boundary_state : BoundaryState | None
        The boundary-observer state, or ``None``.
    comparison_label : str
        Label identifying the comparison.

    Returns
    -------
    SupervisorHandTunedBaselineComparison
        The supervisor-vs-rule-based comparison.

    Raises
    ------
    ValueError
        If the inputs are invalid or inconsistent.
    """
    if not comparison_label:
        raise ValueError("comparison_label must not be empty")
    active_boundary = boundary_state or BoundaryState()
    upde_state = _upde_state_from_supervisor_scenario(scenario)
    legacy_policy = SupervisorPolicy(RegimeManager(cooldown_steps=0))
    policy_actions = legacy_policy.decide(upde_state, active_boundary)
    baseline_action = _supervisor_action_from_control_actions(
        policy_actions, policy.config
    )
    baseline_record = _rollout_static_supervisor_action(
        "hand_tuned_supervisor_policy",
        baseline_action,
        scenario,
        policy.config,
    )
    baseline_record["policy_actions"] = [
        _control_action_record(action) for action in policy_actions
    ]
    supervisor_record = _rollout_static_supervisor_action(
        "differentiable_supervisor",
        policy(scenario),
        scenario,
        policy.config,
    )
    baseline_metrics = _mapping_value(baseline_record, "metrics")
    supervisor_metrics = _mapping_value(supervisor_record, "metrics")
    metrics = _prefixed_float_metrics("baseline", baseline_metrics)
    metrics.update(_prefixed_float_metrics("supervisor", supervisor_metrics))
    metrics["delta_reward"] = metrics["supervisor_reward"] - metrics["baseline_reward"]
    return SupervisorHandTunedBaselineComparison(
        baseline=baseline_record,
        supervisor=supervisor_record,
        scenario_summary=_supervisor_scenario_summary(scenario),
        comparison_label=comparison_label,
        metrics=metrics,
        actuation_permitted=False,
    )

compare_supervisor_learner_proposals 

compare_supervisor_learner_proposals(
    supervisor_proposal: SupervisorReplayProposal
    | SupervisorCorpusReplayProposals,
    learner_proposals: Iterable[Any],
    *,
    comparison_label: str = "learner_proposal_generators",
) -> SupervisorLearnerProposalComparison

Compare supervisor replay output with replay-only autotune learner proposals.

Parameters

supervisor_proposal : SupervisorReplayProposal | SupervisorCorpusReplayProposals The supervisor replay proposal(s). learner_proposals : Iterable[Any] The replay-only learner proposals. comparison_label : str Label identifying the comparison.

Returns

SupervisorLearnerProposalComparison The supervisor-vs-learner comparison.

Raises

ValueError If the inputs are invalid or inconsistent.

Source code in src/scpn_phase_orchestrator/nn/supervisor/comparison.py 
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
def compare_supervisor_learner_proposals(
    supervisor_proposal: SupervisorReplayProposal | SupervisorCorpusReplayProposals,
    learner_proposals: Iterable[Any],
    *,
    comparison_label: str = "learner_proposal_generators",
) -> SupervisorLearnerProposalComparison:
    """Compare supervisor replay output with replay-only autotune learner proposals.

    Parameters
    ----------
    supervisor_proposal : SupervisorReplayProposal | SupervisorCorpusReplayProposals
        The supervisor replay proposal(s).
    learner_proposals : Iterable[Any]
        The replay-only learner proposals.
    comparison_label : str
        Label identifying the comparison.

    Returns
    -------
    SupervisorLearnerProposalComparison
        The supervisor-vs-learner comparison.

    Raises
    ------
    ValueError
        If the inputs are invalid or inconsistent.
    """
    if not comparison_label:
        raise ValueError("comparison_label must not be empty")
    supervisor_record = _json_object(
        supervisor_proposal.to_audit_record(),
        "supervisor proposal audit record",
    )
    learner_records = tuple(
        _learner_proposal_record_from_object(proposal, f"learner proposal {index}")
        for index, proposal in enumerate(learner_proposals)
    )
    if not learner_records:
        raise ValueError("learner comparison requires at least one learner proposal")
    return SupervisorLearnerProposalComparison(
        supervisor=supervisor_record,
        learner_proposals=learner_records,
        comparison_label=comparison_label,
        metrics=_learner_proposal_comparison_metrics(learner_records),
        actuation_permitted=False,
    )

compare_supervisor_random_baseline 

compare_supervisor_random_baseline(
    policy: DifferentiableSupervisorPolicy,
    scenario: KuramotoSupervisorScenario,
    *,
    key: Array,
    comparison_label: str = "bounded_random_action",
) -> SupervisorRandomBaselineComparison

Compare one deterministic supervisor proposal against bounded randomness.

Parameters

policy : DifferentiableSupervisorPolicy The differentiable supervisor policy. scenario : KuramotoSupervisorScenario The Kuramoto supervisor scenario. key : jax.Array JAX PRNG key. comparison_label : str Label identifying the comparison.

Returns

SupervisorRandomBaselineComparison The supervisor-vs-random comparison.

Raises

ValueError If the inputs are invalid or inconsistent.

Source code in src/scpn_phase_orchestrator/nn/supervisor/comparison.py 
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
def compare_supervisor_random_baseline(
    policy: DifferentiableSupervisorPolicy,
    scenario: KuramotoSupervisorScenario,
    *,
    key: jax.Array,
    comparison_label: str = "bounded_random_action",
) -> SupervisorRandomBaselineComparison:
    """Compare one deterministic supervisor proposal against bounded randomness.

    Parameters
    ----------
    policy : DifferentiableSupervisorPolicy
        The differentiable supervisor policy.
    scenario : KuramotoSupervisorScenario
        The Kuramoto supervisor scenario.
    key : jax.Array
        JAX PRNG key.
    comparison_label : str
        Label identifying the comparison.

    Returns
    -------
    SupervisorRandomBaselineComparison
        The supervisor-vs-random comparison.

    Raises
    ------
    ValueError
        If the inputs are invalid or inconsistent.
    """
    if not comparison_label:
        raise ValueError("comparison_label must not be empty")
    random_action = _bounded_random_supervisor_action(policy.config, key)
    baseline_record = _rollout_static_supervisor_action(
        "bounded_random_action",
        random_action,
        scenario,
        policy.config,
    )
    baseline_record["seed"] = _jax_key_record(key)
    supervisor_record = _rollout_static_supervisor_action(
        "differentiable_supervisor",
        policy(scenario),
        scenario,
        policy.config,
    )
    baseline_metrics = _mapping_value(baseline_record, "metrics")
    supervisor_metrics = _mapping_value(supervisor_record, "metrics")
    metrics = _prefixed_float_metrics("baseline", baseline_metrics)
    metrics.update(_prefixed_float_metrics("supervisor", supervisor_metrics))
    metrics["delta_reward"] = metrics["supervisor_reward"] - metrics["baseline_reward"]
    return SupervisorRandomBaselineComparison(
        baseline=baseline_record,
        supervisor=supervisor_record,
        scenario_summary=_supervisor_scenario_summary(scenario),
        comparison_label=comparison_label,
        metrics=metrics,
        actuation_permitted=False,
    )

compare_supervisor_replay_proposal 

compare_supervisor_replay_proposal(
    supervisor_proposal: SupervisorReplayProposal
    | SupervisorCorpusReplayProposals,
    replay_policy_search: Any,
    *,
    comparison_label: str = "replay_policy_search",
) -> SupervisorReplayComparison

Compare supervisor replay proposals with a replay policy-search result.

The result is deliberately audit-only: it records both proposal surfaces and scalar comparison metrics, but it never authorises live actuation.

Parameters

supervisor_proposal : SupervisorReplayProposal | SupervisorCorpusReplayProposals The supervisor replay proposal(s). replay_policy_search : Any The replay policy-search result. comparison_label : str Label identifying the comparison.

Returns

SupervisorReplayComparison The supervisor-vs-policy-search comparison.

Raises

ValueError If the inputs are invalid or inconsistent.

Source code in src/scpn_phase_orchestrator/nn/supervisor/comparison.py 
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
def compare_supervisor_replay_proposal(
    supervisor_proposal: SupervisorReplayProposal | SupervisorCorpusReplayProposals,
    replay_policy_search: Any,
    *,
    comparison_label: str = "replay_policy_search",
) -> SupervisorReplayComparison:
    """Compare supervisor replay proposals with a replay policy-search result.

    The result is deliberately audit-only: it records both proposal surfaces and
    scalar comparison metrics, but it never authorises live actuation.

    Parameters
    ----------
    supervisor_proposal : SupervisorReplayProposal | SupervisorCorpusReplayProposals
        The supervisor replay proposal(s).
    replay_policy_search : Any
        The replay policy-search result.
    comparison_label : str
        Label identifying the comparison.

    Returns
    -------
    SupervisorReplayComparison
        The supervisor-vs-policy-search comparison.

    Raises
    ------
    ValueError
        If the inputs are invalid or inconsistent.
    """
    if not comparison_label:
        raise ValueError("comparison_label must not be empty")
    supervisor_record = _json_object(
        supervisor_proposal.to_audit_record(),
        "supervisor proposal audit record",
    )
    replay_record = _audit_record_from_object(
        replay_policy_search,
        "replay_policy_search",
    )
    metrics = _supervisor_replay_comparison_metrics(
        supervisor_record,
        replay_record,
    )
    return SupervisorReplayComparison(
        supervisor=supervisor_record,
        replay_policy_search=replay_record,
        comparison_label=comparison_label,
        metrics=metrics,
        actuation_permitted=False,
    )

compare_supervisor_static_baseline 

compare_supervisor_static_baseline(
    policy: DifferentiableSupervisorPolicy,
    scenario: KuramotoSupervisorScenario,
    *,
    comparison_label: str = "static_zero_action",
) -> SupervisorStaticBaselineComparison

Compare one deterministic supervisor proposal against zero-action control.

This is a benchmark/audit primitive only. It runs both candidates on the same scenario and records scalar metrics without returning actuation objects or enabling any live adapter handoff.

Parameters

policy : DifferentiableSupervisorPolicy The differentiable supervisor policy. scenario : KuramotoSupervisorScenario The Kuramoto supervisor scenario. comparison_label : str Label identifying the comparison.

Returns

SupervisorStaticBaselineComparison The supervisor-vs-zero-action comparison.

Raises

ValueError If the inputs are invalid or inconsistent.

Source code in src/scpn_phase_orchestrator/nn/supervisor/comparison.py 
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
def compare_supervisor_static_baseline(
    policy: DifferentiableSupervisorPolicy,
    scenario: KuramotoSupervisorScenario,
    *,
    comparison_label: str = "static_zero_action",
) -> SupervisorStaticBaselineComparison:
    """Compare one deterministic supervisor proposal against zero-action control.

    This is a benchmark/audit primitive only. It runs both candidates on the
    same scenario and records scalar metrics without returning actuation
    objects or enabling any live adapter handoff.

    Parameters
    ----------
    policy : DifferentiableSupervisorPolicy
        The differentiable supervisor policy.
    scenario : KuramotoSupervisorScenario
        The Kuramoto supervisor scenario.
    comparison_label : str
        Label identifying the comparison.

    Returns
    -------
    SupervisorStaticBaselineComparison
        The supervisor-vs-zero-action comparison.

    Raises
    ------
    ValueError
        If the inputs are invalid or inconsistent.
    """
    if not comparison_label:
        raise ValueError("comparison_label must not be empty")
    zero_action = SupervisorAction(
        delta_K_global=jnp.array(0.0),
        delta_zeta_global=jnp.array(0.0),
        delta_K_layers=jnp.zeros(policy.config.n_layer_controls),
        value_estimate=jnp.array(0.0),
    )
    baseline_record = _rollout_static_supervisor_action(
        "static_zero_action",
        zero_action,
        scenario,
        policy.config,
    )
    supervisor_record = _rollout_static_supervisor_action(
        "differentiable_supervisor",
        policy(scenario),
        scenario,
        policy.config,
    )
    baseline_metrics = _mapping_value(baseline_record, "metrics")
    supervisor_metrics = _mapping_value(supervisor_record, "metrics")
    metrics = _prefixed_float_metrics("baseline", baseline_metrics)
    metrics.update(_prefixed_float_metrics("supervisor", supervisor_metrics))
    metrics["delta_reward"] = metrics["supervisor_reward"] - metrics["baseline_reward"]
    return SupervisorStaticBaselineComparison(
        baseline=baseline_record,
        supervisor=supervisor_record,
        scenario_summary=_supervisor_scenario_summary(scenario),
        comparison_label=comparison_label,
        metrics=metrics,
        actuation_permitted=False,
    )

apply_supervisor_action 

apply_supervisor_action(
    base_K: Array,
    action: SupervisorAction,
    scenario: KuramotoSupervisorScenario,
) -> jax.Array

Apply continuous supervisor output to a symmetric coupling matrix.

Parameters

base_K : jax.Array Base symmetric coupling matrix. action : SupervisorAction The supervisor control action. scenario : KuramotoSupervisorScenario The Kuramoto supervisor scenario.

Returns

jax.Array The modified symmetric coupling matrix.

Source code in src/scpn_phase_orchestrator/nn/supervisor/policy.py 
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
def apply_supervisor_action(
    base_K: jax.Array,
    action: SupervisorAction,
    scenario: KuramotoSupervisorScenario,
) -> jax.Array:
    """Apply continuous supervisor output to a symmetric coupling matrix.

    Parameters
    ----------
    base_K : jax.Array
        Base symmetric coupling matrix.
    action : SupervisorAction
        The supervisor control action.
    scenario : KuramotoSupervisorScenario
        The Kuramoto supervisor scenario.

    Returns
    -------
    jax.Array
        The modified symmetric coupling matrix.
    """
    n = base_K.shape[0]
    offdiag = 1.0 - jnp.eye(n, dtype=base_K.dtype)
    K = base_K + action.delta_K_global * offdiag

    layer_masks = (scenario.good_mask, scenario.bad_mask)
    for idx, mask in enumerate(layer_masks[: action.delta_K_layers.shape[0]]):
        membership = jnp.outer(mask, mask) * offdiag
        K = K + action.delta_K_layers[idx] * membership

    K = 0.5 * (K + K.T)
    return K * offdiag

closed_loop_supervisor_loss 

closed_loop_supervisor_loss(
    policy: DifferentiableSupervisorPolicy,
    scenario: KuramotoSupervisorScenario,
) -> tuple[jax.Array, SupervisorLossAux]

Differentiable closed-loop objective for Kuramoto supervisor training.

The reward maximises good-partition synchrony while penalising bad-partition synchrony, control energy, and abrupt action changes. The returned value is a minimisation loss suitable for jax.grad or optax.

Parameters

policy : DifferentiableSupervisorPolicy The differentiable supervisor policy. scenario : KuramotoSupervisorScenario The Kuramoto supervisor scenario.

Returns

tuple[jax.Array, SupervisorLossAux] The loss and its auxiliary metrics.

Source code in src/scpn_phase_orchestrator/nn/supervisor/policy.py 
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
def closed_loop_supervisor_loss(
    policy: DifferentiableSupervisorPolicy,
    scenario: KuramotoSupervisorScenario,
) -> tuple[jax.Array, SupervisorLossAux]:
    """Differentiable closed-loop objective for Kuramoto supervisor training.

    The reward maximises good-partition synchrony while penalising bad-partition
    synchrony, control energy, and abrupt action changes. The returned value is
    a minimisation loss suitable for ``jax.grad`` or optax.

    Parameters
    ----------
    policy : DifferentiableSupervisorPolicy
        The differentiable supervisor policy.
    scenario : KuramotoSupervisorScenario
        The Kuramoto supervisor scenario.

    Returns
    -------
    tuple[jax.Array, SupervisorLossAux]
        The loss and its auxiliary metrics.
    """
    zero_action = SupervisorAction(
        delta_K_global=jnp.array(0.0),
        delta_zeta_global=jnp.array(0.0),
        delta_K_layers=jnp.zeros(policy.config.n_layer_controls),
        value_estimate=jnp.array(0.0),
    )

    def body(
        carry: tuple[jax.Array, SupervisorAction],
        _: None,
    ) -> tuple[tuple[jax.Array, SupervisorAction], tuple[jax.Array, jax.Array]]:
        """Apply the policy network body to the input features."""
        phases, previous_action = carry
        step_scenario = scenario._replace(phases=phases)
        action = policy(step_scenario)
        controlled_K = apply_supervisor_action(scenario.base_K, action, scenario)
        final, _traj = kuramoto_forward(
            phases,
            scenario.omegas,
            controlled_K,
            scenario.dt,
            scenario.inner_steps,
        )
        energy = _control_energy(action)
        smoothness = _action_distance(action, previous_action)
        return (final, action), (energy, smoothness)

    (final_phases, _), (energies, smoothnesses) = jax.lax.scan(
        body,
        (scenario.phases, zero_action),
        None,
        length=scenario.horizon,
    )
    final_R_good = masked_order_parameter(final_phases, scenario.good_mask)
    final_R_bad = masked_order_parameter(final_phases, scenario.bad_mask)
    control_energy = jnp.mean(energies)
    smoothness = jnp.mean(smoothnesses)
    reward = final_R_good - policy.config.bad_sync_weight * final_R_bad
    loss = (
        -reward
        + policy.config.control_energy_weight * control_energy
        + policy.config.smoothness_weight * smoothness
    )
    aux = SupervisorLossAux(
        final_R_good=final_R_good,
        final_R_bad=final_R_bad,
        control_energy=control_energy,
        smoothness=smoothness,
    )
    return loss, aux

control_actions_from_supervisor 

control_actions_from_supervisor(
    action: SupervisorAction,
    *,
    ttl_s: float = 5.0,
    include_layer_actions: bool = True,
) -> list[ControlAction]

Convert a detached neural supervisor output into actuation commands.

Parameters

action : SupervisorAction The supervisor control action. ttl_s : float Action time-to-live in seconds. include_layer_actions : bool Whether to include per-layer actions.

Returns

list[ControlAction] The actuation control actions.

Source code in src/scpn_phase_orchestrator/nn/supervisor/policy.py 
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
def control_actions_from_supervisor(
    action: SupervisorAction,
    *,
    ttl_s: float = 5.0,
    include_layer_actions: bool = True,
) -> list[ControlAction]:
    """Convert a detached neural supervisor output into actuation commands.

    Parameters
    ----------
    action : SupervisorAction
        The supervisor control action.
    ttl_s : float
        Action time-to-live in seconds.
    include_layer_actions : bool
        Whether to include per-layer actions.

    Returns
    -------
    list[ControlAction]
        The actuation control actions.
    """
    actions = [
        ControlAction(
            knob="K",
            scope="global",
            value=float(action.delta_K_global),
            ttl_s=ttl_s,
            justification="differentiable supervisor: global coupling proposal",
        ),
        ControlAction(
            knob="zeta",
            scope="global",
            value=float(action.delta_zeta_global),
            ttl_s=ttl_s,
            justification="differentiable supervisor: damping proposal",
        ),
    ]
    if include_layer_actions:
        for idx, value in enumerate(action.delta_K_layers):
            actions.append(
                ControlAction(
                    knob="K",
                    scope=f"layer_{idx}",
                    value=float(value),
                    ttl_s=ttl_s,
                    justification=(
                        "differentiable supervisor: partition coupling proposal"
                    ),
                )
            )
    return actions

pack_supervisor_action 

pack_supervisor_action(
    action: SupervisorAction,
) -> jax.Array

Pack SupervisorAction controls into a flat continuous action vector.

Parameters

action : SupervisorAction The supervisor control action.

Returns

jax.Array The flat continuous action vector.

Source code in src/scpn_phase_orchestrator/nn/supervisor/policy.py 
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
def pack_supervisor_action(action: SupervisorAction) -> jax.Array:
    """Pack ``SupervisorAction`` controls into a flat continuous action vector.

    Parameters
    ----------
    action : SupervisorAction
        The supervisor control action.

    Returns
    -------
    jax.Array
        The flat continuous action vector.
    """
    return jnp.concatenate(
        [
            jnp.atleast_1d(action.delta_K_global),
            jnp.atleast_1d(action.delta_zeta_global),
            action.delta_K_layers,
        ]
    )

sample_supervisor_action 

sample_supervisor_action(
    policy: DifferentiableSupervisorPolicy,
    scenario: KuramotoSupervisorScenario,
    *,
    key: Array,
) -> tuple[SupervisorAction, jax.Array]

Sample a bounded squashed-Gaussian action and its log probability.

Parameters

policy : DifferentiableSupervisorPolicy The differentiable supervisor policy. scenario : KuramotoSupervisorScenario The Kuramoto supervisor scenario. key : jax.Array JAX PRNG key.

Returns

tuple[SupervisorAction, jax.Array] The sampled action and its log probability.

Source code in src/scpn_phase_orchestrator/nn/supervisor/policy.py 
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
def sample_supervisor_action(
    policy: DifferentiableSupervisorPolicy,
    scenario: KuramotoSupervisorScenario,
    *,
    key: jax.Array,
) -> tuple[SupervisorAction, jax.Array]:
    """Sample a bounded squashed-Gaussian action and its log probability.

    Parameters
    ----------
    policy : DifferentiableSupervisorPolicy
        The differentiable supervisor policy.
    scenario : KuramotoSupervisorScenario
        The Kuramoto supervisor scenario.
    key : jax.Array
        JAX PRNG key.

    Returns
    -------
    tuple[SupervisorAction, jax.Array]
        The sampled action and its log probability.
    """
    mean, value = _policy_mean_and_value(policy, scenario)
    std = jnp.exp(policy.log_std)
    pre_squash = mean + std * jax.random.normal(key, mean.shape)
    bounded = jnp.tanh(pre_squash) * _action_bounds(policy.config)
    action = unpack_supervisor_action(
        bounded,
        value_estimate=value,
        config=policy.config,
    )
    return action, _squashed_gaussian_log_prob(mean, policy.log_std, pre_squash)

supervisor_action_bound_penalty 

supervisor_action_bound_penalty(
    action: SupervisorAction,
    config: DifferentiableSupervisorConfig,
) -> jax.Array

Differentiable quadratic penalty for proposals outside action bounds.

Parameters

action : SupervisorAction The supervisor control action. config : DifferentiableSupervisorConfig The supervisor configuration.

Returns

jax.Array The quadratic out-of-bounds penalty.

Source code in src/scpn_phase_orchestrator/nn/supervisor/policy.py 
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
def supervisor_action_bound_penalty(
    action: SupervisorAction,
    config: DifferentiableSupervisorConfig,
) -> jax.Array:
    """Differentiable quadratic penalty for proposals outside action bounds.

    Parameters
    ----------
    action : SupervisorAction
        The supervisor control action.
    config : DifferentiableSupervisorConfig
        The supervisor configuration.

    Returns
    -------
    jax.Array
        The quadratic out-of-bounds penalty.
    """
    values = pack_supervisor_action(action)
    bounds = _action_bounds(config)
    excess = jnp.maximum(jnp.abs(values) - bounds, 0.0)
    return jnp.sum(excess**2)

supervisor_action_log_prob 

supervisor_action_log_prob(
    policy: DifferentiableSupervisorPolicy,
    scenario: KuramotoSupervisorScenario,
    action: SupervisorAction,
) -> tuple[jax.Array, jax.Array, jax.Array]

Return squashed-Gaussian log probability, entropy proxy, and value.

Parameters

policy : DifferentiableSupervisorPolicy The differentiable supervisor policy. scenario : KuramotoSupervisorScenario The Kuramoto supervisor scenario. action : SupervisorAction The supervisor control action.

Returns

tuple[jax.Array, jax.Array, jax.Array] The log probability, entropy proxy, and value.

Source code in src/scpn_phase_orchestrator/nn/supervisor/policy.py 
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
def supervisor_action_log_prob(
    policy: DifferentiableSupervisorPolicy,
    scenario: KuramotoSupervisorScenario,
    action: SupervisorAction,
) -> tuple[jax.Array, jax.Array, jax.Array]:
    """Return squashed-Gaussian log probability, entropy proxy, and value.

    Parameters
    ----------
    policy : DifferentiableSupervisorPolicy
        The differentiable supervisor policy.
    scenario : KuramotoSupervisorScenario
        The Kuramoto supervisor scenario.
    action : SupervisorAction
        The supervisor control action.

    Returns
    -------
    tuple[jax.Array, jax.Array, jax.Array]
        The log probability, entropy proxy, and value.
    """
    mean, value = _policy_mean_and_value(policy, scenario)
    bounds = _action_bounds(policy.config)
    scaled = jnp.clip(pack_supervisor_action(action) / bounds, -0.999999, 0.999999)
    pre_squash = jnp.arctanh(scaled)
    log_prob = _squashed_gaussian_log_prob(mean, policy.log_std, pre_squash)
    entropy = 0.5 * jnp.sum(1.0 + jnp.log(2.0 * jnp.pi) + 2.0 * policy.log_std)
    return log_prob, entropy, value

supervisor_train_step 

supervisor_train_step(
    policy: DifferentiableSupervisorPolicy,
    scenario: KuramotoSupervisorScenario,
    opt_state: Any,
    optimizer: GradientTransformation,
) -> tuple[DifferentiableSupervisorPolicy, Any, jax.Array]

Run one optax update for the differentiable supervisor objective.

Parameters

policy : DifferentiableSupervisorPolicy The differentiable supervisor policy. scenario : KuramotoSupervisorScenario The Kuramoto supervisor scenario. opt_state : Any The optax optimiser state. optimizer : optax.GradientTransformation The optax optimiser.

Returns

tuple[DifferentiableSupervisorPolicy, Any, jax.Array] The updated policy, optimiser state, and loss.

Source code in src/scpn_phase_orchestrator/nn/supervisor/policy.py 
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
def supervisor_train_step(
    policy: DifferentiableSupervisorPolicy,
    scenario: KuramotoSupervisorScenario,
    opt_state: Any,
    optimizer: optax.GradientTransformation,
) -> tuple[DifferentiableSupervisorPolicy, Any, jax.Array]:
    """Run one optax update for the differentiable supervisor objective.

    Parameters
    ----------
    policy : DifferentiableSupervisorPolicy
        The differentiable supervisor policy.
    scenario : KuramotoSupervisorScenario
        The Kuramoto supervisor scenario.
    opt_state : Any
        The optax optimiser state.
    optimizer : optax.GradientTransformation
        The optax optimiser.

    Returns
    -------
    tuple[DifferentiableSupervisorPolicy, Any, jax.Array]
        The updated policy, optimiser state, and loss.
    """

    def loss_fn(model: DifferentiableSupervisorPolicy) -> jax.Array:
        """Return the policy loss for a batch under the current parameters."""
        loss, _ = closed_loop_supervisor_loss(model, scenario)
        return loss

    loss, grads = eqx.filter_value_and_grad(loss_fn)(policy)
    params = eqx.filter(policy, eqx.is_array)
    updates, opt_state = optimizer.update(grads, opt_state, params)
    updated = eqx.apply_updates(policy, updates)
    return updated, opt_state, loss

unpack_supervisor_action 

unpack_supervisor_action(
    values: Array,
    *,
    value_estimate: Array,
    config: DifferentiableSupervisorConfig,
) -> SupervisorAction

Unpack a flat action vector using config.n_layer_controls.

Parameters

values : jax.Array Flat packed action values. value_estimate : jax.Array The critic value estimate. config : DifferentiableSupervisorConfig The supervisor configuration.

Returns

SupervisorAction The reconstructed SupervisorAction.

Source code in src/scpn_phase_orchestrator/nn/supervisor/policy.py 
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
def unpack_supervisor_action(
    values: jax.Array,
    *,
    value_estimate: jax.Array,
    config: DifferentiableSupervisorConfig,
) -> SupervisorAction:
    """Unpack a flat action vector using ``config.n_layer_controls``.

    Parameters
    ----------
    values : jax.Array
        Flat packed action values.
    value_estimate : jax.Array
        The critic value estimate.
    config : DifferentiableSupervisorConfig
        The supervisor configuration.

    Returns
    -------
    SupervisorAction
        The reconstructed ``SupervisorAction``.
    """
    return SupervisorAction(
        delta_K_global=values[0],
        delta_zeta_global=values[1],
        delta_K_layers=values[2 : 2 + config.n_layer_controls],
        value_estimate=value_estimate,
    )

ppo_supervisor_loss 

ppo_supervisor_loss(
    policy: DifferentiableSupervisorPolicy,
    batch: SupervisorPPOBatch,
    *,
    clip_epsilon: float = 0.2,
    value_clip: float | None = None,
    value_weight: float = 0.5,
    entropy_weight: float = 0.01,
) -> tuple[jax.Array, SupervisorPPOAux]

Clipped PPO objective for bounded differentiable supervisor actions.

Parameters

policy : DifferentiableSupervisorPolicy The differentiable supervisor policy. batch : SupervisorPPOBatch The PPO training batch. clip_epsilon : float PPO clipping epsilon. value_clip : float | None Value-function clip range, or None. value_weight : float Weight of the value loss. entropy_weight : float Weight of the entropy bonus.

Returns

tuple[jax.Array, SupervisorPPOAux] The clipped PPO loss and its auxiliary metrics.

Source code in src/scpn_phase_orchestrator/nn/supervisor/ppo.py 
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
def ppo_supervisor_loss(
    policy: DifferentiableSupervisorPolicy,
    batch: SupervisorPPOBatch,
    *,
    clip_epsilon: float = 0.2,
    value_clip: float | None = None,
    value_weight: float = 0.5,
    entropy_weight: float = 0.01,
) -> tuple[jax.Array, SupervisorPPOAux]:
    """Clipped PPO objective for bounded differentiable supervisor actions.

    Parameters
    ----------
    policy : DifferentiableSupervisorPolicy
        The differentiable supervisor policy.
    batch : SupervisorPPOBatch
        The PPO training batch.
    clip_epsilon : float
        PPO clipping epsilon.
    value_clip : float | None
        Value-function clip range, or ``None``.
    value_weight : float
        Weight of the value loss.
    entropy_weight : float
        Weight of the entropy bonus.

    Returns
    -------
    tuple[jax.Array, SupervisorPPOAux]
        The clipped PPO loss and its auxiliary metrics.
    """
    if value_clip is not None:
        value_clip = _non_negative_float(value_clip, "value_clip")

    def item_loss(
        phases: jax.Array,
        omegas: jax.Array,
        base_K: jax.Array,
        good_mask: jax.Array,
        bad_mask: jax.Array,
        action_values: jax.Array,
        old_log_prob: jax.Array,
        advantage: jax.Array,
        ret: jax.Array,
        old_value: jax.Array,
    ) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array, jax.Array]:
        """Return the per-item PPO surrogate loss."""
        scenario = KuramotoSupervisorScenario(
            phases=phases,
            omegas=omegas,
            base_K=base_K,
            good_mask=good_mask,
            bad_mask=bad_mask,
            dt=batch.dt,
            inner_steps=batch.inner_steps,
            horizon=batch.horizon,
        )
        action = unpack_supervisor_action(
            action_values,
            value_estimate=jnp.array(0.0),
            config=policy.config,
        )
        log_prob, entropy, value = supervisor_action_log_prob(policy, scenario, action)
        ratio = jnp.exp(log_prob - old_log_prob)
        unclipped = ratio * advantage
        clipped = jnp.clip(ratio, 1.0 - clip_epsilon, 1.0 + clip_epsilon) * advantage
        policy_loss = -jnp.minimum(unclipped, clipped)
        unclipped_value_loss = (value - ret) ** 2
        if value_clip is not None:
            clipped_value = old_value + jnp.clip(
                value - old_value, -value_clip, value_clip
            )
            clipped_value_loss = (clipped_value - ret) ** 2
            value_loss = jnp.maximum(unclipped_value_loss, clipped_value_loss)
        else:
            value_loss = unclipped_value_loss
        approx_kl = old_log_prob - log_prob
        clipped_flag = jnp.abs(ratio - 1.0) > clip_epsilon
        return (
            policy_loss,
            value_loss,
            entropy,
            approx_kl,
            clipped_flag.astype(jnp.float32),
        )

    policy_losses, value_losses, entropies, approx_kls, clip_flags = jax.vmap(
        item_loss
    )(
        batch.phases,
        batch.omegas,
        batch.base_K,
        batch.good_mask,
        batch.bad_mask,
        batch.actions,
        batch.old_log_probs,
        batch.advantages,
        batch.returns,
        batch.values,
    )
    policy_loss = jnp.mean(policy_losses)
    value_loss = jnp.mean(value_losses)
    entropy = jnp.mean(entropies)
    total = policy_loss + value_weight * value_loss - entropy_weight * entropy
    aux = SupervisorPPOAux(
        policy_loss=policy_loss,
        value_loss=value_loss,
        entropy=entropy,
        approx_kl=jnp.mean(approx_kls),
        clip_fraction=jnp.mean(clip_flags),
    )
    return total, aux

ppo_supervisor_train_epochs 

ppo_supervisor_train_epochs(
    policy: DifferentiableSupervisorPolicy,
    batch: SupervisorPPOBatch,
    key: Array,
    opt_state: Any,
    optimizer: GradientTransformation,
    *,
    n_epochs: int,
    minibatch_size: int = 32,
    clip_epsilon: float = 0.2,
    value_clip: float | None = None,
    value_weight: float = 0.5,
    entropy_weight: float = 0.01,
    entropy_schedule: tuple[float, ...] | None = None,
    max_grad_norm: float | None = None,
    kl_early_stop: float | None = None,
) -> tuple[
    DifferentiableSupervisorPolicy, Any, jax.Array, int
]

Run PPO training for multiple epochs with deterministic minibatching.

Parameters

policy : DifferentiableSupervisorPolicy Optimisable differentiable supervisor policy. batch : SupervisorPPOBatch Flattened PPO batch from rollout collection. key : jax.Array JAX PRNG key used only for shuffle ordering. opt_state : Any Optimiser state. optimizer : optax.GradientTransformation Optax optimiser transformation. n_epochs : int Number of passes over the dataset. minibatch_size : int Per-update minibatch size. clip_epsilon : float PPO clipping radius. value_weight : float Value-function loss coefficient. entropy_weight : float Entropy bonus coefficient. entropy_schedule : tuple[float, ...] | None Optional non-negative per-update entropy weights. When shorter than the number of updates, the final value is held. max_grad_norm : float | None Optional global gradient-norm clip radius. kl_early_stop : float | None Optional KL threshold for early stopping.

Returns

tuple[DifferentiableSupervisorPolicy, Any, jax.Array, int] (policy, opt_state, loss_history, n_updates).

Source code in src/scpn_phase_orchestrator/nn/supervisor/ppo.py 
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
def ppo_supervisor_train_epochs(
    policy: DifferentiableSupervisorPolicy,
    batch: SupervisorPPOBatch,
    key: jax.Array,
    opt_state: Any,
    optimizer: optax.GradientTransformation,
    *,
    n_epochs: int,
    minibatch_size: int = 32,
    clip_epsilon: float = 0.2,
    value_clip: float | None = None,
    value_weight: float = 0.5,
    entropy_weight: float = 0.01,
    entropy_schedule: tuple[float, ...] | None = None,
    max_grad_norm: float | None = None,
    kl_early_stop: float | None = None,
) -> tuple[DifferentiableSupervisorPolicy, Any, jax.Array, int]:
    """Run PPO training for multiple epochs with deterministic minibatching.

    Parameters
    ----------
    policy : DifferentiableSupervisorPolicy
        Optimisable differentiable supervisor policy.
    batch : SupervisorPPOBatch
        Flattened PPO batch from rollout collection.
    key : jax.Array
        JAX PRNG key used only for shuffle ordering.
    opt_state : Any
        Optimiser state.
    optimizer : optax.GradientTransformation
        Optax optimiser transformation.
    n_epochs : int
        Number of passes over the dataset.
    minibatch_size : int
        Per-update minibatch size.
    clip_epsilon : float
        PPO clipping radius.
    value_weight : float
        Value-function loss coefficient.
    entropy_weight : float
        Entropy bonus coefficient.
    entropy_schedule : tuple[float, ...] | None
        Optional non-negative per-update entropy weights. When shorter than the number
        of updates, the final value is held.
    max_grad_norm : float | None
        Optional global gradient-norm clip radius.
    kl_early_stop : float | None
        Optional KL threshold for early stopping.

    Returns
    -------
    tuple[DifferentiableSupervisorPolicy, Any, jax.Array, int]
        (policy, opt_state, loss_history, n_updates).
    """
    policy, opt_state, loss_history, n_updates, _ = _ppo_supervisor_train_epochs_impl(
        policy,
        batch,
        key,
        opt_state,
        optimizer,
        n_epochs=n_epochs,
        minibatch_size=minibatch_size,
        clip_epsilon=clip_epsilon,
        value_clip=value_clip,
        value_weight=value_weight,
        entropy_weight=entropy_weight,
        entropy_schedule=entropy_schedule,
        max_grad_norm=max_grad_norm,
        kl_early_stop=kl_early_stop,
    )
    return policy, opt_state, loss_history, n_updates

ppo_supervisor_train_step 

ppo_supervisor_train_step(
    policy: DifferentiableSupervisorPolicy,
    batch: SupervisorPPOBatch,
    opt_state: Any,
    optimizer: GradientTransformation,
    *,
    clip_epsilon: float = 0.2,
    value_clip: float | None = None,
    value_weight: float = 0.5,
    entropy_weight: float = 0.01,
    max_grad_norm: float | None = None,
) -> tuple[DifferentiableSupervisorPolicy, Any, jax.Array]

Run one optax update using the clipped PPO supervisor objective.

Parameters

policy : DifferentiableSupervisorPolicy The differentiable supervisor policy. batch : SupervisorPPOBatch The PPO training batch. opt_state : Any The optax optimiser state. optimizer : optax.GradientTransformation The optax optimiser. clip_epsilon : float PPO clipping epsilon. value_clip : float | None Value-function clip range, or None. value_weight : float Weight of the value loss. entropy_weight : float Weight of the entropy bonus. max_grad_norm : float | None Gradient-norm clip threshold, or None.

Returns

tuple[DifferentiableSupervisorPolicy, Any, jax.Array] The updated policy, optimiser state, and loss.

Source code in src/scpn_phase_orchestrator/nn/supervisor/ppo.py 
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
def ppo_supervisor_train_step(
    policy: DifferentiableSupervisorPolicy,
    batch: SupervisorPPOBatch,
    opt_state: Any,
    optimizer: optax.GradientTransformation,
    *,
    clip_epsilon: float = 0.2,
    value_clip: float | None = None,
    value_weight: float = 0.5,
    entropy_weight: float = 0.01,
    max_grad_norm: float | None = None,
) -> tuple[DifferentiableSupervisorPolicy, Any, jax.Array]:
    """Run one optax update using the clipped PPO supervisor objective.

    Parameters
    ----------
    policy : DifferentiableSupervisorPolicy
        The differentiable supervisor policy.
    batch : SupervisorPPOBatch
        The PPO training batch.
    opt_state : Any
        The optax optimiser state.
    optimizer : optax.GradientTransformation
        The optax optimiser.
    clip_epsilon : float
        PPO clipping epsilon.
    value_clip : float | None
        Value-function clip range, or ``None``.
    value_weight : float
        Weight of the value loss.
    entropy_weight : float
        Weight of the entropy bonus.
    max_grad_norm : float | None
        Gradient-norm clip threshold, or ``None``.

    Returns
    -------
    tuple[DifferentiableSupervisorPolicy, Any, jax.Array]
        The updated policy, optimiser state, and loss.
    """

    def loss_fn(model: DifferentiableSupervisorPolicy) -> jax.Array:
        """Return the PPO loss for a batch under the current parameters."""
        loss, _ = ppo_supervisor_loss(
            model,
            batch,
            clip_epsilon=clip_epsilon,
            value_clip=value_clip,
            value_weight=value_weight,
            entropy_weight=entropy_weight,
        )
        return loss

    loss, grads = eqx.filter_value_and_grad(loss_fn)(policy)
    params = eqx.filter(policy, eqx.is_array)
    updates, opt_state = optimizer.update(grads, opt_state, params)
    if max_grad_norm is not None:
        max_grad_norm = _positive_float(max_grad_norm, "max_grad_norm")
        updates = _clip_updates_by_global_norm(updates, max_grad_norm)
    updated = eqx.apply_updates(policy, updates)
    return updated, opt_state, loss

ppo_supervisor_train_with_checkpoint 

ppo_supervisor_train_with_checkpoint(
    policy: DifferentiableSupervisorPolicy,
    batch: SupervisorPPOBatch,
    key: Array,
    opt_state: Any,
    optimizer: GradientTransformation,
    *,
    n_epochs: int,
    checkpoint_dir: str | Path | None = None,
    resume: bool = False,
    minibatch_size: int = 32,
    clip_epsilon: float = 0.2,
    value_clip: float | None = None,
    value_weight: float = 0.5,
    entropy_weight: float = 0.01,
    entropy_schedule: tuple[float, ...] | None = None,
    max_grad_norm: float | None = None,
    kl_early_stop: float | None = None,
    metadata: dict[str, Any] | None = None,
) -> SupervisorPPOTrainResult

Run PPO epochs and optionally checkpoint a deterministic resume state.

Parameters

policy : DifferentiableSupervisorPolicy The differentiable supervisor policy. batch : SupervisorPPOBatch The PPO training batch. key : jax.Array JAX PRNG key. opt_state : Any The optax optimiser state. optimizer : optax.GradientTransformation The optax optimiser. n_epochs : int Number of training epochs. checkpoint_dir : str | Path | None Directory for training checkpoints, or None. resume : bool Whether to resume from a checkpoint. minibatch_size : int Minibatch size. clip_epsilon : float PPO clipping epsilon. value_clip : float | None Value-function clip range, or None. value_weight : float Weight of the value loss. entropy_weight : float Weight of the entropy bonus. entropy_schedule : tuple[float, ...] | None Per-epoch entropy-weight schedule, or None. max_grad_norm : float | None Gradient-norm clip threshold, or None. kl_early_stop : float | None KL early-stopping threshold, or None. metadata : dict[str, Any] | None Associated metadata mapping, or None.

Returns

SupervisorPPOTrainResult The PPO training result.

Raises

ValueError If the inputs are invalid or inconsistent.

Source code in src/scpn_phase_orchestrator/nn/supervisor/ppo.py 
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
def ppo_supervisor_train_with_checkpoint(
    policy: DifferentiableSupervisorPolicy,
    batch: SupervisorPPOBatch,
    key: jax.Array,
    opt_state: Any,
    optimizer: optax.GradientTransformation,
    *,
    n_epochs: int,
    checkpoint_dir: str | Path | None = None,
    resume: bool = False,
    minibatch_size: int = 32,
    clip_epsilon: float = 0.2,
    value_clip: float | None = None,
    value_weight: float = 0.5,
    entropy_weight: float = 0.01,
    entropy_schedule: tuple[float, ...] | None = None,
    max_grad_norm: float | None = None,
    kl_early_stop: float | None = None,
    metadata: dict[str, Any] | None = None,
) -> SupervisorPPOTrainResult:
    """Run PPO epochs and optionally checkpoint a deterministic resume state.

    Parameters
    ----------
    policy : DifferentiableSupervisorPolicy
        The differentiable supervisor policy.
    batch : SupervisorPPOBatch
        The PPO training batch.
    key : jax.Array
        JAX PRNG key.
    opt_state : Any
        The optax optimiser state.
    optimizer : optax.GradientTransformation
        The optax optimiser.
    n_epochs : int
        Number of training epochs.
    checkpoint_dir : str | Path | None
        Directory for training checkpoints, or ``None``.
    resume : bool
        Whether to resume from a checkpoint.
    minibatch_size : int
        Minibatch size.
    clip_epsilon : float
        PPO clipping epsilon.
    value_clip : float | None
        Value-function clip range, or ``None``.
    value_weight : float
        Weight of the value loss.
    entropy_weight : float
        Weight of the entropy bonus.
    entropy_schedule : tuple[float, ...] | None
        Per-epoch entropy-weight schedule, or ``None``.
    max_grad_norm : float | None
        Gradient-norm clip threshold, or ``None``.
    kl_early_stop : float | None
        KL early-stopping threshold, or ``None``.
    metadata : dict[str, Any] | None
        Associated metadata mapping, or ``None``.

    Returns
    -------
    SupervisorPPOTrainResult
        The PPO training result.

    Raises
    ------
    ValueError
        If the inputs are invalid or inconsistent.
    """
    prior_losses = jnp.asarray([])
    prior_updates = 0
    checkpoint_path: Path | None = None
    if checkpoint_dir is not None:
        checkpoint_path = Path(checkpoint_dir)

    if resume:
        if checkpoint_path is None:
            raise ValueError("checkpoint_dir is required when resume=True")
        checkpoint = load_supervisor_ppo_checkpoint(
            checkpoint_path,
            template_policy=policy,
            template_opt_state=opt_state,
        )
        policy = checkpoint.policy
        opt_state = checkpoint.opt_state
        key = checkpoint.key
        prior_losses = checkpoint.loss_history
        prior_updates = checkpoint.n_updates

    policy, opt_state, new_losses, n_updates, next_key = (
        _ppo_supervisor_train_epochs_impl(
            policy,
            batch,
            key,
            opt_state,
            optimizer,
            n_epochs=n_epochs,
            minibatch_size=minibatch_size,
            clip_epsilon=clip_epsilon,
            value_clip=value_clip,
            value_weight=value_weight,
            entropy_weight=entropy_weight,
            entropy_schedule=entropy_schedule,
            max_grad_norm=max_grad_norm,
            kl_early_stop=kl_early_stop,
            initial_update_index=prior_updates,
        )
    )
    loss_history = (
        new_losses
        if prior_losses.size == 0
        else jnp.concatenate([prior_losses, new_losses])
    )
    total_updates = prior_updates + n_updates

    if checkpoint_path is not None:
        save_supervisor_ppo_checkpoint(
            checkpoint_path,
            policy=policy,
            opt_state=opt_state,
            key=next_key,
            n_updates=total_updates,
            loss_history=loss_history,
            metadata=metadata,
            overwrite=True,
        )

    return SupervisorPPOTrainResult(
        policy=policy,
        opt_state=opt_state,
        key=next_key,
        loss_history=loss_history,
        n_updates=total_updates,
        checkpoint_path=checkpoint_path,
    )

project_supervisor_action_for_audit 

project_supervisor_action_for_audit(
    action: SupervisorAction,
    config: DifferentiableSupervisorConfig,
    *,
    previous_action: SupervisorAction | None = None,
    ttl_s: float = 5.0,
    max_ttl_s: float = 5.0,
    rate_limit_fraction: float = 1.0,
    include_layer_actions: bool = True,
    regime_churn_score: float | None = None,
    max_regime_churn: float | None = None,
) -> SupervisorActionProjection

Project a neural proposal into replay-safe bounds with audit metadata.

This is intentionally non-actuating. It creates the explicit audit envelope that callers can inspect before converting a proposal into ControlAction objects for any live adapter path.

Parameters

action : SupervisorAction The supervisor control action. config : DifferentiableSupervisorConfig The supervisor configuration. previous_action : SupervisorAction | None The previous supervisor action, or None. ttl_s : float Action time-to-live in seconds. max_ttl_s : float Maximum action time-to-live in seconds. rate_limit_fraction : float Maximum fractional change per step. include_layer_actions : bool Whether to include per-layer actions. regime_churn_score : float | None The regime-churn score, or None. max_regime_churn : float | None Maximum allowed regime churn, or None.

Returns

SupervisorActionProjection The replay-safe action projection with audit metadata.

Source code in src/scpn_phase_orchestrator/nn/supervisor/projection.py 
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
def project_supervisor_action_for_audit(
    action: SupervisorAction,
    config: DifferentiableSupervisorConfig,
    *,
    previous_action: SupervisorAction | None = None,
    ttl_s: float = 5.0,
    max_ttl_s: float = 5.0,
    rate_limit_fraction: float = 1.0,
    include_layer_actions: bool = True,
    regime_churn_score: float | None = None,
    max_regime_churn: float | None = None,
) -> SupervisorActionProjection:
    """Project a neural proposal into replay-safe bounds with audit metadata.

    This is intentionally non-actuating. It creates the explicit audit envelope
    that callers can inspect before converting a proposal into ``ControlAction``
    objects for any live adapter path.

    Parameters
    ----------
    action : SupervisorAction
        The supervisor control action.
    config : DifferentiableSupervisorConfig
        The supervisor configuration.
    previous_action : SupervisorAction | None
        The previous supervisor action, or ``None``.
    ttl_s : float
        Action time-to-live in seconds.
    max_ttl_s : float
        Maximum action time-to-live in seconds.
    rate_limit_fraction : float
        Maximum fractional change per step.
    include_layer_actions : bool
        Whether to include per-layer actions.
    regime_churn_score : float | None
        The regime-churn score, or ``None``.
    max_regime_churn : float | None
        Maximum allowed regime churn, or ``None``.

    Returns
    -------
    SupervisorActionProjection
        The replay-safe action projection with audit metadata.
    """
    ttl_s = _positive_float(ttl_s, "ttl_s")
    max_ttl_s = _positive_float(max_ttl_s, "max_ttl_s")
    rate_limit_fraction = _bounded_unit_scalar(
        rate_limit_fraction,
        "rate_limit_fraction",
    )
    if regime_churn_score is not None:
        regime_churn_score = _non_negative_float(
            regime_churn_score,
            "regime_churn_score",
        )
    if max_regime_churn is not None:
        max_regime_churn = _positive_float(max_regime_churn, "max_regime_churn")
    projected_ttl = min(ttl_s, max_ttl_s)
    bounds = _action_bounds(config)
    proposed_values = pack_supervisor_action(action)
    bounded_values = jnp.clip(proposed_values, -bounds, bounds)
    rate_limited_values = bounded_values

    if previous_action is not None:
        previous_values = pack_supervisor_action(previous_action)
        max_delta = bounds * rate_limit_fraction
        lower = previous_values - max_delta
        upper = previous_values + max_delta
        rate_limited_values = jnp.clip(bounded_values, lower, upper)

    if not include_layer_actions:
        rate_limited_values = rate_limited_values.at[2:].set(0.0)

    rejection_reasons: list[str] = []
    if (
        regime_churn_score is not None
        and max_regime_churn is not None
        and regime_churn_score > max_regime_churn
    ):
        rejection_reasons.append("regime_churn")
        rate_limited_values = jnp.zeros_like(rate_limited_values)

    projected_action = unpack_supervisor_action(
        rate_limited_values,
        value_estimate=action.value_estimate,
        config=config,
    )
    controls = _supervisor_projection_control_records(
        proposed_values=proposed_values,
        projected_values=rate_limited_values,
        bounds=bounds,
    )
    clipped = projected_ttl != ttl_s or any(
        bool(control["clipped"]) for control in controls
    )
    audit_record = {
        "proposal_type": "differentiable_supervisor_action_projection",
        "non_actuating": True,
        "rejected": bool(rejection_reasons),
        "rejection_reasons": rejection_reasons,
        "clipped": clipped,
        "ttl_s": projected_ttl,
        "requested_ttl_s": ttl_s,
        "constraints": {
            "max_ttl_s": max_ttl_s,
            "rate_limit_fraction": rate_limit_fraction,
            "include_layer_actions": include_layer_actions,
            "previous_action": previous_action is not None,
            "regime_churn_score": regime_churn_score,
            "max_regime_churn": max_regime_churn,
        },
        "controls": controls,
    }
    return SupervisorActionProjection(
        action=projected_action,
        ttl_s=projected_ttl,
        audit_record=audit_record,
    )

build_supervisor_baseline_report 

build_supervisor_baseline_report(
    comparisons: Iterable[Any],
    *,
    report_label: str = "supervisor_baseline_report",
) -> SupervisorBaselineReport

Aggregate already-generated supervisor comparison records for review.

Parameters

comparisons : Iterable[Any] The supervisor comparison records. report_label : str Label for the baseline report.

Returns

SupervisorBaselineReport The aggregated baseline report.

Raises

ValueError If the inputs are invalid or inconsistent.

Source code in src/scpn_phase_orchestrator/nn/supervisor/replay.py 
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
def build_supervisor_baseline_report(
    comparisons: Iterable[Any],
    *,
    report_label: str = "supervisor_baseline_report",
) -> SupervisorBaselineReport:
    """Aggregate already-generated supervisor comparison records for review.

    Parameters
    ----------
    comparisons : Iterable[Any]
        The supervisor comparison records.
    report_label : str
        Label for the baseline report.

    Returns
    -------
    SupervisorBaselineReport
        The aggregated baseline report.

    Raises
    ------
    ValueError
        If the inputs are invalid or inconsistent.
    """
    if not report_label:
        raise ValueError("report_label must not be empty")
    records = tuple(
        _comparison_record_from_object(comparison, f"comparison {index}")
        for index, comparison in enumerate(comparisons)
    )
    if not records:
        raise ValueError("baseline report requires at least one comparison")
    return SupervisorBaselineReport(
        comparisons=records,
        summary=_baseline_report_summary(records),
        report_label=report_label,
        actuation_permitted=False,
    )

build_supervisor_corpus_replay_proposals 

build_supervisor_corpus_replay_proposals(
    policy: DifferentiableSupervisorPolicy,
    corpus: SupervisorScenarioCorpus,
    *,
    previous_action: SupervisorAction | None = None,
    ttl_s: float = 5.0,
    max_ttl_s: float = 5.0,
    rate_limit_fraction: float = 1.0,
    include_layer_actions: bool = True,
    regime_churn_scores: tuple[float, ...] | None = None,
    max_regime_churn: float | None = None,
) -> SupervisorCorpusReplayProposals

Build deterministic replay-only proposals for every corpus scenario.

Parameters

policy : DifferentiableSupervisorPolicy The differentiable supervisor policy. corpus : SupervisorScenarioCorpus The validated supervisor scenario corpus. previous_action : SupervisorAction | None The previous supervisor action, or None. ttl_s : float Action time-to-live in seconds. max_ttl_s : float Maximum action time-to-live in seconds. rate_limit_fraction : float Maximum fractional change per step. include_layer_actions : bool Whether to include per-layer actions. regime_churn_scores : tuple[float, ...] | None Per-scenario regime-churn scores, or None. max_regime_churn : float | None Maximum allowed regime churn, or None.

Returns

SupervisorCorpusReplayProposals The replay-only proposals for every corpus scenario.

Raises

ValueError If the inputs are invalid or inconsistent.

Source code in src/scpn_phase_orchestrator/nn/supervisor/replay.py 
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
def build_supervisor_corpus_replay_proposals(
    policy: DifferentiableSupervisorPolicy,
    corpus: SupervisorScenarioCorpus,
    *,
    previous_action: SupervisorAction | None = None,
    ttl_s: float = 5.0,
    max_ttl_s: float = 5.0,
    rate_limit_fraction: float = 1.0,
    include_layer_actions: bool = True,
    regime_churn_scores: tuple[float, ...] | None = None,
    max_regime_churn: float | None = None,
) -> SupervisorCorpusReplayProposals:
    """Build deterministic replay-only proposals for every corpus scenario.

    Parameters
    ----------
    policy : DifferentiableSupervisorPolicy
        The differentiable supervisor policy.
    corpus : SupervisorScenarioCorpus
        The validated supervisor scenario corpus.
    previous_action : SupervisorAction | None
        The previous supervisor action, or ``None``.
    ttl_s : float
        Action time-to-live in seconds.
    max_ttl_s : float
        Maximum action time-to-live in seconds.
    rate_limit_fraction : float
        Maximum fractional change per step.
    include_layer_actions : bool
        Whether to include per-layer actions.
    regime_churn_scores : tuple[float, ...] | None
        Per-scenario regime-churn scores, or ``None``.
    max_regime_churn : float | None
        Maximum allowed regime churn, or ``None``.

    Returns
    -------
    SupervisorCorpusReplayProposals
        The replay-only proposals for every corpus scenario.

    Raises
    ------
    ValueError
        If the inputs are invalid or inconsistent.
    """
    if not corpus.scenarios:
        raise ValueError("scenario corpus requires at least one scenario")
    if len(corpus.metadata) != len(corpus.scenarios):
        raise ValueError("scenario corpus metadata must match scenario count")
    if regime_churn_scores is not None and len(regime_churn_scores) != len(
        corpus.scenarios
    ):
        raise ValueError("regime_churn_scores must match scenario count")

    proposals = []
    for index, (scenario, metadata) in enumerate(
        zip(corpus.scenarios, corpus.metadata, strict=True)
    ):
        regime_churn_score = (
            None if regime_churn_scores is None else regime_churn_scores[index]
        )
        proposal_metadata = _json_object(
            {**metadata, "corpus_index": index},
            f"corpus metadata {index}",
        )
        proposals.append(
            build_supervisor_replay_proposal(
                policy,
                scenario,
                scenario_metadata=proposal_metadata,
                previous_action=previous_action,
                ttl_s=ttl_s,
                max_ttl_s=max_ttl_s,
                rate_limit_fraction=rate_limit_fraction,
                include_layer_actions=include_layer_actions,
                regime_churn_score=regime_churn_score,
                max_regime_churn=max_regime_churn,
            )
        )
    return SupervisorCorpusReplayProposals(
        proposals=tuple(proposals),
        actuation_permitted=False,
    )

build_supervisor_experiment_manifest 

build_supervisor_experiment_manifest(
    baseline_report: SupervisorBaselineReport,
    *,
    command: str,
    git_sha: str,
    dependency_lock: Mapping[str, Any],
    device_info: Mapping[str, Any],
    seed_list: Iterable[int],
    config_json_path: str | None = None,
    metrics_jsonl_path: str | None = None,
    summary_table_path: str | None = None,
    checkpoint_manifest_path: str | None = None,
    plot_manifest_path: str | None = None,
) -> SupervisorExperimentManifest

Build a reproducibility manifest for a supervisor baseline report.

Parameters

baseline_report : SupervisorBaselineReport The aggregated baseline report. command : str The command line recorded with the run. git_sha : str Git commit SHA recorded with the run. dependency_lock : Mapping[str, Any] Dependency lock mapping recorded with the run. device_info : Mapping[str, Any] Device information recorded with the run. seed_list : Iterable[int] Seeds for the experiment runs. config_json_path : str | None Filesystem path to the config json, or None. metrics_jsonl_path : str | None Filesystem path to the metrics jsonl, or None. summary_table_path : str | None Filesystem path to the summary table, or None. checkpoint_manifest_path : str | None Filesystem path to the checkpoint manifest, or None. plot_manifest_path : str | None Filesystem path to the plot manifest, or None.

Returns

SupervisorExperimentManifest The reproducibility manifest.

Raises

ValueError If the inputs are invalid or inconsistent.

Source code in src/scpn_phase_orchestrator/nn/supervisor/replay.py 
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
def build_supervisor_experiment_manifest(
    baseline_report: SupervisorBaselineReport,
    *,
    command: str,
    git_sha: str,
    dependency_lock: Mapping[str, Any],
    device_info: Mapping[str, Any],
    seed_list: Iterable[int],
    config_json_path: str | None = None,
    metrics_jsonl_path: str | None = None,
    summary_table_path: str | None = None,
    checkpoint_manifest_path: str | None = None,
    plot_manifest_path: str | None = None,
) -> SupervisorExperimentManifest:
    """Build a reproducibility manifest for a supervisor baseline report.

    Parameters
    ----------
    baseline_report : SupervisorBaselineReport
        The aggregated baseline report.
    command : str
        The command line recorded with the run.
    git_sha : str
        Git commit SHA recorded with the run.
    dependency_lock : Mapping[str, Any]
        Dependency lock mapping recorded with the run.
    device_info : Mapping[str, Any]
        Device information recorded with the run.
    seed_list : Iterable[int]
        Seeds for the experiment runs.
    config_json_path : str | None
        Filesystem path to the config json, or ``None``.
    metrics_jsonl_path : str | None
        Filesystem path to the metrics jsonl, or ``None``.
    summary_table_path : str | None
        Filesystem path to the summary table, or ``None``.
    checkpoint_manifest_path : str | None
        Filesystem path to the checkpoint manifest, or ``None``.
    plot_manifest_path : str | None
        Filesystem path to the plot manifest, or ``None``.

    Returns
    -------
    SupervisorExperimentManifest
        The reproducibility manifest.

    Raises
    ------
    ValueError
        If the inputs are invalid or inconsistent.
    """
    if not command:
        raise ValueError("command must not be empty")
    if not git_sha:
        raise ValueError("git_sha must not be empty")
    seeds = tuple(
        _non_negative_int(seed, f"seed_list[{index}]")
        for index, seed in enumerate(seed_list)
    )
    if not seeds:
        raise ValueError("seed_list must contain at least one seed")
    lock_record = _json_object(dict(dependency_lock), "dependency_lock")
    if not lock_record:
        raise ValueError("dependency_lock must not be empty")
    device_record = _json_object(dict(device_info), "device_info")
    if not device_record:
        raise ValueError("device_info must not be empty")
    report_record = _json_object(
        baseline_report.to_audit_record(),
        "baseline_report",
    )
    if report_record.get("actuation_permitted") is not False:
        raise ValueError("baseline_report must be non-actuating")
    artifacts = _json_object(
        {
            "config_json_path": config_json_path,
            "metrics_jsonl_path": metrics_jsonl_path,
            "summary_table_path": summary_table_path,
            "checkpoint_manifest_path": checkpoint_manifest_path,
            "plot_manifest_path": plot_manifest_path,
        },
        "artifacts",
    )
    return SupervisorExperimentManifest(
        baseline_report=report_record,
        command=command,
        git_sha=git_sha,
        dependency_lock=lock_record,
        device_info=device_record,
        seed_list=seeds,
        artifacts=artifacts,
        actuation_permitted=False,
    )

build_supervisor_replay_proposal 

build_supervisor_replay_proposal(
    policy: DifferentiableSupervisorPolicy,
    scenario: KuramotoSupervisorScenario,
    *,
    scenario_metadata: dict[str, Any] | None = None,
    previous_action: SupervisorAction | None = None,
    ttl_s: float = 5.0,
    max_ttl_s: float = 5.0,
    rate_limit_fraction: float = 1.0,
    include_layer_actions: bool = True,
    regime_churn_score: float | None = None,
    max_regime_churn: float | None = None,
) -> SupervisorReplayProposal

Build a deterministic replay-only proposal from a neural supervisor.

The proposal is an audit artefact only. It does not return ControlAction objects and carries actuation_permitted=False so that downstream replay/autotune surfaces can review it without enabling live actuation.

Parameters

policy : DifferentiableSupervisorPolicy The differentiable supervisor policy. scenario : KuramotoSupervisorScenario The Kuramoto supervisor scenario. scenario_metadata : dict[str, Any] | None Scenario metadata, or None. previous_action : SupervisorAction | None The previous supervisor action, or None. ttl_s : float Action time-to-live in seconds. max_ttl_s : float Maximum action time-to-live in seconds. rate_limit_fraction : float Maximum fractional change per step. include_layer_actions : bool Whether to include per-layer actions. regime_churn_score : float | None The regime-churn score, or None. max_regime_churn : float | None Maximum allowed regime churn, or None.

Returns

SupervisorReplayProposal The replay-only supervisor proposal.

Source code in src/scpn_phase_orchestrator/nn/supervisor/replay.py 
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
def build_supervisor_replay_proposal(
    policy: DifferentiableSupervisorPolicy,
    scenario: KuramotoSupervisorScenario,
    *,
    scenario_metadata: dict[str, Any] | None = None,
    previous_action: SupervisorAction | None = None,
    ttl_s: float = 5.0,
    max_ttl_s: float = 5.0,
    rate_limit_fraction: float = 1.0,
    include_layer_actions: bool = True,
    regime_churn_score: float | None = None,
    max_regime_churn: float | None = None,
) -> SupervisorReplayProposal:
    """Build a deterministic replay-only proposal from a neural supervisor.

    The proposal is an audit artefact only. It does not return
    ``ControlAction`` objects and carries ``actuation_permitted=False`` so that
    downstream replay/autotune surfaces can review it without enabling live
    actuation.

    Parameters
    ----------
    policy : DifferentiableSupervisorPolicy
        The differentiable supervisor policy.
    scenario : KuramotoSupervisorScenario
        The Kuramoto supervisor scenario.
    scenario_metadata : dict[str, Any] | None
        Scenario metadata, or ``None``.
    previous_action : SupervisorAction | None
        The previous supervisor action, or ``None``.
    ttl_s : float
        Action time-to-live in seconds.
    max_ttl_s : float
        Maximum action time-to-live in seconds.
    rate_limit_fraction : float
        Maximum fractional change per step.
    include_layer_actions : bool
        Whether to include per-layer actions.
    regime_churn_score : float | None
        The regime-churn score, or ``None``.
    max_regime_churn : float | None
        Maximum allowed regime churn, or ``None``.

    Returns
    -------
    SupervisorReplayProposal
        The replay-only supervisor proposal.
    """
    metadata = _json_object(scenario_metadata, "scenario_metadata")
    action = policy(scenario)
    projection = project_supervisor_action_for_audit(
        action,
        policy.config,
        previous_action=previous_action,
        ttl_s=ttl_s,
        max_ttl_s=max_ttl_s,
        rate_limit_fraction=rate_limit_fraction,
        include_layer_actions=include_layer_actions,
        regime_churn_score=regime_churn_score,
        max_regime_churn=max_regime_churn,
    )
    metrics = {
        "current_R_global": float(order_parameter(scenario.phases)),
        "current_R_good": float(
            masked_order_parameter(scenario.phases, scenario.good_mask)
        ),
        "current_R_bad": float(
            masked_order_parameter(scenario.phases, scenario.bad_mask)
        ),
        "value_estimate": float(action.value_estimate),
        "bound_penalty": float(supervisor_action_bound_penalty(action, policy.config)),
    }
    scenario_summary = {
        "n_oscillators": int(scenario.phases.shape[0]),
        "dt": float(scenario.dt),
        "inner_steps": int(scenario.inner_steps),
        "horizon": int(scenario.horizon),
    }
    return SupervisorReplayProposal(
        action=action,
        projection=projection,
        scenario_summary=scenario_summary,
        scenario_metadata=metadata,
        metrics=metrics,
        actuation_permitted=False,
    )

build_supervisor_scenario_corpus 

build_supervisor_scenario_corpus(
    records: Iterable[Mapping[str, Any]],
    *,
    dtype: Any = jnp.float32,
) -> SupervisorScenarioCorpus

Convert replay/audit records into validated supervisor scenarios.

Parameters

records : Iterable[Mapping[str, Any]] Replay/audit records to convert. dtype : Any Target array dtype.

Returns

SupervisorScenarioCorpus The validated supervisor scenario corpus.

Raises

ValueError If the inputs are invalid or inconsistent.

Source code in src/scpn_phase_orchestrator/nn/supervisor/rollouts.py 
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
def build_supervisor_scenario_corpus(
    records: Iterable[Mapping[str, Any]],
    *,
    dtype: Any = jnp.float32,
) -> SupervisorScenarioCorpus:
    """Convert replay/audit records into validated supervisor scenarios.

    Parameters
    ----------
    records : Iterable[Mapping[str, Any]]
        Replay/audit records to convert.
    dtype : Any
        Target array dtype.

    Returns
    -------
    SupervisorScenarioCorpus
        The validated supervisor scenario corpus.

    Raises
    ------
    ValueError
        If the inputs are invalid or inconsistent.
    """
    scenarios: list[KuramotoSupervisorScenario] = []
    metadata: list[dict[str, Any]] = []
    for index, record in enumerate(records):
        if not isinstance(record, Mapping):
            raise ValueError(f"record {index} must be a mapping")
        scenario, record_metadata = _supervisor_scenario_from_record(
            record,
            index=index,
            dtype=dtype,
        )
        scenarios.append(scenario)
        metadata.append(record_metadata)
    if not scenarios:
        raise ValueError("scenario corpus requires at least one record")
    return SupervisorScenarioCorpus(
        scenarios=tuple(scenarios),
        metadata=tuple(metadata),
    )

collect_supervisor_corpus_rollouts 

collect_supervisor_corpus_rollouts(
    policy: DifferentiableSupervisorPolicy,
    corpus: SupervisorScenarioCorpus,
    *,
    key: Array,
    n_episodes_per_scenario: int,
    gamma: float = 0.99,
    gae_lambda: float = 0.95,
    trajectory_jitter: float = 0.0,
) -> SupervisorPPOCorpusRollout

Collect replay-only PPO rollouts across a validated scenario corpus.

The returned batch is a flat concatenation suitable for PPO epochs. All corpus scenarios must share dt, inner_steps, horizon, and oscillator tensor shapes because SupervisorPPOBatch stores timing fields once for the full batch.

Parameters

policy : DifferentiableSupervisorPolicy The differentiable supervisor policy. corpus : SupervisorScenarioCorpus The validated supervisor scenario corpus. key : jax.Array JAX PRNG key. n_episodes_per_scenario : int Number of episodes per corpus scenario. gamma : float Flow-dependent elimination rate. gae_lambda : float Generalised-advantage-estimation lambda. trajectory_jitter : float Trajectory jitter magnitude.

Returns

SupervisorPPOCorpusRollout The collected corpus PPO rollouts.

Raises

ValueError If the inputs are invalid or inconsistent.

Source code in src/scpn_phase_orchestrator/nn/supervisor/rollouts.py 
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
def collect_supervisor_corpus_rollouts(
    policy: DifferentiableSupervisorPolicy,
    corpus: SupervisorScenarioCorpus,
    *,
    key: jax.Array,
    n_episodes_per_scenario: int,
    gamma: float = 0.99,
    gae_lambda: float = 0.95,
    trajectory_jitter: float = 0.0,
) -> SupervisorPPOCorpusRollout:
    """Collect replay-only PPO rollouts across a validated scenario corpus.

    The returned batch is a flat concatenation suitable for PPO epochs. All
    corpus scenarios must share ``dt``, ``inner_steps``, ``horizon``, and
    oscillator tensor shapes because ``SupervisorPPOBatch`` stores timing
    fields once for the full batch.

    Parameters
    ----------
    policy : DifferentiableSupervisorPolicy
        The differentiable supervisor policy.
    corpus : SupervisorScenarioCorpus
        The validated supervisor scenario corpus.
    key : jax.Array
        JAX PRNG key.
    n_episodes_per_scenario : int
        Number of episodes per corpus scenario.
    gamma : float
        Flow-dependent elimination rate.
    gae_lambda : float
        Generalised-advantage-estimation lambda.
    trajectory_jitter : float
        Trajectory jitter magnitude.

    Returns
    -------
    SupervisorPPOCorpusRollout
        The collected corpus PPO rollouts.

    Raises
    ------
    ValueError
        If the inputs are invalid or inconsistent.
    """
    n_episodes_per_scenario = _positive_int(
        n_episodes_per_scenario,
        "n_episodes_per_scenario",
    )
    if not corpus.scenarios:
        raise ValueError("scenario corpus requires at least one scenario")
    if len(corpus.metadata) != len(corpus.scenarios):
        raise ValueError("scenario corpus metadata must match scenario count")
    _validate_supervisor_corpus_rollout_compatibility(corpus.scenarios)

    scenario_keys = jax.random.split(key, len(corpus.scenarios))
    scenario_rollouts: list[SupervisorPPORollout] = []
    scenario_indices: list[jax.Array] = []
    for scenario_index, (scenario, scenario_key) in enumerate(
        zip(corpus.scenarios, scenario_keys, strict=True)
    ):
        rollout = collect_supervisor_rollouts(
            policy,
            scenario,
            key=scenario_key,
            n_episodes=n_episodes_per_scenario,
            gamma=gamma,
            gae_lambda=gae_lambda,
            trajectory_jitter=trajectory_jitter,
        )
        scenario_rollouts.append(rollout)
        scenario_indices.append(
            jnp.full(
                (n_episodes_per_scenario,),
                scenario_index,
                dtype=jnp.int32,
            )
        )

    batch = _concatenate_supervisor_ppo_batches(
        tuple(rollout.batch for rollout in scenario_rollouts)
    )
    episode_returns = jnp.concatenate(
        [rollout.episode_returns for rollout in scenario_rollouts],
        axis=0,
    )
    return SupervisorPPOCorpusRollout(
        batch=batch,
        episode_returns=episode_returns,
        episode_return_mean=jnp.mean(episode_returns),
        episode_return_std=jnp.std(episode_returns),
        scenario_indices=jnp.concatenate(scenario_indices, axis=0),
        metadata=tuple(dict(item) for item in corpus.metadata),
    )

collect_supervisor_rollouts 

collect_supervisor_rollouts(
    policy: DifferentiableSupervisorPolicy,
    scenario: KuramotoSupervisorScenario,
    *,
    key: Array,
    n_episodes: int,
    gamma: float = 0.99,
    gae_lambda: float = 0.95,
    trajectory_jitter: float = 0.0,
) -> SupervisorPPORollout

Collect deterministic, replay-only PPO rollouts from a starting scenario.

Parameters

policy : DifferentiableSupervisorPolicy The differentiable supervisor policy. scenario : KuramotoSupervisorScenario The Kuramoto supervisor scenario. key : jax.Array JAX PRNG key. n_episodes : int Number of rollout episodes. gamma : float Flow-dependent elimination rate. gae_lambda : float Generalised-advantage-estimation lambda. trajectory_jitter : float Trajectory jitter magnitude.

Returns

SupervisorPPORollout The collected PPO rollouts.

Raises

ValueError If the inputs are invalid or inconsistent.

Source code in src/scpn_phase_orchestrator/nn/supervisor/rollouts.py 
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
def collect_supervisor_rollouts(
    policy: DifferentiableSupervisorPolicy,
    scenario: KuramotoSupervisorScenario,
    *,
    key: jax.Array,
    n_episodes: int,
    gamma: float = 0.99,
    gae_lambda: float = 0.95,
    trajectory_jitter: float = 0.0,
) -> SupervisorPPORollout:
    """Collect deterministic, replay-only PPO rollouts from a starting scenario.

    Parameters
    ----------
    policy : DifferentiableSupervisorPolicy
        The differentiable supervisor policy.
    scenario : KuramotoSupervisorScenario
        The Kuramoto supervisor scenario.
    key : jax.Array
        JAX PRNG key.
    n_episodes : int
        Number of rollout episodes.
    gamma : float
        Flow-dependent elimination rate.
    gae_lambda : float
        Generalised-advantage-estimation lambda.
    trajectory_jitter : float
        Trajectory jitter magnitude.

    Returns
    -------
    SupervisorPPORollout
        The collected PPO rollouts.

    Raises
    ------
    ValueError
        If the inputs are invalid or inconsistent.
    """
    n_episodes = _positive_int(n_episodes, "n_episodes")
    horizon = _positive_int(scenario.horizon, "scenario.horizon")
    inner_steps = _positive_int(scenario.inner_steps, "scenario.inner_steps")
    gamma = _bounded_unit_scalar(gamma, "gamma")
    gae_lambda = _bounded_unit_scalar(gae_lambda, "gae_lambda")
    if isinstance(trajectory_jitter, bool) or not isinstance(
        trajectory_jitter, int | float
    ):
        raise ValueError("trajectory_jitter must be a finite float")
    trajectory_jitter = float(trajectory_jitter)
    if not isfinite(trajectory_jitter) or trajectory_jitter < 0.0:
        raise ValueError("trajectory_jitter must be non-negative")

    rollout_keys = jax.random.split(key, n_episodes + 1)[1:]
    all_phases: list[jax.Array] = []
    all_omegas: list[jax.Array] = []
    all_base_k: list[jax.Array] = []
    all_good_masks: list[jax.Array] = []
    all_bad_masks: list[jax.Array] = []
    all_actions: list[jax.Array] = []
    all_old_log_probs: list[jax.Array] = []
    all_advantages: list[jax.Array] = []
    all_returns: list[jax.Array] = []
    all_values: list[jax.Array] = []
    episode_returns: list[jax.Array] = []

    for episode_key in rollout_keys:
        state = scenario
        if trajectory_jitter > 0.0:
            jitter_key, step_key = jax.random.split(episode_key)
            state = state._replace(
                phases=state.phases
                + trajectory_jitter
                * jax.random.normal(jitter_key, shape=state.phases.shape),
            )
            step_key = jax.random.fold_in(step_key, 1)
        else:
            step_key = episode_key

        step_rewards: list[jax.Array] = []
        step_values: list[jax.Array] = []
        step_actions: list[jax.Array] = []
        step_log_probs: list[jax.Array] = []
        step_phases: list[jax.Array] = []

        for _ in range(horizon):
            current_key, step_key = jax.random.split(step_key)
            action, action_log_prob = sample_supervisor_action(
                policy, state, key=current_key
            )
            controlled_K = apply_supervisor_action(scenario.base_K, action, state)
            next_phases, _ = kuramoto_forward(
                state.phases,
                state.omegas,
                controlled_K,
                state.dt,
                inner_steps,
            )
            reward = masked_order_parameter(next_phases, state.good_mask) - (
                policy.config.bad_sync_weight
                * masked_order_parameter(next_phases, state.bad_mask)
            )

            step_phases.append(state.phases)
            step_rewards.append(reward)
            step_values.append(action.value_estimate)
            step_actions.append(pack_supervisor_action(action))
            step_log_probs.append(action_log_prob)
            state = state._replace(phases=next_phases)

        terminal_value = _policy_mean_and_value(policy, state)[1]
        trajectory_values = jnp.stack(step_values + [terminal_value])
        rewards = jnp.stack(step_rewards)
        deltas = rewards + gamma * trajectory_values[1:] - trajectory_values[:-1]

        episode_advantages = [jnp.array(0.0)] * horizon
        running_advantage = jnp.array(0.0)
        for index in reversed(range(horizon)):
            running_advantage = deltas[index] + gamma * gae_lambda * running_advantage
            episode_advantages[index] = running_advantage
        advantages = jnp.stack(episode_advantages)
        returns = advantages + trajectory_values[:-1]

        episode_returns.append(jnp.sum(rewards))
        for index in range(horizon):
            all_phases.append(step_phases[index])
            all_omegas.append(scenario.omegas)
            all_base_k.append(scenario.base_K)
            all_good_masks.append(scenario.good_mask)
            all_bad_masks.append(scenario.bad_mask)
            all_actions.append(step_actions[index])
            all_old_log_probs.append(step_log_probs[index])
            all_advantages.append(advantages[index])
            all_returns.append(returns[index])
            all_values.append(step_values[index])

    episode_returns_array = jnp.stack(episode_returns)
    batch = SupervisorPPOBatch(
        phases=jnp.stack(all_phases),
        omegas=jnp.stack(all_omegas),
        base_K=jnp.stack(all_base_k),
        good_mask=jnp.stack(all_good_masks),
        bad_mask=jnp.stack(all_bad_masks),
        actions=jnp.stack(all_actions),
        old_log_probs=jnp.stack(all_old_log_probs),
        advantages=jnp.stack(all_advantages),
        returns=jnp.stack(all_returns),
        values=jnp.stack(all_values),
        dt=scenario.dt,
        inner_steps=inner_steps,
        horizon=horizon,
    )
    return SupervisorPPORollout(
        batch=batch,
        episode_returns=episode_returns_array,
        episode_return_mean=jnp.mean(episode_returns_array),
        episode_return_std=jnp.std(episode_returns_array),
    )

Physics validation

The nn module includes 13 physics validation test files (test_nn_physics_validation_p1 through _p13) verifying:

  • Energy conservation under Hamiltonian coupling
  • Gradient correctness via finite-difference comparison
  • Order parameter convergence for strong coupling
  • Stuart-Landau bifurcation (subcritical → supercritical)
  • Simplicial explosive synchronisation
  • BOLD hemodynamic response shape
  • Reservoir echo state property
  • UDE residual convergence
  • OIM graph colouring correctness
  • Inverse coupling recovery (r > 0.95)