Added a Probana backend and a real device solver

p_kit/backends/probana_backend.py

@@ -0,0 +1,209 @@
+"""
+Probana backend for p-kit.
+
+Probana is an experimental physical probabilistic computer:
+https://github.com/toncho11/probana
+
+Communicates with the Probana physical p-bit computer over USB serial.
+
+Protocol commands:
+  PING
+  CLEAR <n>
+  H <i> <value>
+  J <from> <to> <value>
+  ANNEAL CONSTANT <scale>
+  ANNEAL LINEAR <start> <end> <steps>
+  COMMIT
+  RUN <samples> <burn_in> <thin>
+
+J and h are uploaded once. During RUN, the p-bit update loop,
+calibration correction and physical sampling are performed locally
+on the Probana board.
+"""
+
+import time
+import numpy as np
+from .numpy_backend import NumpyBackend
+
+
+class ProbanaBackend(NumpyBackend):
+    """Backend for the Probana physical p-bit computer."""
+
+    protocol_name = "PROBANA"
+    protocol_version = 1
+    supports_annealing = True
+    supports_native_pbits = True
+
+    def __init__(self, port=None, baudrate=115200, timeout=2.0,
+                 startup_timeout=60.0, dtype=None):
+        super().__init__(dtype=dtype)
+        try:
+            import serial
+            from serial.tools import list_ports
+        except ImportError as e:
+            raise ImportError("Install pyserial: pip install pyserial") from e
+
+        self.serial, self.list_ports = serial, list_ports
+        self.timeout = timeout
+        self.port = port or self._find_port()
+        self.ser = serial.Serial(self.port, baudrate, timeout=0.2)
+        self.n_pbits = self._wait_ready(startup_timeout)
+        self.ser.timeout = timeout
+
+        self.supports_pkit_annealing_func = False
+
+    def _find_port(self):
+        ports = list(self.list_ports.comports())
+        if not ports:
+            raise RuntimeError("No serial device found")
+        if len(ports) == 1:
+            return ports[0].device
+
+        keys = ("arduino", "rp2", "pico", "esp32", "cp210", "ch340")
+        found = [p for p in ports
+                 if any(k in (p.description or "").lower() for k in keys)]
+        if len(found) == 1:
+            return found[0].device
+
+        raise RuntimeError(
+            "Multiple serial devices found; specify port: " +
+            ", ".join(p.device for p in ports)
+        )
+
+    def _send(self, cmd):
+        self.ser.write((cmd + "\n").encode("ascii"))
+        self.ser.flush()
+
+    def _read(self, timeout=None):
+        end = time.monotonic() + (timeout or self.timeout)
+        while time.monotonic() < end:
+            raw = self.ser.readline()
+            if raw:
+                return raw.decode("ascii", errors="replace").strip()
+        raise TimeoutError("Probana did not respond")
+
+    def _parse_ready(self, line):
+        p = line.split()
+        if len(p) != 5 or p[0] not in ("READY", "OK") or \
+           p[1] != self.protocol_name or p[4] != "WORKING":
+            return None
+
+        if int(p[2]) != self.protocol_version:
+            raise RuntimeError(f"Unsupported protocol version {p[2]}")
+        return int(p[3])
+
+    def _wait_ready(self, timeout):
+        end, next_ping = time.monotonic() + timeout, 0
+        while time.monotonic() < end:
+            now = time.monotonic()
+            if now >= next_ping:
+                self._send("PING")
+                next_ping = now + 0.5
+
+            raw = self.ser.readline()
+            if not raw:
+                continue
+
+            line = raw.decode("ascii", errors="replace").strip()
+            n = self._parse_ready(line)
+            if n is not None:
+                return n
+            if line.startswith("ERR "):
+                raise RuntimeError(line)
+
+        raise TimeoutError("Probana did not enter WORKING mode")
+
+    def _ok(self):
+        while True:
+            line = self._read()
+            if line == "OK":
+                return
+            if line.startswith("ERR "):
+                raise RuntimeError(line)
+
+    def ping(self):
+        self._send("PING")
+        line = self._read()
+        n = self._parse_ready(line)
+        if n is None:
+            raise RuntimeError(f"Invalid PING response: {line}")
+        return n
+
+    def set_annealing(self, mode="constant", start=1.0, end=None, steps=1):
+        mode = mode.lower()
+
+        if mode == "constant":
+            self._send(f"ANNEAL CONSTANT {float(start):.9g}")
+        elif mode == "linear":
+            if end is None or int(steps) < 1:
+                raise ValueError("Linear annealing requires end and steps >= 1")
+            self._send(
+                f"ANNEAL LINEAR {float(start):.9g} "
+                f"{float(end):.9g} {int(steps)}"
+            )
+        else:
+            raise ValueError(f"Unsupported annealing mode: {mode}")
+
+        self._ok()
+
+    def load_circuit(self, J, h, annealing=("constant", 1.0)):
+        J = np.asarray(J, dtype=float)
+        h = np.asarray(h, dtype=float).reshape(-1)
+
+        if J.ndim != 2 or J.shape[0] != J.shape[1]:
+            raise ValueError("J must be square")
+        if J.shape[0] != h.size:
+            raise ValueError("J and h sizes do not match")
+        if h.size > self.n_pbits:
+            raise ValueError(
+                f"Circuit needs {h.size} p-bits; Probana has {self.n_pbits}"
+            )
+
+        self._send(f"CLEAR {h.size}"); self._ok()
+
+        for i, v in enumerate(h):
+            if v != 0:
+                self._send(f"H {i} {v:.9g}"); self._ok()
+
+        rows, cols = np.nonzero(J)
+        for src, dst in zip(rows, cols):
+            self._send(f"J {src} {dst} {J[src,dst]:.9g}"); self._ok()
+
+        self._send("COMMIT"); self._ok()
+        self.set_annealing(*annealing)
+
+    def run_circuit(self, J, h, samples, burn_in=100, thin=1, annealing=("constant", 1.0)):
+        self.load_circuit(J, h, annealing)
+        n = len(h)
+
+        self._send(f"RUN {int(samples)} {int(burn_in)} {int(thin)}")
+        self._ok()
+        states = []
+
+        while True:
+            line = self._read(max(self.timeout, self.timeout * samples))
+            if line == "DONE":
+                break
+            if line.startswith("ERR "):
+                raise RuntimeError(line)
+            if not line.startswith("S "):
+                continue
+
+            bits = line[2:].strip()
+            if len(bits) != n or any(b not in "01" for b in bits):
+                raise RuntimeError(f"Invalid sample: {line}")
+            states.append([1 if b == "1" else -1 for b in bits])
+
+        if len(states) != samples:
+            raise RuntimeError(f"Expected {samples} samples, got {len(states)}")
+        return np.asarray(states, dtype=np.int8)
+
+    def close(self):
+        if getattr(self, "ser", None) and self.ser.is_open:
+            self.ser.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        self.close()

p_kit/solver/real_device_solver.py

@@ -0,0 +1,112 @@
+"""
+Solver for physical p-bit devices.
+
+For example, it can be used with the physical probabilistic computer Probana:
+https://github.com/toncho11/probana
+
+RealDeviceSolver connects p-kit's PCircuit interface to a physical-device
+backend. The solver passes J and h to the backend and returns sampled states
+in the standard p-kit solver format.
+
+The Probana backend handles USB communication, while Probana performs the
+stochastic p-bit updates, calibration correction, annealing, and sampling
+locally on the board.
+
+Multiple shots are executed sequentially on the physical device, avoiding
+parallel access to a single USB connection.
+"""
+
+import numpy as np
+
+from .base_solver import Solver
+from .annealing import constant, linear
+
+
+class RealDeviceSolver(Solver):
+    def __init__(self, Nt, backend, i0=1.0, burn_in=100, thin=1):
+        if not getattr(backend, "supports_native_pbits", False):
+            raise TypeError("Backend does not support physical p-bits")
+        if burn_in < 0 or thin < 1:
+            raise ValueError("burn_in must be >= 0 and thin must be >= 1")
+
+        super().__init__(Nt=Nt, dt=1.0, i0=i0, backend=backend)
+        self.burn_in = int(burn_in)
+        self.thin = int(thin)
+
+    def _annealing(self, func):
+        if func is None or func is constant:
+            return ("constant", self.i0)
+
+        if func is linear:
+            start = float(func(self, 0))
+            end = float(func(self, self.Nt - 1))
+            steps = self.burn_in + self.Nt * self.thin
+            return ("linear", start, end, max(1, steps))
+
+        raise NotImplementedError(
+            "RealDeviceSolver currently supports constant and linear annealing"
+        )
+
+    def _sample_scales(self, annealing):
+        if annealing[0] == "constant":
+            return np.full(self.Nt, annealing[1], dtype=float)
+
+        _, start, end, steps = annealing
+        idx = self.burn_in + (np.arange(self.Nt) + 1) * self.thin - 1
+        x = np.minimum(1.0, idx / max(1, steps - 1))
+        return start + x * (end - start)
+
+    def solve(self, c, annealing_func=constant, n_shots=1,
+              bias_func=None, return_filtered=False,
+              initial_state=None, return_final=False):
+
+        if n_shots < 1:
+            raise ValueError("n_shots must be >= 1")
+
+        if bias_func is not None or return_filtered or initial_state is not None:
+            raise NotImplementedError(
+                "Dynamic bias, filtering and initial_state are not supported"
+            )
+
+        if (annealing_func not in (None, constant) and not getattr(self.backend, "supports_pkit_annealing_func", True)):
+            raise NotImplementedError(
+                "This p-kit annealing schedule cannot be mapped directly to this "
+                "hardware backend. Support for a custom annealing_func is disabled."
+            )
+
+        J = np.asarray(c.J, dtype=float)
+        h = np.asarray(c.h, dtype=float).reshape(-1)
+        annealing = self._annealing(annealing_func)
+
+        shots = [
+            self.backend.run_circuit(
+                J, h, samples=self.Nt,
+                burn_in=self.burn_in, thin=self.thin,
+                annealing=annealing
+            )
+            for _ in range(n_shots)
+        ]
+
+        all_m = np.stack(shots, axis=1)  # (Nt, n_shots, n_pbits)
+
+        if return_final:
+            return all_m[-1, 0] if n_shots == 1 else all_m[-1]
+
+        if n_shots > 1:
+            return all_m
+
+        m = all_m[:, 0]
+        scale = self._sample_scales(annealing)
+        I = scale[:, None] * (m @ J + h)
+
+        E = self.i0 * (
+            m @ h + 0.5 * np.einsum("bi,ij,bj->b", m, J, m)
+        )
+
+        return I, m, E
+
+    def copy(self):
+        return RealDeviceSolver(
+            Nt=self.Nt, backend=self.backend, i0=self.i0,
+            burn_in=self.burn_in, thin=self.thin
+        )