diff --git a/p_kit/backends/probana_backend.py b/p_kit/backends/probana_backend.py new file mode 100644 index 0000000..66b4533 --- /dev/null +++ b/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 + H + J + ANNEAL CONSTANT + ANNEAL LINEAR + COMMIT + RUN + +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() \ No newline at end of file diff --git a/p_kit/solver/real_device_solver.py b/p_kit/solver/real_device_solver.py new file mode 100644 index 0000000..199c5dc --- /dev/null +++ b/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 + ) \ No newline at end of file