softopt 0.7.0__tar.gz → 0.7.3__tar.gz

{softopt-0.7.0/softopt.egg-info → softopt-0.7.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: softopt
-Version: 0.7.0
+Version: 0.7.3
 Summary: An optimization layer that closes the gap between your model and the real system, built on Klein–Maimon soft-number calculus
 Author: Ido Angel
 License: MIT

{softopt-0.7.0 → softopt-0.7.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "softopt"
-version = "0.7.0"
+version = "0.7.3"
 description = "An optimization layer that closes the gap between your model and the real system, built on Klein–Maimon soft-number calculus"
 readme = "README.md"
 license = { text = "MIT" }

{softopt-0.7.0 → softopt-0.7.3}/softopt/__init__.py

@@ -6,6 +6,7 @@ model (a simulator, a circuit, a physical law) and the real system you measure.
     from softopt import spsa_gradient    # gradient, direction and measured slope from 2 readings
     from softopt import SoftOptTorch     # PyTorch version (torch.optim.Optimizer)
     from softopt import soft_compile     # turn a plain-Python model into g_delta
+    from softopt import torch_compile    # the same for a PyTorch model (forward-mode = soft numbers)
     from softopt import SoftNumber       # the algebraic primitive itself
 
 calibrate=True (new in 0.7.0): soft_compile(model, n_params, n_bias=k) for model(theta, beta), then
@@ -25,12 +26,15 @@ def __getattr__(name):
     if name == "SoftOptTorch":
         from .torch_optimizer import SoftOptTorch
         return SoftOptTorch
+    if name == "torch_compile":
+        from ._torch_compile import torch_compile
+        return torch_compile
     raise AttributeError(f"module 'softopt' has no attribute {name!r}")
 
 
 __all__ = [
-    "SoftOpt", "SoftOptTorch", "SoftNumber", "soft_compile", "spsa_gradient",
+    "SoftOpt", "SoftOptTorch", "SoftNumber", "soft_compile", "torch_compile", "spsa_gradient",
     "sadd", "sneg", "ssub", "smul", "sinv", "sdiv",
     "spow", "sroot", "ssqrt", "ssin", "scos", "sexp", "slog",
 ]
-__version__ = "0.7.0"
+__version__ = "0.7.3"

softopt-0.7.3/softopt/_torch_compile.py

@@ -0,0 +1,115 @@
+"""
+torch_compile: soft_compile for models written in PyTorch -- a neural network, a differentiable sampler, any
+torch computation -- so SoftOpt (default and calibrate=True) can use them.
+
+    from softopt import SoftOpt, torch_compile
+
+    def model(theta, beta):            # torch tensors in, a scalar tensor out
+        ...
+
+    g = torch_compile(model, n_params=n, n_bias=k)
+    opt = SoftOpt(n, g, lr=0.02, calibrate=True, bias_scale=0.1)
+
+How the soft algebra is carried
+  A soft number a*eps + b (eps^2 = 0) is what PyTorch's forward-mode automatic differentiation propagates:
+  every value b travels with a tangent a, and every operation applies f(a eps + b) = a f'(b) eps + f(b)
+  (Lemma 6.1). torch.func.jvp is that single-axis soft evaluation; a jvp inside a jvp is the NESTED soft number
+  ((e, a), (c, v)) with two axes eps1, eps2:
+    eps1 and eps2 both along d         ->  v = f, c = a = D1, e = D2           (d1_d2)
+    eps1 along d, eps2 along beta_j     ->  e = d D1 / d beta_j                (d1_and_jac: the mixed part)
+  d1_and_jac carries all beta_j axes in one vectorized pass (0.7.3); it falls back to one pass per beta_j for
+  models whose operations cannot be batched.
+  The same quantities softopt.calibrate.CompiledModel computes with SoftNumber objects, so the calibrator and
+  the optimizer are unchanged. Checked against CompiledModel in tests/test_torch_compile.py.
+"""
+import numpy as np
+
+
+def _torch():
+    import torch
+    from torch.func import jvp
+    return torch, jvp
+
+
+class TorchModel:
+    """model_fn(theta, beta) (or model_fn(theta) when n_bias=0) on nested soft numbers via nested jvp."""
+
+    def __init__(self, model_fn, n_bias=0, dtype=None, device=None):
+        torch, _ = _torch()
+        self.f, self.nb = model_fn, int(n_bias)
+        self.dtype = torch.float64 if dtype is None else dtype
+        self.device = device
+        self._vectorized = True
+
+    def _t(self, v):
+        torch, _ = _torch()
+        return torch.as_tensor(np.asarray(v, dtype=float), dtype=self.dtype, device=self.device)
+
+    def _call(self, th, bs):
+        return self.f(th, bs) if self.nb else self.f(th)
+
+    def value(self, x, beta=()):
+        torch, _ = _torch()
+        with torch.no_grad():
+            return float(self._call(self._t(x), self._t(beta)))
+
+    def _d1(self, x, d, b):
+        _, jvp = _torch()
+        return jvp(lambda t: self._call(t, b), (x,), (d,))[1]
+
+    def d1_d2(self, x, d, beta=()):
+        """D1 and D2 along d: eps1 and eps2 both along d."""
+        _, jvp = _torch()
+        x, d, b = self._t(x), self._t(d), self._t(beta)
+        D1, D2 = jvp(lambda t: self._d1(t, d, b), (x,), (d,))
+        return float(D1), float(D2)
+
+    def d1_and_jac(self, x, d, beta=()):
+        """D1 along d, and d D1 / d beta_j for every j: eps1 along d, eps2 along beta_j (the mixed part)."""
+        torch, jvp = _torch()
+        x, d, b = self._t(x), self._t(d), self._t(beta)
+        if not self.nb:
+            return float(self._d1(x, d, b)), np.zeros(0)
+        if self._vectorized:
+            # all bias tangents in one pass: a soft number carrying one eps2 axis per bias parameter (vmap of the
+            # forward-mode slope over beta). Same numbers as the loop below, about nb times fewer passes.
+            try:
+                from torch.func import jacfwd
+                def slope(bb):
+                    y = self._d1(x, d, bb)
+                    return y, y
+                jac, D1 = jacfwd(slope, has_aux=True)(b)
+                return float(D1), jac.detach().cpu().numpy().astype(float).reshape(self.nb)
+            except Exception:
+                self._vectorized = False      # a model op without batching support: use one pass per bias
+        jac, D1 = np.zeros(self.nb), None
+        for j in range(self.nb):
+            e = torch.zeros_like(b); e[j] = 1.0
+            D1, jac[j] = jvp(lambda bb: self._d1(x, d, bb), (b,), (e,))
+        return float(D1), jac
+
+
+def torch_compile(model_fn, n_params=None, n_bias=0, dtype=None, device=None, verify=True):
+    """model_fn(theta) -> scalar tensor, or model_fn(theta, beta) with n_bias > 0 (beta = offsets of the
+    model's uncertain parameters, 0 = nominal). Returns g_delta for SoftOpt, carrying .d1_d2 (exact D2) and
+    .calibration (for calibrate=True). dtype defaults to float64; pass torch.float32 for large networks."""
+    M = TorchModel(model_fn, n_bias, dtype, device)
+    zero = np.zeros(M.nb)
+
+    def g_delta(theta, delta):
+        return M.d1_and_jac(theta, delta, zero)[0]
+
+    g_delta.d1_d2 = lambda theta, delta: M.d1_d2(theta, delta, zero)
+    g_delta.calibration = M
+    g_delta.value = lambda theta: M.value(theta, zero)
+    if verify and n_params:
+        rng = np.random.default_rng(0)
+        x, d = rng.normal(size=n_params), rng.choice([-1.0, 1.0], size=n_params)
+        h = 1e-4
+        fd = (M.value(x + h * d, zero) - M.value(x - h * d, zero)) / (2 * h)
+        D1 = g_delta(x, d)
+        tol = 1e-4 if M.dtype.itemsize == 8 else 1e-2
+        if not np.isfinite(D1) or abs(D1 - fd) > tol * max(1.0, abs(fd)):
+            raise ValueError(f'torch_compile: the forward-mode slope {D1:.6g} disagrees with finite differences '
+                             f'{fd:.6g} -- the model may not be differentiable here, or it detaches theta')
+    return g_delta

{softopt-0.7.0 → softopt-0.7.3}/softopt/calibrate.py

@@ -28,6 +28,10 @@ Two ways to use it
     their expected size (a number or one per parameter), on the PHYSICAL scale -- a prior far too wide lets the
     filter fit noise.
 
+The discrepancy dg enters the slope only as far as it stands out of its own uncertainty,
+wg = (dg.e)^2 / ((dg.e)^2 + var(dg.e)) (new in 0.7.2): with noisy readings dg would otherwise absorb the noise
+and cancel the model's soft-number slope.
+
 Validated (development + held-out seeds) where the bias structure is right: gate calibration, adaptive optics,
 H2 with coherent gate errors, robot sim-to-real. When the model lacks a major error type the median is still at
 or below Adam but single runs can be worse; when the model is already exact, calibrate=False is better.
@@ -135,8 +139,15 @@ class Calibrator:
         beta, dg = st[self.ib], st[self.ig]
         D1m, D2m = self.M.d1_d2(x, e, beta)
         _, jac_e = self.M.d1_and_jac(x, e, beta)
-        D1 = D1m + dg @ e
-        h = np.r_[jac_e, 0.0, e]
+        # dg is the part of the slope the model cannot explain. Under noisy readings it also absorbs the noise,
+        # so it enters only as far as it stands out of its own uncertainty (wg); the soft-number slope of the
+        # calibrated model is always kept.
+        dge = float(dg @ e)
+        hg = np.r_[np.zeros_like(jac_e), 0.0, e]
+        varg = max(float(hg @ P @ hg), 0.0)
+        wg = dge ** 2 / (dge ** 2 + varg) if (dge ** 2 + varg) > 0 else 0.0
+        D1 = D1m + wg * dge
+        h = np.r_[jac_e, 0.0, wg * e]
         var = max(float(h @ P @ h), 0.0)
         t = -D1 / D2m if D2m > 1e-8 else -eta_fallback * D1
         w = D1 ** 2 / (D1 ** 2 + var) if (D1 ** 2 + var) > 0 else 0.0
@@ -144,7 +155,7 @@ class Calibrator:
         if step != 0.0:
             f0 = self.M.value(x, beta)
             for _ in range(20):
-                df = self.M.value(x + step * e, beta) - f0 + step * (dg @ e)
+                df = self.M.value(x + step * e, beta) - f0 + step * wg * dge
                 if df < 0:
                     break
                 c = (df - D1 * step) / step ** 2

{softopt-0.7.0 → softopt-0.7.3/softopt.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: softopt
-Version: 0.7.0
+Version: 0.7.3
 Summary: An optimization layer that closes the gap between your model and the real system, built on Klein–Maimon soft-number calculus
 Author: Ido Angel
 License: MIT

{softopt-0.7.0 → softopt-0.7.3}/softopt.egg-info/SOURCES.txt

@@ -2,6 +2,7 @@ LICENSE
 README.md
 pyproject.toml
 softopt/__init__.py
+softopt/_torch_compile.py
 softopt/calibrate.py
 softopt/compile.py
 softopt/numpy_optimizer.py
@@ -19,4 +20,5 @@ tests/test_compile.py
 tests/test_exact_d2.py
 tests/test_measured_slope.py
 tests/test_soft_number.py
-tests/test_torch.py
+tests/test_torch.py
+tests/test_torch_compile.py

softopt-0.7.3/tests/test_torch_compile.py

@@ -0,0 +1,97 @@
+"""torch_compile (nested jvp) must give exactly what soft_compile (nested SoftNumber) gives."""
+import numpy as np
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from softopt import SoftOpt, soft_compile, spsa_gradient   # noqa: E402
+from softopt import torch_compile                          # noqa: E402
+
+N, K = 5, 3
+
+
+def model_soft(theta, beta):
+    x = theta
+    s = x[0] * x[0] * (1 + beta[0]) + np.sin(x[1] - 0.3 * x[2]) * (2 + beta[1])
+    s = s + np.exp(x[3] * x[4] + beta[2] * x[0]) * 0.1 + x[2] * x[2] * x[2] / (2 + x[4] * x[4])
+    return s
+
+
+def model_torch(theta, beta):
+    x = theta
+    s = x[0] * x[0] * (1 + beta[0]) + torch.sin(x[1] - 0.3 * x[2]) * (2 + beta[1])
+    s = s + torch.exp(x[3] * x[4] + beta[2] * x[0]) * 0.1 + x[2] ** 3 / (2 + x[4] * x[4])
+    return s
+
+
+def test_matches_soft_compile():
+    gs = soft_compile(model_soft, n_params=N, n_bias=K)
+    gt = torch_compile(model_torch, n_params=N, n_bias=K)
+    rng = np.random.default_rng(1)
+    for _ in range(5):
+        x, d, b = rng.normal(size=N), rng.choice([-1.0, 1.0], size=N), 0.1 * rng.normal(size=K)
+        assert abs(gs(x, d) - gt(x, d)) < 1e-10
+        Ms, Mt = gs.calibration, gt.calibration
+        for a, c in zip(Ms.d1_d2(x, d, b), Mt.d1_d2(x, d, b)):
+            assert abs(a - c) < 1e-10
+        (D1s, Js), (D1t, Jt) = Ms.d1_and_jac(x, d, b), Mt.d1_and_jac(x, d, b)
+        assert abs(D1s - D1t) < 1e-10 and np.max(np.abs(Js - Jt)) < 1e-10
+        assert abs(Ms.value(x, b) - Mt.value(x, b)) < 1e-12
+
+
+def test_calibrate_trajectory_identical():
+    """The whole calibrate=True optimisation is the same with either compiler."""
+    true_beta = np.array([0.2, -0.3, 0.1])
+
+    def measure(x, rng):
+        return model_torch(torch.tensor(x), torch.tensor(true_beta)).item() + 0.01 * rng.normal()
+
+    xs = []
+    for comp, fn in ((soft_compile, model_soft), (torch_compile, model_torch)):
+        g = comp(fn, n_params=N, n_bias=K)
+        opt = SoftOpt(N, g, lr=0.05, calibrate=True, bias_scale=0.3, seed=3)
+        x, rng = np.full(N, 0.4), np.random.default_rng(7)
+        for _ in range(25):
+            grad, d, slope = spsa_gradient(lambda t: measure(t, rng), x, c=0.1, rng=rng)
+            x = opt.step(x, grad, direction=d, measured_slope=slope)
+        xs.append(x)
+    assert np.max(np.abs(xs[0] - xs[1])) < 1e-8
+
+
+def test_plain_model_and_exact_d2():
+    g = torch_compile(lambda t: (t ** 2).sum() + torch.cos(t[0]), n_params=3)
+    D1, D2 = g.d1_d2(np.array([0.5, 1.0, -1.0]), np.array([1.0, -1.0, 1.0]))
+    assert abs(D2 - (6 - np.cos(0.5))) < 1e-12
+    assert abs(D1 - (1.0 - 2.0 - 2.0 - np.sin(0.5))) < 1e-12
+
+
+def test_detached_model_is_caught():
+    with pytest.raises(ValueError):
+        torch_compile(lambda t: torch.tensor(float(t.sum().detach())), n_params=2)
+
+
+def test_vectorized_bias_jacobian_equals_loop():
+    """0.7.3: d1_and_jac carries all bias axes in one vectorized pass; it must equal one pass per bias."""
+    gv = torch_compile(model_torch, n_params=N, n_bias=K)
+    gl = torch_compile(model_torch, n_params=N, n_bias=K)
+    gl.calibration._vectorized = False
+    rng = np.random.default_rng(2)
+    for _ in range(5):
+        x, d, b = rng.normal(size=N), rng.choice([-1.0, 1.0], size=N), 0.1 * rng.normal(size=K)
+        (D1v, Jv), (D1l, Jl) = gv.calibration.d1_and_jac(x, d, b), gl.calibration.d1_and_jac(x, d, b)
+        assert gv.calibration._vectorized
+        assert abs(D1v - D1l) < 1e-12 and np.max(np.abs(Jv - Jl)) < 1e-12
+
+
+def test_unbatchable_model_falls_back_to_loop():
+    """A model op that cannot be batched (here .item() on a bias) must fall back to the per-bias loop."""
+    def model_item(theta, beta):
+        if float(beta[0].detach()) > 1e9:           # data-dependent Python branch: not vmap-able
+            return theta.sum()
+        return model_torch(theta, beta)
+    g = torch_compile(model_item, n_params=N, n_bias=K, verify=False)
+    ref = torch_compile(model_torch, n_params=N, n_bias=K)
+    rng = np.random.default_rng(3)
+    x, d, b = rng.normal(size=N), rng.choice([-1.0, 1.0], size=N), 0.1 * rng.normal(size=K)
+    (D1, J), (D1r, Jr) = g.calibration.d1_and_jac(x, d, b), ref.calibration.d1_and_jac(x, d, b)
+    assert abs(D1 - D1r) < 1e-12 and np.max(np.abs(J - Jr)) < 1e-12