spinoct 0.10.0__tar.gz

spinoct-0.10.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Felipe Santibanez-Leal
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

spinoct-0.10.0/PKG-INFO

@@ -0,0 +1,68 @@
+Metadata-Version: 2.4
+Name: spinoct
+Version: 0.10.0
+Summary: Optimal control paths and energy-efficient switching pulses for classical spin dynamics (Landau-Lifshitz-Gilbert)
+Author-email: Felipe Santibanez-Leal <fsantibanez@gmail.com>
+License: MIT
+Project-URL: Repository, https://github.com/fsantibanezleal/CAOS_SpinOCT
+Keywords: optimal-control,magnetization-switching,landau-lifshitz-gilbert,spintronics,spin-dynamics,micromagnetics,pulse-shaping,energy-efficient-memory
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Physics
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.24
+Requires-Dist: scipy>=1.10
+Provides-Extra: torch
+Requires-Dist: torch>=2.2; extra == "torch"
+Provides-Extra: dev
+Requires-Dist: pytest>=8; extra == "dev"
+Requires-Dist: ruff>=0.4; extra == "dev"
+Dynamic: license-file
+
+# spinoct
+
+Optimal control paths and energy-efficient switching pulses for classical spin dynamics
+(Landau-Lifshitz-Gilbert).
+
+`spinoct` computes the control (an applied magnetic field, an electric current, or both) that
+drives a magnetic moment from one state to another in a given time for the least dissipated energy.
+It is the reusable engine behind [Espira](https://github.com/fsantibanezleal/CAOS_RES_Espira), and
+it is deliberately independent of any material database, so it works on any spin Hamiltonian.
+
+## Why this exists
+
+The optimal control of magnetization switching has a small, rigorous literature (Kwiatkowski,
+Badarneh, Berkov and Bessarab, Phys. Rev. Lett. 126, 177206 (2021), and the papers that follow it),
+but no open implementation. The analytic solutions live only as equations in journal articles, and
+the group's own code is not public. `spinoct` is that implementation, with the closed-form results
+built in as positive controls that every numerical solver is validated against before it is trusted.
+
+## Install
+
+```bash
+pip install spinoct              # core: numpy + scipy
+pip install "spinoct[torch]"     # add the batched GPU solvers
+```
+
+## The dimensional contract
+
+Read `spinoct.units` first. The same physical quantity is written four different ways across this
+literature, and the central quantity of the package, the switching cost `Phi = int |b|^2 dt`, is in
+tesla-squared-seconds, not joules. It becomes an energy only through an explicit circuit model. The
+package refuses to hide that assumption: `spinoct.units.CircuitModel` is a required, described
+object, never a buried constant.
+
+## Status
+
+Pre-1.0, under active development. The analytic uniaxial optimal control path, its closed-form pulse
+and cost, and the negative-parameter Jacobi elliptic machinery it needs are complete and validated.
+The numerical image-based solver, the GRAPE and CRAB constrained solvers, and the batched GPU lane
+are in progress.
+
+## License
+
+MIT. See `LICENSE`.

spinoct-0.10.0/README.md

@@ -0,0 +1,43 @@
+# spinoct
+
+Optimal control paths and energy-efficient switching pulses for classical spin dynamics
+(Landau-Lifshitz-Gilbert).
+
+`spinoct` computes the control (an applied magnetic field, an electric current, or both) that
+drives a magnetic moment from one state to another in a given time for the least dissipated energy.
+It is the reusable engine behind [Espira](https://github.com/fsantibanezleal/CAOS_RES_Espira), and
+it is deliberately independent of any material database, so it works on any spin Hamiltonian.
+
+## Why this exists
+
+The optimal control of magnetization switching has a small, rigorous literature (Kwiatkowski,
+Badarneh, Berkov and Bessarab, Phys. Rev. Lett. 126, 177206 (2021), and the papers that follow it),
+but no open implementation. The analytic solutions live only as equations in journal articles, and
+the group's own code is not public. `spinoct` is that implementation, with the closed-form results
+built in as positive controls that every numerical solver is validated against before it is trusted.
+
+## Install
+
+```bash
+pip install spinoct              # core: numpy + scipy
+pip install "spinoct[torch]"     # add the batched GPU solvers
+```
+
+## The dimensional contract
+
+Read `spinoct.units` first. The same physical quantity is written four different ways across this
+literature, and the central quantity of the package, the switching cost `Phi = int |b|^2 dt`, is in
+tesla-squared-seconds, not joules. It becomes an energy only through an explicit circuit model. The
+package refuses to hide that assumption: `spinoct.units.CircuitModel` is a required, described
+object, never a buried constant.
+
+## Status
+
+Pre-1.0, under active development. The analytic uniaxial optimal control path, its closed-form pulse
+and cost, and the negative-parameter Jacobi elliptic machinery it needs are complete and validated.
+The numerical image-based solver, the GRAPE and CRAB constrained solvers, and the batched GPU lane
+are in progress.
+
+## License
+
+MIT. See `LICENSE`.

spinoct-0.10.0/pyproject.toml

@@ -0,0 +1,51 @@
+[build-system]
+requires = ["setuptools>=68"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "spinoct"
+version = "0.10.0"
+description = "Optimal control paths and energy-efficient switching pulses for classical spin dynamics (Landau-Lifshitz-Gilbert)"
+readme = "README.md"
+license = { text = "MIT" }
+authors = [{ name = "Felipe Santibanez-Leal", email = "fsantibanez@gmail.com" }]
+requires-python = ">=3.10"
+dependencies = ["numpy>=1.24", "scipy>=1.10"]
+classifiers = [
+  "Development Status :: 3 - Alpha",
+  "Intended Audience :: Science/Research",
+  "License :: OSI Approved :: MIT License",
+  "Programming Language :: Python :: 3",
+  "Topic :: Scientific/Engineering :: Physics",
+]
+keywords = [
+  "optimal-control",
+  "magnetization-switching",
+  "landau-lifshitz-gilbert",
+  "spintronics",
+  "spin-dynamics",
+  "micromagnetics",
+  "pulse-shaping",
+  "energy-efficient-memory",
+]
+
+[project.urls]
+Repository = "https://github.com/fsantibanezleal/CAOS_SpinOCT"
+
+[project.optional-dependencies]
+torch = ["torch>=2.2"]
+dev = ["pytest>=8", "ruff>=0.4"]
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = "-q"
+
+[tool.ruff]
+line-length = 110
+target-version = "py310"
+
+[tool.ruff.lint]
+select = ["E", "F", "W", "I", "UP", "B"]

spinoct-0.10.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

spinoct-0.10.0/src/spinoct/__init__.py

@@ -0,0 +1,24 @@
+"""spinoct: optimal control paths and energy-efficient switching pulses for classical spin dynamics.
+
+The package solves one problem: given a magnetic system, an initial and a final state, and a
+switching time, find the control (an applied magnetic field, an electric current, or both) that
+drives the transition for the least dissipated energy.
+
+It is the engine behind Espira (https://github.com/fsantibanezleal/CAOS_RES_Espira) and is
+deliberately independent of any material database, so it can be used on any spin Hamiltonian.
+
+Start here
+----------
+- :mod:`spinoct.units` is the dimensional contract. Read it before anything else; the same symbol
+  means four different things across this literature.
+- :mod:`spinoct.dynamics` carries the system definition and the Landau-Lifshitz-Gilbert equation.
+- :mod:`spinoct.analytic` carries the closed-form optimal control paths, which are the positive
+  controls every numerical result is checked against.
+"""
+
+from __future__ import annotations
+
+__version__ = "0.10.0"
+__display_version__ = "0.10.000"
+
+__all__ = ["__display_version__", "__version__"]

spinoct-0.10.0/src/spinoct/adjoint.py

@@ -0,0 +1,251 @@
+"""Exact gradient-based pulse optimization by the discrete adjoint method (R10).
+
+The finite-difference gradient of the constrained solvers costs one forward integration per control
+parameter. The adjoint method computes the exact gradient with respect to every control parameter in a
+single backward pass, at the cost of one extra integration, independent of the number of parameters.
+For a pulse with many slices this is the difference between a solver that scales and one that does not,
+and it is the classical route to gradient-based optimal control.
+
+We differentiate through a norm-projected forward Euler integration of the Landau-Lifshitz-Gilbert
+equation analytically (a hand-derived reverse-mode pass), so the gradient is exact for the discretized
+system and needs no automatic-differentiation dependency; the core stays pure numpy. The same reverse
+pass runs elementwise, so it batches and ports to a GPU tensor library unchanged, which is the
+``[torch]`` extra.
+
+Forward step, with control field ``b_k`` piecewise constant on ``N`` slices:
+
+    r_k = s_k + dt * f(s_k, b_k),    s_{k+1} = r_k / |r_k|,
+
+with ``f`` the LLG right-hand side. The objective is the switching cost plus a reversal-fidelity
+penalty:
+
+    C = sum_k |b_k|^2 dt  +  lambda (1 + s_{N,z}) / 2.
+
+The adjoint recursion propagates the cost sensitivity ``g_k = dC/ds_k`` backward and reads off
+``dC/db_k`` at each step. All Jacobians are analytic because the LLG right-hand side is a sum of cross
+products, which are linear in each argument.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+
+from .dynamics.system import MacrospinSystem
+
+__all__ = ["AdjointResult", "adjoint_gradient", "optimize_pulse_adjoint"]
+
+
+def _cross_matrix(v: np.ndarray) -> np.ndarray:
+    """The skew matrix ``[v]_x`` such that ``[v]_x u = v x u``."""
+    return np.array(
+        [[0.0, -v[2], v[1]], [v[2], 0.0, -v[0]], [-v[1], v[0], 0.0]]
+    )
+
+
+def _f(s: np.ndarray, b_applied: np.ndarray, system: MacrospinSystem) -> np.ndarray:
+    """The LLG right-hand side for one moment."""
+    b_total = system.internal_field(s) + b_applied
+    alpha, gamma = system.alpha, system.gamma
+    precession = np.cross(s, b_total)
+    damping = np.cross(s, precession)
+    return (-gamma * precession - alpha * gamma * damping) / (1.0 + alpha**2)
+
+
+def _f_jacobians(
+    s: np.ndarray, b_applied: np.ndarray, system: MacrospinSystem
+) -> tuple[np.ndarray, np.ndarray]:
+    """The Jacobians ``df/ds`` and ``df/db`` at a point, each a 3x3 matrix.
+
+    The internal field is linear in s (a diagonal anisotropy matrix), so
+    ``b_total = A s + b_applied`` with ``A = diag(0, 0, 2K/mu)`` for the uniaxial easy axis. The LLG
+    right-hand side is a sum of cross products, each bilinear, so the Jacobians are assembled from
+    skew matrices.
+    """
+    alpha, gamma = system.alpha, system.gamma
+    scale = 1.0 / (1.0 + alpha**2)
+    anisotropy = np.diag([0.0, 0.0, 2.0 * system.anisotropy_j / system.mu])
+    b_total = anisotropy @ s + b_applied
+
+    sx = _cross_matrix(s)
+    btx = _cross_matrix(b_total)
+    # d/ds of (s x b_total) = [s]_x (A) - [b_total]_x, since b_total depends on s through A.
+    dprec_ds = sx @ anisotropy - btx
+    prec = np.cross(s, b_total)
+    px = _cross_matrix(prec)
+    # damping = s x prec; d/ds = [s]_x dprec_ds - [prec]_x
+    ddamp_ds = sx @ dprec_ds - px
+    df_ds = scale * (-gamma * dprec_ds - alpha * gamma * ddamp_ds)
+
+    # d/db_applied: b_total depends on b_applied as identity.
+    dprec_db = -btx * 0.0 + sx  # d(s x b_total)/db_applied = [s]_x
+    ddamp_db = sx @ dprec_db
+    df_db = scale * (-gamma * dprec_db - alpha * gamma * ddamp_db)
+    return df_ds, df_db
+
+
+@dataclass(frozen=True)
+class AdjointResult:
+    """The outcome of an adjoint-gradient pulse optimization.
+
+    Attributes:
+        times: the slice times, s, shape ``(N,)``.
+        field: the optimized field per slice, shape ``(N, 3)``, T.
+        cost: the switching cost, T^2 s.
+        final_sz: the final z-component.
+        switched: whether the moment reversed.
+        objective: the final penalized objective value.
+        iterations: the number of gradient steps taken.
+    """
+
+    times: np.ndarray
+    field: np.ndarray
+    cost: float
+    final_sz: float
+    switched: bool
+    objective: float
+    iterations: int
+
+
+def _forward(
+    field: np.ndarray, times: np.ndarray, system: MacrospinSystem
+) -> tuple[np.ndarray, np.ndarray]:
+    """Forward integration; returns the trajectory and the pre-normalization vectors."""
+    n = times.size
+    s = np.array([0.0, 0.0, 1.0])
+    trajectory = np.empty((n, 3))
+    pre_norm = np.empty((n, 3))
+    trajectory[0] = s
+    pre_norm[0] = s
+    for k in range(n - 1):
+        dt = times[k + 1] - times[k]
+        r = s + dt * _f(s, field[k], system)
+        pre_norm[k + 1] = r
+        s = r / np.linalg.norm(r)
+        trajectory[k + 1] = s
+    return trajectory, pre_norm
+
+
+def adjoint_gradient(
+    field: np.ndarray,
+    times: np.ndarray,
+    system: MacrospinSystem,
+    fidelity_weight: float,
+) -> tuple[float, float, np.ndarray]:
+    """The objective and its exact gradient with respect to the per-slice field.
+
+    Args:
+        field: the field per slice, shape ``(N, 3)``, T.
+        times: the slice times, s, shape ``(N,)``.
+        system: the macrospin.
+        fidelity_weight: the penalty weight on ``(1 + s_z(T)) / 2``.
+
+    Returns:
+        ``(objective, final_sz, gradient)`` with ``gradient`` of shape ``(N, 3)``.
+    """
+    n = times.size
+    trajectory, pre_norm = _forward(field, times, system)
+    final_sz = float(trajectory[-1, 2])
+
+    steps = np.diff(times)
+    cost = float(np.sum(np.sum(field[:-1] ** 2, axis=1) * steps))
+    objective = cost + fidelity_weight * 0.5 * (1.0 + final_sz)
+
+    grad = np.zeros_like(field)
+    # Terminal adjoint: dC/ds_N from the fidelity term.
+    g = np.array([0.0, 0.0, 0.5 * fidelity_weight])
+
+    for k in range(n - 2, -1, -1):
+        dt = times[k + 1] - times[k]
+        # Through the normalization s_{k+1} = r / |r|.
+        r = pre_norm[k + 1]
+        norm = np.linalg.norm(r)
+        s_next = r / norm
+        d_normalize = (np.eye(3) - np.outer(s_next, s_next)) / norm  # ds_{k+1}/dr
+        g_r = d_normalize.T @ g  # dC/dr_k
+
+        df_ds, df_db = _f_jacobians(trajectory[k], field[k], system)
+        # r = s_k + dt f(s_k, b_k): dr/ds_k = I + dt df_ds; dr/db_k = dt df_db.
+        # Direct cost term in b_k (this slice contributes |b_k|^2 dt).
+        grad[k] += 2.0 * field[k] * dt + dt * (df_db.T @ g_r)
+        # Propagate to s_k.
+        g = (np.eye(3) + dt * df_ds).T @ g_r
+
+    return objective, final_sz, grad
+
+
+def optimize_pulse_adjoint(
+    system: MacrospinSystem,
+    switching_time: float,
+    n_slices: int = 60,
+    fidelity_weight: float | None = None,
+    learning_rate: float | None = None,
+    max_iterations: int = 400,
+    seed: int = 0,
+) -> AdjointResult:
+    """Optimize a transverse field pulse by adjoint-gradient descent.
+
+    Args:
+        system: the macrospin (uniaxial).
+        switching_time: ``T`` in s.
+        n_slices: the number of piecewise-constant field slices.
+        fidelity_weight: the reversal penalty weight; a scaled default is used if omitted.
+        learning_rate: the gradient-descent step; a scaled default is used if omitted.
+        max_iterations: the descent iteration cap.
+        seed: the initial-guess seed.
+
+    Returns:
+        The :class:`AdjointResult`.
+    """
+    from scipy.optimize import minimize
+
+    from .analytic.uniaxial import cost_free_macrospin
+
+    del learning_rate  # the exact gradient is fed to L-BFGS, which sets its own step
+    times = np.linspace(0.0, switching_time, n_slices)
+    free = cost_free_macrospin(switching_time, system.alpha, system.gamma)
+    if fidelity_weight is None:
+        fidelity_weight = 50.0 * free
+
+    # Non-dimensionalize so L-BFGS works in O(1) variables: the field in units of the anisotropy
+    # field, the objective in units of the free-macrospin cost. Otherwise the objective and gradient
+    # are order 1e-12 and the optimizer's default tolerances are met before it takes a single step.
+    field_scale = system.anisotropy_field
+    objective_scale = max(free, 1e-300)
+
+    rng = np.random.default_rng(seed)
+    # Only the two transverse components per slice are optimized; the drive stays perpendicular.
+    initial = rng.normal(scale=0.3, size=(n_slices, 2)).reshape(-1)
+
+    def unpack(vector: np.ndarray) -> np.ndarray:
+        field = np.zeros((n_slices, 3))
+        field[:, :2] = vector.reshape(n_slices, 2) * field_scale
+        return field
+
+    def objective_and_grad(vector: np.ndarray) -> tuple[float, np.ndarray]:
+        field = unpack(vector)
+        objective, _final_sz, grad = adjoint_gradient(field, times, system, fidelity_weight)
+        scaled_grad = grad[:, :2].reshape(-1) * field_scale / objective_scale
+        return objective / objective_scale, scaled_grad
+
+    result = minimize(
+        objective_and_grad,
+        initial,
+        method="L-BFGS-B",
+        jac=True,
+        options={"maxiter": max_iterations, "ftol": 1e-12, "gtol": 1e-9},
+    )
+
+    field = unpack(result.x)
+    objective, final_sz, _ = adjoint_gradient(field, times, system, fidelity_weight)
+    cost = float(np.sum(np.sum(field[:-1] ** 2, axis=1) * np.diff(times)))
+    return AdjointResult(
+        times=times,
+        field=field,
+        cost=cost,
+        final_sz=final_sz,
+        switched=final_sz < 0.0,
+        objective=objective,
+        iterations=int(result.nit),
+    )

spinoct-0.10.0/src/spinoct/amortized.py

@@ -0,0 +1,218 @@
+"""An amortized learned policy for optimal switching pulses (R15).
+
+Every solver in this package re-optimizes from scratch for each new set of material parameters. An
+amortized policy instead learns the map from parameters to the optimal pulse once, and then emits a
+near-optimal pulse for any new parameters instantly, with no optimization at inference. This is useful
+when a pulse must be chosen online, and it is a clean testbed for a learned controller because the
+uniaxial optimum is known in closed form, so the policy's claim is checkable rather than merely
+plausible.
+
+The output representation matters, and getting it wrong is instructive. Regressing a learned model onto
+the raw pulse amplitude profile fails: switching is a threshold phenomenon, and a profile fit in a
+least-squares sense is often a few percent too weak to complete the reversal, so the emitted pulse does
+not switch even though it looks close. The fix is to amortize the map onto the pulse's single
+physically-meaningful shape parameter, the elliptic parameter ``p`` that the closed-form solution is
+built from. Any ``p`` yields a genuine optimal-control pulse that reverses the moment; a slightly wrong
+``p`` yields the optimal pulse for a slightly different switching time, which still switches and is only
+slightly suboptimal. So the shape-parameter policy is both reliable and near-optimal, while the
+profile policy is neither. That contrast is the finding.
+
+The policy is a small multilayer perceptron in numpy, so the package keeps no autodiff dependency. It
+maps the two dimensionless parameters that determine the optimal uniaxial pulse, the damping ``alpha``
+and the log switching time ``ln(T/tau0)``, to ``ln(p)``, and is trained by regression against the
+analytic shape parameter on a grid. It is then validated the only honest way: the emitted pulse is fed
+to the equation of motion and its cost and outcome are compared to the analytic optimum on parameter
+combinations the policy never saw in training.
+
+The pre-declared acceptance criterion (from the project's research dossier): a learned controller that
+cannot reach the analytic optimum on the uniaxial case, where the optimum is known, has no business
+being trusted on the harder cases where it is not. This module is that gate, and the shape-parameter
+policy passes it.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+
+from .analytic.uniaxial import UniaxialOptimalControl, solve_shape_parameter
+from .dynamics.llg import integrate_llg_tabulated, switching_cost
+from .dynamics.system import MacrospinSystem
+
+__all__ = ["AmortizedPolicy", "PolicyEvaluation", "evaluate_policy", "train_amortized_policy"]
+
+#: The number of amplitude samples the emitted pulse is reconstructed on for integration.
+_N_SAMPLES = 200
+
+
+def _features(system: MacrospinSystem, switching_time: float) -> np.ndarray:
+    """The two dimensionless inputs the policy sees: ``[alpha, ln(T/tau0)]``."""
+    return np.array([system.alpha, np.log(switching_time / system.tau0)])
+
+
+@dataclass
+class AmortizedPolicy:
+    """A trained policy that emits the optimal-pulse shape parameter for given parameters.
+
+    Attributes:
+        w1, b1, w2, b2: the two-layer network weights and biases.
+        input_mean, input_std: the input normalization.
+        target_mean, target_std: the ``ln(p)`` target normalization.
+    """
+
+    w1: np.ndarray
+    b1: np.ndarray
+    w2: np.ndarray
+    b2: np.ndarray
+    input_mean: np.ndarray
+    input_std: np.ndarray
+    target_mean: float
+    target_std: float
+
+    def _forward(self, features: np.ndarray) -> float:
+        x = (features - self.input_mean) / self.input_std
+        h = np.tanh(x @ self.w1 + self.b1)
+        y = float(h @ self.w2 + self.b2)
+        return y * self.target_std + self.target_mean
+
+    def shape_parameter(self, system: MacrospinSystem, switching_time: float) -> float:
+        """The predicted shape parameter ``p`` for a system and switching time."""
+        return float(np.exp(self._forward(_features(system, switching_time))))
+
+    def pulse(self, system: MacrospinSystem, switching_time: float) -> UniaxialOptimalControl:
+        """The optimal-control pulse the policy emits, built from the predicted shape parameter."""
+        return UniaxialOptimalControl(
+            system=system, switching_time=switching_time, p=self.shape_parameter(system, switching_time)
+        )
+
+
+@dataclass(frozen=True)
+class PolicyEvaluation:
+    """The result of applying the policy's pulse to the equation of motion.
+
+    Attributes:
+        cost: the switching cost of the emitted pulse, T^2 s.
+        analytic_cost: the closed-form optimal cost for the same system and switching time, T^2 s.
+        cost_ratio: ``cost / analytic_cost``, at least 1 for a valid pulse up to discretization.
+        predicted_p: the shape parameter the policy emitted.
+        true_p: the exact shape parameter for these parameters.
+        final_sz: the final z-component under the emitted pulse.
+        switched: whether the emitted pulse reversed the moment.
+    """
+
+    cost: float
+    analytic_cost: float
+    cost_ratio: float
+    predicted_p: float
+    true_p: float
+    final_sz: float
+    switched: bool
+
+
+def evaluate_policy(
+    policy: AmortizedPolicy, system: MacrospinSystem, switching_time: float
+) -> PolicyEvaluation:
+    """Apply the policy's emitted pulse to the equation of motion and score it against the optimum.
+
+    Args:
+        policy: the trained policy.
+        system: the macrospin (uniaxial).
+        switching_time: ``T`` in s.
+
+    Returns:
+        The :class:`PolicyEvaluation`.
+    """
+    emitted = policy.pulse(system, switching_time)
+    optimal = UniaxialOptimalControl.for_switching_time(system, switching_time)
+    grid = np.linspace(0.0, switching_time, _N_SAMPLES)
+    field_table = emitted.field_vector(grid)
+    trajectory = integrate_llg_tabulated(np.array([0.0, 0.0, 1.0]), field_table, grid, system)
+    final_sz = float(trajectory[-1, 2])
+    cost = switching_cost(grid, field_table)
+    analytic_cost = optimal.cost()
+    return PolicyEvaluation(
+        cost=cost,
+        analytic_cost=analytic_cost,
+        cost_ratio=cost / analytic_cost if analytic_cost > 0 else float("inf"),
+        predicted_p=emitted.p,
+        true_p=optimal.p,
+        final_sz=final_sz,
+        switched=final_sz < 0.0,
+    )
+
+
+def train_amortized_policy(
+    mu: float,
+    anisotropy_j: float,
+    alphas: np.ndarray,
+    switching_times_tau0: np.ndarray,
+    hidden: int = 24,
+    epochs: int = 6000,
+    learning_rate: float = 0.05,
+    seed: int = 0,
+) -> AmortizedPolicy:
+    """Train the amortized policy by regression onto the analytic shape parameter on a grid.
+
+    Args:
+        mu: magnetic moment, J/T (fixed across the training set).
+        anisotropy_j: anisotropy energy, J (fixed).
+        alphas: the damping values to train on.
+        switching_times_tau0: the switching times to train on, in units of tau0.
+        hidden: the hidden-layer width.
+        epochs: the number of full-batch gradient steps.
+        learning_rate: the step size.
+        seed: the weight-initialization seed.
+
+    Returns:
+        The trained :class:`AmortizedPolicy`.
+    """
+    features = []
+    targets = []
+    for alpha in alphas:
+        system = MacrospinSystem(mu=mu, anisotropy_j=anisotropy_j, alpha=float(alpha))
+        for t_tau0 in switching_times_tau0:
+            switching_time = system.switching_time_from_tau0(float(t_tau0))
+            features.append(_features(system, switching_time))
+            targets.append(np.log(solve_shape_parameter(switching_time, system)))
+    features = np.array(features)
+    targets = np.array(targets)
+
+    input_mean, input_std = features.mean(0), features.std(0) + 1e-12
+    target_mean, target_std = float(targets.mean()), float(targets.std() + 1e-12)
+    x = (features - input_mean) / input_std
+    y = (targets - target_mean) / target_std
+
+    rng = np.random.default_rng(seed)
+    n_in = x.shape[1]
+    w1 = rng.normal(scale=1.0 / np.sqrt(n_in), size=(n_in, hidden))
+    b1 = np.zeros(hidden)
+    w2 = rng.normal(scale=1.0 / np.sqrt(hidden), size=hidden)
+    b2 = 0.0
+
+    n = x.shape[0]
+    for _ in range(epochs):
+        h_pre = x @ w1 + b1
+        h = np.tanh(h_pre)
+        pred = h @ w2 + b2
+        error = pred - y
+        grad_w2 = h.T @ error / n
+        grad_b2 = float(error.mean())
+        grad_h = np.outer(error, w2) * (1.0 - h**2)
+        grad_w1 = x.T @ grad_h / n
+        grad_b1 = grad_h.mean(0)
+        w1 -= learning_rate * grad_w1
+        b1 -= learning_rate * grad_b1
+        w2 -= learning_rate * grad_w2
+        b2 -= learning_rate * grad_b2
+
+    return AmortizedPolicy(
+        w1=w1,
+        b1=b1,
+        w2=w2,
+        b2=b2,
+        input_mean=input_mean,
+        input_std=input_std,
+        target_mean=target_mean,
+        target_std=target_std,
+    )