From 8fc753367bf5f1a94d7153a99c94baf0ceaf813d Mon Sep 17 00:00:00 2001 From: Anton Date: Thu, 3 Sep 2026 15:33:02 +0200 Subject: [PATCH 1/7] Added new correlation primitives. Added a new solver that can take advantage of these primitives. A new image classification demo. --- examples/image_classification.py | 268 +++++++++++++++++++++++ p_kit/psl/p_circuit.py | 97 ++++++++- p_kit/solver/corr_ann_solver.py | 352 +++++++++++++++++++++++++++++++ 3 files changed, 707 insertions(+), 10 deletions(-) create mode 100644 examples/image_classification.py create mode 100644 p_kit/solver/corr_ann_solver.py diff --git a/examples/image_classification.py b/examples/image_classification.py new file mode 100644 index 0000000..87c122a --- /dev/null +++ b/examples/image_classification.py @@ -0,0 +1,268 @@ +""" +Fashion-MNIST controlled-correlation p-bit demo. + +The image-classification task is expressed largely as a probabilistic circuit, +which stores the model through base h/J terms and named correlation components. +The solver controls how these components are activated over time. + +Image evidence is applied through h, while each class is represented by a learned +J/h correlation component that is gradually activated by CorrelationAnnealingSolver. + +The current model uses 720 p-bits. The implementation is intentionally simple +and currently unoptimized for speed. + +This is POC demo and classification accuracy can also be improved. +""" +from pathlib import Path +from urllib.request import urlretrieve +import gzip, struct +import numpy as np +from sklearn.cluster import MiniBatchKMeans + +from p_kit.psl import PCircuit +from p_kit.solver.corr_ann_solver import CorrelationAnnealingSolver, staged_ramp + +# ---------------------------------------------------------------------- +# Configuration +# ---------------------------------------------------------------------- + +SIZE = 28 +N_FEATURES = 10 +PATCH_H, PATCH_W = 6, 7 +STRIDE = 3 +FEATURE_PATCHES = 40000 +FEATURE_BETA = 2.0 + +TRAIN_PER_CLASS = 2000 +OFFSETS = [(0,1),(0,2),(0,3),(1,0),(2,0),(3,0), + (1,1),(2,2),(1,-1),(2,-2),(1,2),(2,1),(1,-2),(2,-1)] + +I0 = 2.0 +SAMPLES = 10 +NT = 60 +TEST_LIMIT = 1000 # max test images +SEED = 1234 + +rng = np.random.default_rng(SEED) + +# ---------------------------------------------------------------------- +# Fashion-MNIST +# ---------------------------------------------------------------------- + +DATA = Path("data/fashion") +BASE = "https://raw.githubusercontent.com/zalandoresearch/fashion-mnist/master/data/fashion/" +FILES = ["train-images-idx3-ubyte.gz", "train-labels-idx1-ubyte.gz", + "t10k-images-idx3-ubyte.gz", "t10k-labels-idx1-ubyte.gz"] + +DATA.mkdir(parents=True, exist_ok=True) +for name in FILES: + p = DATA / name + if not p.exists(): + print("Downloading", name) + urlretrieve(BASE + name, p) + +def load_images(path): + with gzip.open(path, "rb") as f: + _, n, r, c = struct.unpack(">IIII", f.read(16)) + return np.frombuffer(f.read(), np.uint8).reshape(n, r, c) / 255.0 + +def load_labels(path): + with gzip.open(path, "rb") as f: + _, n = struct.unpack(">II", f.read(8)) + return np.frombuffer(f.read(), np.uint8) + +x_train = load_images(DATA / FILES[0]) +y_train = load_labels(DATA / FILES[1]) +x_test = load_images(DATA / FILES[2]) +y_test = load_labels(DATA / FILES[3]) + +# ---------------------------------------------------------------------- +# Shared local categorical features +# ---------------------------------------------------------------------- + +def positions(size, patch): + p = list(range(0, size - patch + 1, STRIDE)) + if p[-1] != size - patch: + p.append(size - patch) + return p + +YS, XS = positions(SIZE, PATCH_H), positions(SIZE, PATCH_W) +GH, GW = len(YS), len(XS) +GRID_N = GH * GW +N_PBITS = GRID_N * N_FEATURES + +patches = [] +while len(patches) < FEATURE_PATCHES: + im = x_train[rng.integers(len(x_train))] + y, x = rng.choice(YS), rng.choice(XS) + p = im[y:y+PATCH_H, x:x+PATCH_W] + if p.mean() > .04: + patches.append(p.ravel()) + +km = MiniBatchKMeans(N_FEATURES - 1, batch_size=1024, n_init=5, + random_state=SEED).fit(patches) +templates = np.vstack([np.zeros(PATCH_H * PATCH_W), km.cluster_centers_]) + +def feature_prob(images): + images = np.asarray(images) + if images.ndim == 2: + images = images[None] + P = np.stack([images[:, y:y+PATCH_H, x:x+PATCH_W].reshape(len(images), -1) + for y in YS for x in XS], axis=1) + d = ((P[:,:,None] - templates[None,None]) ** 2).mean(-1) + z = (d - d.min(2, keepdims=True)) / (d.std(2, keepdims=True) + 1e-6) + z = -FEATURE_BETA * z + z -= z.max(2, keepdims=True) + e = np.exp(z) + return e / e.sum(2, keepdims=True) + +def image_state(image): + p = feature_prob(image)[0] + u = np.log(np.clip(p, 1e-6, 1)) + u -= u.mean(1, keepdims=True) + q = np.zeros_like(p) + q[np.arange(GRID_N), p.argmax(1)] = 1 + return (u / 2).ravel(), (2 * q - 1).ravel() + +# ---------------------------------------------------------------------- +# Learn class-dependent feature correlations +# ---------------------------------------------------------------------- + +def shifted(F, dy, dx): + ya, yb = (slice(0,-dy), slice(dy,None)) if dy else (slice(None), slice(None)) + if dx > 0: + xa, xb = slice(0,-dx), slice(dx,None) + elif dx < 0: + xa, xb = slice(-dx,None), slice(0,dx) + else: + xa = xb = slice(None) + return F[:,ya,xa], F[:,yb,xb] + +priors = np.zeros((10, N_FEATURES)) +joints = np.zeros((10, len(OFFSETS), N_FEATURES, N_FEATURES)) +acount = np.zeros((10, len(OFFSETS), N_FEATURES)) + +print("Learning class correlations...") + +for c in range(10): + ids = rng.choice(np.flatnonzero(y_train == c), TRAIN_PER_CLASS, replace=False) + F = feature_prob(x_train[ids]).reshape(-1, GH, GW, N_FEATURES) + priors[c] = F.sum((0,1,2)) + + for o, (dy, dx) in enumerate(OFFSETS): + a, b = shifted(F, dy, dx) + a, b = a.reshape(-1, N_FEATURES), b.reshape(-1, N_FEATURES) + acount[c,o] = a.sum(0) + joints[c,o] = a.T @ b + +alpha = 1.0 +prior_p = (priors + alpha/N_FEATURES) / (priors.sum(1, keepdims=True) + alpha) +global_prior = priors.sum(0) / priors.sum() + +cond = (joints + alpha/N_FEATURES) / (acount[:,:,:,None] + alpha) +global_cond = (joints.sum(0) + alpha/N_FEATURES) / (acount.sum(0)[:,:,None] + alpha) + +prior_w = np.clip(np.log((prior_p + 1e-12) / (global_prior + 1e-12)), -2.5, 2.5) +rel_w = np.clip(np.log((cond + 1e-12) / (global_cond[None] + 1e-12)), -2.5, 2.5) + +# ---------------------------------------------------------------------- +# Convert categorical correlations to Ising J/h +# ---------------------------------------------------------------------- + +def block(y, x): + s = (y * GW + x) * N_FEATURES + return np.arange(s, s + N_FEATURES) + +def pairs(dy, dx): + for y in range(GH): + for x in range(GW): + yy, xx = y + dy, x + dx + if 0 <= yy < GH and 0 <= xx < GW: + yield y, x, yy, xx + +def build_model(c): + J = np.zeros((N_PBITS, N_PBITS)) + h = np.zeros(N_PBITS) + offset = 0.0 + U = np.tile(prior_w[c], (GRID_N, 1)) + + for o, (dy, dx) in enumerate(OFFSETS): + W = rel_w[c,o] + row, col, mean = W.mean(1), W.mean(0), W.mean() + W0 = W - row[:,None] - col[None,:] + mean + + for y, x, yy, xx in pairs(dy, dx): + a, b = block(y,x), block(yy,xx) + J[np.ix_(a,b)] += W0 / 4 + J[np.ix_(b,a)] += W0.T / 4 + U[y*GW+x] += row + U[yy*GW+xx] += col + offset += mean + + for s in range(GRID_N): + u = U[s] + mean = u.mean() + h[s*N_FEATURES:(s+1)*N_FEATURES] = (u - mean) / 2 + offset -= mean + + return J, h, offset + +models = [build_model(c) for c in range(10)] +radius = max(np.max(np.abs(np.linalg.eigvalsh(J))) for J, _, _ in models) +scale = 1.0 / radius + +circuits, solvers, offsets = [], [], [] + +for c, (J, h, offset) in enumerate(models): + circuit = PCircuit(N_PBITS) + circuit.set_correlation_component("class", J=scale*J, h=scale*h) + + solver = CorrelationAnnealingSolver( + Nt=NT, dt=.1667, i0=I0, seed=SEED+c, + block_size=N_FEATURES, + component_schedules={"class": staged_ramp(.2, .8)} + ) + + circuits.append(circuit) + solvers.append(solver) + offsets.append(scale * offset) + +# ---------------------------------------------------------------------- +# Controlled-correlation classification +# ---------------------------------------------------------------------- + +def classify(image): + h_image, initial = image_state(image) + E = np.zeros(10) + + for c in range(10): + circuits[c].h = h_image + + _, best_E = solvers[c].solve( + circuits[c], + n_shots=SAMPLES, + initial_state=initial, + return_best=True, + target_scales={"class": 1.0} + ) + + E[c] = best_E.mean() + offsets[c] + + return E.argmin() + + +# ---------------------------------------------------------------------- +# Main code +# ---------------------------------------------------------------------- + +print(f"Running controlled-correlation demo ({N_PBITS} p-bits)...") + +correct = 0 +n = min(TEST_LIMIT, len(x_test)) + +for i in range(n): + correct += classify(x_test[i]) == y_test[i] + if (i + 1) % 10 == 0: + print(f"\r{i+1}/{n} accuracy: {100*correct/(i+1):.2f}%", end="", flush=True) + +print(f"\nFinal accuracy: {100*correct/n:.2f}%") \ No newline at end of file diff --git a/p_kit/psl/p_circuit.py b/p_kit/psl/p_circuit.py index d191da8..9547d30 100644 --- a/p_kit/psl/p_circuit.py +++ b/p_kit/psl/p_circuit.py @@ -18,15 +18,15 @@ class PCircuit: biases J : np.array((n_pbits, n_pbits)) weights - ports: dict[str, object] - Circuit ports + ports : Dict[str, Any] | None + circuit ports """ def __init__(self, n_pbits: int, ports: Dict[str, Any] = None): self.n_pbits = n_pbits - self.ports = ports #Kept for copy behavior - + self.ports = ports + self.h = np.zeros((n_pbits,)) self.J = np.zeros((n_pbits, n_pbits)) self._connections = {} @@ -34,7 +34,8 @@ def __init__(self, n_pbits: int, ports: Dict[str, Any] = None): # identity (e.g. CaSuDaSolver's cache_J) can detect staleness # even though id(self.J) doesn't change across set_weight() calls. self._j_version = 0 - + self._correlation_components = {} + if ports: self._initialize_ports(ports) @@ -46,16 +47,14 @@ def _initialize_ports(self, port_attrs: Dict[str, Any]) -> None: for name, port in port_attrs.items(): port_indices[name] = idx idx += port.width - + # Set up each port for name, port in port_attrs.items(): new_port = Port(name=port.name, width=port.width) new_port.circuit = self new_port.index = port_indices[name] # Check ports name doesn't conflict with reserved attributes. - assert name != "ports" - assert name != "h" - assert name != "J" + assert name not in ("ports", "h", "J") setattr(self, name, new_port) def set_weight(self, from_pbit, to_pbit, weight, sym=True): @@ -64,8 +63,86 @@ def set_weight(self, from_pbit, to_pbit, weight, sym=True): self.J[to_pbit, from_pbit] = weight self._j_version += 1 + def add_correlation(self, i, j, weight=1.0, sym=True): + """Add positive coupling. i,j: p-bit indices; weight: strength; sym: symmetric coupling.""" + w = abs(weight) + self.J[i, j] += w + if sym: + self.J[j, i] += w + self._j_version += 1 + + def add_anticorrelation(self, i, j, weight=1.0, sym=True): + """Add negative coupling. i,j: p-bit indices; weight: strength; sym: symmetric coupling.""" + w = -abs(weight) + self.J[i, j] += w + if sym: + self.J[j, i] += w + self._j_version += 1 + + def add_group_coupling(self, group_a, group_b, weights=1.0, sym=True): + """Couple two groups. group_a/group_b: indices; weights: scalar or matrix; sym: symmetric coupling.""" + a = np.asarray(group_a, dtype=int) + b = np.asarray(group_b, dtype=int) + W = np.full((len(a), len(b)), float(weights)) if np.isscalar(weights) else np.asarray(weights, dtype=float) + if W.shape != (len(a), len(b)): + raise ValueError(f"weights shape {W.shape}, expected {(len(a), len(b))}") + self.J[np.ix_(a, b)] += W + if sym: + self.J[np.ix_(b, a)] += W.T + self._j_version += 1 + + def add_group_competition(self, group, strength=1.0): + """Add pairwise anticorrelation within a group. group: indices; strength: inhibition strength.""" + g = np.asarray(group, dtype=int) + w = abs(float(strength)) + for i in range(len(g)): + for j in range(i + 1, len(g)): + self.J[g[i], g[j]] -= w + self.J[g[j], g[i]] -= w + self._j_version += 1 + + def set_correlation_component(self, name, J=None, h=None): + """Set named dynamic component. name: identifier; J: coupling matrix; h: bias vector.""" + J = np.zeros_like(self.J) if J is None else np.asarray(J, dtype=float) + h = np.zeros_like(self.h) if h is None else np.asarray(h, dtype=float).reshape(-1) + if J.shape != self.J.shape: + raise ValueError(f"J shape {J.shape}, expected {self.J.shape}") + if h.shape != self.h.shape: + raise ValueError(f"h shape {h.shape}, expected {self.h.shape}") + self._correlation_components[name] = {"J": J.copy(), "h": h.copy()} + + def get_correlation_component(self, name): + """Return named component. name: component identifier.""" + return self._correlation_components[name] + + def remove_correlation_component(self, name): + """Remove named component. name: component identifier.""" + del self._correlation_components[name] + + def clear_correlation_components(self): + self._correlation_components.clear() + + def correlation_components(self): + return tuple(self._correlation_components) + + def effective_parameters(self, component_scales=None): + """Return effective J,h. component_scales: {name: scale} for active components.""" + J = self.J.copy() + h = self.h.copy() + if component_scales: + for name, scale in component_scales.items(): + c = self._correlation_components[name] + J += scale * c["J"] + h += scale * c["h"] + return J, h + def copy(self): new_circuit = PCircuit(self.n_pbits, self.ports) new_circuit.J = self.J new_circuit.h = self.h - return new_circuit + new_circuit._j_version = self._j_version + new_circuit._correlation_components = { + name: {"J": c["J"].copy(), "h": c["h"].copy()} + for name, c in self._correlation_components.items() + } + return new_circuit \ No newline at end of file diff --git a/p_kit/solver/corr_ann_solver.py b/p_kit/solver/corr_ann_solver.py new file mode 100644 index 0000000..4fa3836 --- /dev/null +++ b/p_kit/solver/corr_ann_solver.py @@ -0,0 +1,352 @@ +import numpy as np + +from p_kit.solver.annealing import constant +from p_kit.solver.base_solver import Solver + + +def hold(value=1.0): + """Return a component schedule that holds one constant scale.""" + value = float(value) + + def schedule(_solver, _run): + return value + + return schedule + + +def linear_ramp(start_value=0.0, end_value=1.0, + start_fraction=0.0, end_fraction=1.0): + """Ramp a correlation component over a fraction of the solver run.""" + start_value = float(start_value) + end_value = float(end_value) + start_fraction = float(start_fraction) + end_fraction = float(end_fraction) + + if not 0.0 <= start_fraction <= 1.0: + raise ValueError("start_fraction must be in [0, 1]") + if not 0.0 <= end_fraction <= 1.0: + raise ValueError("end_fraction must be in [0, 1]") + if end_fraction <= start_fraction: + raise ValueError("end_fraction must be greater than start_fraction") + + def schedule(solver, run): + if solver.Nt <= 1: + x = 1.0 + else: + x = run / (solver.Nt - 1) + + if x <= start_fraction: + return start_value + if x >= end_fraction: + return end_value + + t = (x - start_fraction) / (end_fraction - start_fraction) + return start_value + t * (end_value - start_value) + + return schedule + + +def staged_ramp(start_fraction=0.20, end_fraction=0.80): + """0 -> 1 schedule: evidence-only, ramp correlations, full relaxation.""" + return linear_ramp( + 0.0, 1.0, + start_fraction=start_fraction, + end_fraction=end_fraction, + ) + + +class CorrelationAnnealingSolver(Solver): + """Gibbs solver with dynamically controlled named J/h components. + + A ``PCircuit`` may contain a static base ``J``/``h`` plus named + correlation components created with ``set_correlation_component()`` or + the correlation primitives on ``PCircuit``. Each component receives an + independent time-dependent scale during a solve: + + J(t) = J_base + sum_r alpha_r(t) J_r + h(t) = h_base + sum_r alpha_r(t) h_r + + The same component scale is applied to its J and h parts, so a Potts/Ising + model can be annealed without breaking its exact J/h conversion. + + With ``block_size=None`` the solver performs ordinary sequential binary + Gibbs p-bit updates. With ``block_size=K`` it performs exact categorical + block updates: each contiguous K-p-bit group always contains exactly one + +1 state. The latter is useful for one-hot feature variables while still + allowing the inter-block correlations to be controlled through J. + + This first implementation deliberately uses the NumPy backend only. It is + intended to establish the primitive and algorithm cleanly before adding + optimized Torch/CuPy kernels. + """ + + def __init__(self, Nt, dt, i0, expected_mean=0, seed=None, backend=None, + tau=0.1, component_schedules=None, block_size=None): + super().__init__(Nt, dt, i0, expected_mean, seed, backend, tau) + + if self.backend.xp is not np: + raise NotImplementedError( + "CorrelationAnnealingSolver currently supports NumpyBackend only" + ) + + if block_size is not None and int(block_size) < 2: + raise ValueError("block_size must be >= 2") + + self.block_size = None if block_size is None else int(block_size) + self.component_schedules = dict(component_schedules or {}) + + def _components(self, c): + return { + name: c.get_correlation_component(name) + for name in c.correlation_components() + } + + def _schedule_value(self, name, run): + schedule = self.component_schedules.get(name, 1.0) + if callable(schedule): + return float(schedule(self, run)) + return float(schedule) + + def _compose(self, c, run): + J = np.asarray(c.J, dtype=float).copy() + h = np.asarray(c.h, dtype=float).reshape(-1).copy() + scales = {} + + for name, comp in self._components(c).items(): + scale = self._schedule_value(name, run) + scales[name] = scale + J += scale * np.asarray(comp["J"], dtype=float) + h += scale * np.asarray(comp["h"], dtype=float).reshape(-1) + + return J, h, scales + + def _compose_target(self, c, target_scales): + components = self._components(c) + scales = {} if target_scales is None else dict(target_scales) + + J = np.asarray(c.J, dtype=float).copy() + h = np.asarray(c.h, dtype=float).reshape(-1).copy() + + for name, comp in components.items(): + scale = float(scales.get(name, 1.0)) + J += scale * np.asarray(comp["J"], dtype=float) + h += scale * np.asarray(comp["h"], dtype=float).reshape(-1) + + return J, h + + @staticmethod + def _energy(m, J, h): + return ( + -0.5 * np.einsum("bi,ij,bj->b", m, J, m) + - m @ h + ) + + def _validate_binary_initial(self, initial_state, n_shots, n_pbits): + if initial_state is None: + return np.where( + self.random((n_shots, n_pbits)) < 0.5, + -1.0, + 1.0, + ) + + state = np.asarray(initial_state, dtype=float) + if state.ndim == 1: + if state.size != n_pbits: + raise ValueError(f"initial_state must contain {n_pbits} values") + state = np.tile(state.reshape(1, -1), (n_shots, 1)) + elif state.shape != (n_shots, n_pbits): + raise ValueError( + f"initial_state shape {state.shape}, expected {(n_shots, n_pbits)}" + ) + + if not np.all((state == -1) | (state == 1)): + raise ValueError("initial_state must contain only -1 or +1") + return state.copy() + + def _validate_block_initial(self, initial_state, n_shots, n_pbits): + K = self.block_size + if n_pbits % K: + raise ValueError( + f"n_pbits ({n_pbits}) must be divisible by block_size ({K})" + ) + + n_blocks = n_pbits // K + + if initial_state is None: + m = -np.ones((n_shots, n_pbits), dtype=float) + winners = self._generator.integers( + 0, K, size=(n_shots, n_blocks) + ) + for s in range(n_shots): + for blk in range(n_blocks): + m[s, blk * K + winners[s, blk]] = 1.0 + return m + + m = self._validate_binary_initial(initial_state, n_shots, n_pbits) + blocks = m.reshape(n_shots, n_blocks, K) + active = (blocks > 0).sum(axis=2) + if not np.all(active == 1): + raise ValueError( + "block initial_state must contain exactly one +1 per block" + ) + return m + + def _validate_J_components(self, c, tol=1e-12): + matrices = [("base", np.asarray(c.J, dtype=float))] + matrices.extend( + (name, np.asarray(comp["J"], dtype=float)) + for name, comp in self._components(c).items() + ) + + for name, J in matrices: + if not np.allclose(J, J.T, atol=tol, rtol=0): + raise ValueError( + f"component '{name}' must be symmetric for Gibbs energy sampling" + ) + + if self.block_size is None: + if np.max(np.abs(np.diag(J))) > tol: + raise ValueError( + f"component '{name}' must have zero diagonal for binary Gibbs updates" + ) + continue + + K = self.block_size + n = c.n_pbits + for start in range(0, n, K): + block = J[start:start + K, start:start + K] + if np.max(np.abs(block)) > tol: + raise ValueError( + f"component '{name}' has intra-block J couplings; " + "exact categorical block updates require zero " + "coupling within each block" + ) + + def _binary_sweep(self, m, J, h, beta): + n_shots, n_pbits = m.shape + field = m @ J + h + + for i in self._generator.permutation(n_pbits): + i = int(i) + logits = 2.0 * beta * field[:, i] + logits = np.clip(logits, -60.0, 60.0) + p_plus = 1.0 / (1.0 + np.exp(-logits)) + new = np.where(self.random((n_shots,)) < p_plus, 1.0, -1.0) + delta = new - m[:, i] + changed = delta != 0 + if np.any(changed): + m[:, i] = new + field += delta[:, None] * J[i, :][None, :] + + return m + + def _block_sweep(self, m, J, h, beta): + K = self.block_size + n_shots, n_pbits = m.shape + n_blocks = n_pbits // K + field = m @ J + h + + for blk in self._generator.permutation(n_blocks): + blk = int(blk) + start = blk * K + idx = slice(start, start + K) + + # For one +1 and K-1 -1 states, changing the winner changes + # the Ising energy by -2*field_k up to a block-independent + # constant. Hence the exact categorical conditional is + # softmax(2*beta*field_k). + logits = 2.0 * beta * field[:, idx] + logits -= logits.max(axis=1, keepdims=True) + p = np.exp(logits) + p /= p.sum(axis=1, keepdims=True) + + J_block = J[start:start + K, :] + + for s in range(n_shots): + k_new = self._generator.choice(K, p=p[s]) + old = m[s, idx].copy() + if old[k_new] > 0: + continue + + new = -np.ones(K, dtype=float) + new[k_new] = 1.0 + delta = new - old + m[s, idx] = new + field[s] += delta @ J_block + + return m + + def solve(self, c, annealing_func=constant, n_shots=1, + initial_state=None, return_final=False, return_best=False, + target_scales=None): + """Run controlled-correlation Gibbs annealing. + + ``component_schedules`` control how each named J/h component is + introduced. ``annealing_func`` still controls the ordinary p-kit + inverse-temperature/current scale (``constant`` means ``i0``). + + ``return_best=True`` tracks, for every shot, the lowest energy state + seen under the FULL target model. The initial state is included, so + the stochastic refinement never loses the best state it started from. + """ + if return_final and return_best: + raise ValueError("return_final and return_best are mutually exclusive") + + n_pbits = c.n_pbits + self._validate_J_components(c) + if self.block_size is None: + m = self._validate_binary_initial(initial_state, n_shots, n_pbits) + else: + m = self._validate_block_initial(initial_state, n_shots, n_pbits) + + target_J, target_h = self._compose_target(c, target_scales) + best_m = m.copy() + best_E = self._energy(m, target_J, target_h) + + if not return_final and not return_best: + all_m = np.zeros((self.Nt, n_shots, n_pbits), dtype=float) + all_E = np.zeros((self.Nt, n_shots), dtype=float) + all_scales = [] + + for run in range(self.Nt): + J, h, scales = self._compose(c, run) + beta = float(annealing_func(self, run)) + + if self.block_size is None: + m = self._binary_sweep(m, J, h, beta) + else: + m = self._block_sweep(m, J, h, beta) + + target_E = self._energy(m, target_J, target_h) + improved = target_E < best_E + if np.any(improved): + best_E[improved] = target_E[improved] + best_m[improved] = m[improved] + + if not return_final and not return_best: + all_m[run] = m + all_E[run] = target_E + all_scales.append(scales) + + if return_best: + return best_m, best_E + if return_final: + return m[0] if n_shots == 1 else m + + if n_shots == 1: + return all_m[:, 0, :], all_E[:, 0], all_scales + return all_m, all_E, all_scales + + def copy(self): + return CorrelationAnnealingSolver( + Nt=self.Nt, + dt=self.dt, + i0=self.i0, + expected_mean=self.expected_mean, + seed=self.seed, + backend=self.backend, + tau=self.tau, + component_schedules=self.component_schedules, + block_size=self.block_size, + ) + From 5252748d267b338335e14916a29b77059b75740a Mon Sep 17 00:00:00 2001 From: Anton Date: Thu, 3 Sep 2026 16:22:52 +0200 Subject: [PATCH 2/7] moved to subfolder --- examples/{ => advanced}/image_classification.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename examples/{ => advanced}/image_classification.py (100%) diff --git a/examples/image_classification.py b/examples/advanced/image_classification.py similarity index 100% rename from examples/image_classification.py rename to examples/advanced/image_classification.py From 87874b4c95114fb49c2f92128077d720236433de Mon Sep 17 00:00:00 2001 From: Anton Date: Thu, 3 Sep 2026 17:58:44 +0200 Subject: [PATCH 3/7] Added cached to solver. --- examples/advanced/image_classification.py | 1 + p_kit/solver/corr_ann_solver.py | 407 ++++++++++------------ 2 files changed, 176 insertions(+), 232 deletions(-) diff --git a/examples/advanced/image_classification.py b/examples/advanced/image_classification.py index 87c122a..ca7202e 100644 --- a/examples/advanced/image_classification.py +++ b/examples/advanced/image_classification.py @@ -12,6 +12,7 @@ and currently unoptimized for speed. This is POC demo and classification accuracy can also be improved. +Inference is slower/heavier than training in general. """ from pathlib import Path from urllib.request import urlretrieve diff --git a/p_kit/solver/corr_ann_solver.py b/p_kit/solver/corr_ann_solver.py index 4fa3836..73a755a 100644 --- a/p_kit/solver/corr_ann_solver.py +++ b/p_kit/solver/corr_ann_solver.py @@ -7,21 +7,16 @@ def hold(value=1.0): """Return a component schedule that holds one constant scale.""" value = float(value) - def schedule(_solver, _run): return value - return schedule def linear_ramp(start_value=0.0, end_value=1.0, start_fraction=0.0, end_fraction=1.0): """Ramp a correlation component over a fraction of the solver run.""" - start_value = float(start_value) - end_value = float(end_value) - start_fraction = float(start_fraction) - end_fraction = float(end_fraction) - + start_value, end_value = float(start_value), float(end_value) + start_fraction, end_fraction = float(start_fraction), float(end_fraction) if not 0.0 <= start_fraction <= 1.0: raise ValueError("start_fraction must be in [0, 1]") if not 0.0 <= end_fraction <= 1.0: @@ -30,323 +25,271 @@ def linear_ramp(start_value=0.0, end_value=1.0, raise ValueError("end_fraction must be greater than start_fraction") def schedule(solver, run): - if solver.Nt <= 1: - x = 1.0 - else: - x = run / (solver.Nt - 1) - + x = 1.0 if solver.Nt <= 1 else run / (solver.Nt - 1) if x <= start_fraction: return start_value if x >= end_fraction: return end_value - t = (x - start_fraction) / (end_fraction - start_fraction) return start_value + t * (end_value - start_value) - return schedule def staged_ramp(start_fraction=0.20, end_fraction=0.80): """0 -> 1 schedule: evidence-only, ramp correlations, full relaxation.""" - return linear_ramp( - 0.0, 1.0, - start_fraction=start_fraction, - end_fraction=end_fraction, - ) + return linear_ramp(0.0, 1.0, start_fraction, end_fraction) class CorrelationAnnealingSolver(Solver): """Gibbs solver with dynamically controlled named J/h components. - A ``PCircuit`` may contain a static base ``J``/``h`` plus named - correlation components created with ``set_correlation_component()`` or - the correlation primitives on ``PCircuit``. Each component receives an - independent time-dependent scale during a solve: + J(t) = J_base + sum_r alpha_r(t) J_r + h(t) = h_base + sum_r alpha_r(t) h_r - J(t) = J_base + sum_r alpha_r(t) J_r - h(t) = h_base + sum_r alpha_r(t) h_r - - The same component scale is applied to its J and h parts, so a Potts/Ising - model can be annealed without breaking its exact J/h conversion. - - With ``block_size=None`` the solver performs ordinary sequential binary - Gibbs p-bit updates. With ``block_size=K`` it performs exact categorical - block updates: each contiguous K-p-bit group always contains exactly one - +1 state. The latter is useful for one-hot feature variables while still - allowing the inter-block correlations to be controlled through J. - - This first implementation deliberately uses the NumPy backend only. It is - intended to establish the primitive and algorithm cleanly before adding - optimized Torch/CuPy kernels. + block_size=None uses binary Gibbs updates; block_size=K uses exact + categorical K-p-bit updates. Validation and block-sparse J layouts are + cached across solves; component fields are updated incrementally. """ def __init__(self, Nt, dt, i0, expected_mean=0, seed=None, backend=None, tau=0.1, component_schedules=None, block_size=None): super().__init__(Nt, dt, i0, expected_mean, seed, backend, tau) - if self.backend.xp is not np: - raise NotImplementedError( - "CorrelationAnnealingSolver currently supports NumpyBackend only" - ) - + raise NotImplementedError("CorrelationAnnealingSolver currently supports NumpyBackend only") if block_size is not None and int(block_size) < 2: raise ValueError("block_size must be >= 2") - self.block_size = None if block_size is None else int(block_size) self.component_schedules = dict(component_schedules or {}) + self._structure_key = None + self._layouts = None def _components(self, c): - return { - name: c.get_correlation_component(name) - for name in c.correlation_components() - } + return {name: c.get_correlation_component(name) + for name in c.correlation_components()} def _schedule_value(self, name, run): schedule = self.component_schedules.get(name, 1.0) - if callable(schedule): - return float(schedule(self, run)) - return float(schedule) - - def _compose(self, c, run): - J = np.asarray(c.J, dtype=float).copy() - h = np.asarray(c.h, dtype=float).reshape(-1).copy() - scales = {} + return float(schedule(self, run) if callable(schedule) else schedule) - for name, comp in self._components(c).items(): - scale = self._schedule_value(name, run) - scales[name] = scale - J += scale * np.asarray(comp["J"], dtype=float) - h += scale * np.asarray(comp["h"], dtype=float).reshape(-1) + def _cache_key(self, c, components): + return (id(c), id(c.J), getattr(c, "_j_version", 0), self.block_size, + tuple((name, id(comp["J"])) for name, comp in components.items())) - return J, h, scales - - def _compose_target(self, c, target_scales): - components = self._components(c) - scales = {} if target_scales is None else dict(target_scales) - - J = np.asarray(c.J, dtype=float).copy() - h = np.asarray(c.h, dtype=float).reshape(-1).copy() - - for name, comp in components.items(): - scale = float(scales.get(name, 1.0)) - J += scale * np.asarray(comp["J"], dtype=float) - h += scale * np.asarray(comp["h"], dtype=float).reshape(-1) - - return J, h - - @staticmethod - def _energy(m, J, h): - return ( - -0.5 * np.einsum("bi,ij,bj->b", m, J, m) - - m @ h - ) + def _validate_J_components(self, c, components, tol=1e-12): + matrices = [("base", np.asarray(c.J, dtype=float))] + matrices += [(name, np.asarray(comp["J"], dtype=float)) + for name, comp in components.items()] + for name, J in matrices: + if not np.allclose(J, J.T, atol=tol, rtol=0): + raise ValueError(f"component '{name}' must be symmetric for Gibbs energy sampling") + if self.block_size is None: + if np.max(np.abs(np.diag(J))) > tol: + raise ValueError(f"component '{name}' must have zero diagonal for binary Gibbs updates") + else: + K = self.block_size + for start in range(0, c.n_pbits, K): + if np.max(np.abs(J[start:start+K, start:start+K])) > tol: + raise ValueError( + f"component '{name}' has intra-block J couplings; " + "exact categorical block updates require zero coupling within each block" + ) + + def _block_layout(self, J): + K, n = self.block_size, len(J) + nb = n // K + rows, count = [], 0 + for a in range(nb): + row = [] + for b in range(nb): + W = J[a*K:(a+1)*K, b*K:(b+1)*K] + if np.any(W): + row.append((b, W)) + count += 1 + rows.append(row) + return rows if count < nb * nb / 2 else None + + def _prepare(self, c, components): + key = self._cache_key(c, components) + if key == self._structure_key: + return self._layouts + self._validate_J_components(c, components) + if self.block_size is None: + layouts = None + else: + layouts = {"base": self._block_layout(np.asarray(c.J, dtype=float))} + layouts.update({name: self._block_layout(np.asarray(comp["J"], dtype=float)) + for name, comp in components.items()}) + self._structure_key, self._layouts = key, layouts + return layouts def _validate_binary_initial(self, initial_state, n_shots, n_pbits): if initial_state is None: - return np.where( - self.random((n_shots, n_pbits)) < 0.5, - -1.0, - 1.0, - ) - - state = np.asarray(initial_state, dtype=float) - if state.ndim == 1: - if state.size != n_pbits: + return np.where(self.random((n_shots, n_pbits)) < 0.5, -1.0, 1.0) + m = np.asarray(initial_state, dtype=float) + if m.ndim == 1: + if m.size != n_pbits: raise ValueError(f"initial_state must contain {n_pbits} values") - state = np.tile(state.reshape(1, -1), (n_shots, 1)) - elif state.shape != (n_shots, n_pbits): - raise ValueError( - f"initial_state shape {state.shape}, expected {(n_shots, n_pbits)}" - ) - - if not np.all((state == -1) | (state == 1)): + m = np.tile(m, (n_shots, 1)) + elif m.shape != (n_shots, n_pbits): + raise ValueError(f"initial_state shape {m.shape}, expected {(n_shots, n_pbits)}") + if not np.all((m == -1) | (m == 1)): raise ValueError("initial_state must contain only -1 or +1") - return state.copy() + return m.copy() def _validate_block_initial(self, initial_state, n_shots, n_pbits): K = self.block_size if n_pbits % K: - raise ValueError( - f"n_pbits ({n_pbits}) must be divisible by block_size ({K})" - ) - - n_blocks = n_pbits // K - + raise ValueError(f"n_pbits ({n_pbits}) must be divisible by block_size ({K})") + nb = n_pbits // K if initial_state is None: - m = -np.ones((n_shots, n_pbits), dtype=float) - winners = self._generator.integers( - 0, K, size=(n_shots, n_blocks) - ) - for s in range(n_shots): - for blk in range(n_blocks): - m[s, blk * K + winners[s, blk]] = 1.0 - return m - + m = -np.ones((n_shots, nb, K)) + winners = self._generator.integers(0, K, size=(n_shots, nb)) + m[np.arange(n_shots)[:, None], np.arange(nb), winners] = 1.0 + return m.reshape(n_shots, n_pbits) m = self._validate_binary_initial(initial_state, n_shots, n_pbits) - blocks = m.reshape(n_shots, n_blocks, K) - active = (blocks > 0).sum(axis=2) - if not np.all(active == 1): - raise ValueError( - "block initial_state must contain exactly one +1 per block" - ) + if not np.all((m.reshape(n_shots, nb, K) > 0).sum(axis=2) == 1): + raise ValueError("block initial_state must contain exactly one +1 per block") return m - def _validate_J_components(self, c, tol=1e-12): - matrices = [("base", np.asarray(c.J, dtype=float))] - matrices.extend( - (name, np.asarray(comp["J"], dtype=float)) - for name, comp in self._components(c).items() - ) - - for name, J in matrices: - if not np.allclose(J, J.T, atol=tol, rtol=0): - raise ValueError( - f"component '{name}' must be symmetric for Gibbs energy sampling" - ) - - if self.block_size is None: - if np.max(np.abs(np.diag(J))) > tol: - raise ValueError( - f"component '{name}' must have zero diagonal for binary Gibbs updates" - ) - continue - - K = self.block_size - n = c.n_pbits - for start in range(0, n, K): - block = J[start:start + K, start:start + K] - if np.max(np.abs(block)) > tol: - raise ValueError( - f"component '{name}' has intra-block J couplings; " - "exact categorical block updates require zero " - "coupling within each block" - ) - - def _binary_sweep(self, m, J, h, beta): - n_shots, n_pbits = m.shape - field = m @ J + h - - for i in self._generator.permutation(n_pbits): - i = int(i) - logits = 2.0 * beta * field[:, i] - logits = np.clip(logits, -60.0, 60.0) - p_plus = 1.0 / (1.0 + np.exp(-logits)) - new = np.where(self.random((n_shots,)) < p_plus, 1.0, -1.0) + @staticmethod + def _target_energy(m, fields, base_h, comp_h, scales): + field = fields["base"].copy() + h = base_h.copy() + for name, scale in scales.items(): + field += scale * fields[name] + h += scale * comp_h[name] + return -0.5 * np.sum(m * field, axis=1) - m @ h + + def _apply_block_delta(self, field, live, delta, blk, J, layout, scale=1.0): + K = self.block_size + if layout is None: + d = delta @ J[blk*K:(blk+1)*K, :] + field += d + live += scale * d + return + for dst, W in layout[blk]: + idx = slice(dst*K, (dst+1)*K) + d = delta @ W + field[:, idx] += d + live[:, idx] += scale * d + + def _binary_sweep(self, m, base_J, comp_J, base_h, comp_h, fields, scales, beta): + live = fields["base"] + base_h + for name, scale in scales.items(): + live += scale * (fields[name] + comp_h[name]) + + for i in self._generator.permutation(m.shape[1]): + logits = np.clip(2.0 * beta * live[:, i], -60.0, 60.0) + p = 1.0 / (1.0 + np.exp(-logits)) + new = np.where(self.random((len(m),)) < p, 1.0, -1.0) delta = new - m[:, i] - changed = delta != 0 - if np.any(changed): - m[:, i] = new - field += delta[:, None] * J[i, :][None, :] - + m[:, i] = new + + d = delta[:, None] * base_J[i][None, :] + fields["base"] += d + live += d + for name, scale in scales.items(): + d = delta[:, None] * comp_J[name][i][None, :] + fields[name] += d + live += scale * d return m - def _block_sweep(self, m, J, h, beta): - K = self.block_size - n_shots, n_pbits = m.shape - n_blocks = n_pbits // K - field = m @ J + h - - for blk in self._generator.permutation(n_blocks): - blk = int(blk) - start = blk * K - idx = slice(start, start + K) - - # For one +1 and K-1 -1 states, changing the winner changes - # the Ising energy by -2*field_k up to a block-independent - # constant. Hence the exact categorical conditional is - # softmax(2*beta*field_k). - logits = 2.0 * beta * field[:, idx] + def _block_sweep(self, m, base_J, comp_J, base_h, comp_h, + fields, scales, layouts, beta): + K, nb = self.block_size, m.shape[1] // self.block_size + live = fields["base"] + base_h + for name, scale in scales.items(): + live += scale * (fields[name] + comp_h[name]) + + for blk in self._generator.permutation(nb): + idx = slice(blk*K, (blk+1)*K) + logits = 2.0 * beta * live[:, idx] logits -= logits.max(axis=1, keepdims=True) p = np.exp(logits) p /= p.sum(axis=1, keepdims=True) - J_block = J[start:start + K, :] - - for s in range(n_shots): - k_new = self._generator.choice(K, p=p[s]) - old = m[s, idx].copy() - if old[k_new] > 0: - continue + cdf = np.cumsum(p, axis=1) + cdf[:, -1] = 1.0 + winner = (cdf < self.random((len(m), 1))).sum(axis=1) - new = -np.ones(K, dtype=float) - new[k_new] = 1.0 - delta = new - old - m[s, idx] = new - field[s] += delta @ J_block + new = -np.ones((len(m), K)) + new[np.arange(len(m)), winner] = 1.0 + delta = new - m[:, idx] + m[:, idx] = new + self._apply_block_delta(fields["base"], live, delta, blk, + base_J, layouts["base"]) + for name, scale in scales.items(): + self._apply_block_delta(fields[name], live, delta, blk, + comp_J[name], layouts[name], scale) return m def solve(self, c, annealing_func=constant, n_shots=1, initial_state=None, return_final=False, return_best=False, target_scales=None): - """Run controlled-correlation Gibbs annealing. - - ``component_schedules`` control how each named J/h component is - introduced. ``annealing_func`` still controls the ordinary p-kit - inverse-temperature/current scale (``constant`` means ``i0``). - - ``return_best=True`` tracks, for every shot, the lowest energy state - seen under the FULL target model. The initial state is included, so - the stochastic refinement never loses the best state it started from. - """ + """Run controlled-correlation Gibbs annealing.""" if return_final and return_best: raise ValueError("return_final and return_best are mutually exclusive") - n_pbits = c.n_pbits - self._validate_J_components(c) + components = self._components(c) + layouts = self._prepare(c, components) + base_J = np.asarray(c.J, dtype=float) + base_h = np.asarray(c.h, dtype=float).reshape(-1) + comp_J = {name: np.asarray(comp["J"], dtype=float) for name, comp in components.items()} + comp_h = {name: np.asarray(comp["h"], dtype=float).reshape(-1) + for name, comp in components.items()} + if self.block_size is None: - m = self._validate_binary_initial(initial_state, n_shots, n_pbits) + m = self._validate_binary_initial(initial_state, n_shots, c.n_pbits) else: - m = self._validate_block_initial(initial_state, n_shots, n_pbits) + m = self._validate_block_initial(initial_state, n_shots, c.n_pbits) + + fields = {"base": m @ base_J} + fields.update({name: m @ J for name, J in comp_J.items()}) + target = {name: float((target_scales or {}).get(name, 1.0)) + for name in components} - target_J, target_h = self._compose_target(c, target_scales) best_m = m.copy() - best_E = self._energy(m, target_J, target_h) + best_E = self._target_energy(m, fields, base_h, comp_h, target) if not return_final and not return_best: - all_m = np.zeros((self.Nt, n_shots, n_pbits), dtype=float) - all_E = np.zeros((self.Nt, n_shots), dtype=float) + all_m = np.zeros((self.Nt, n_shots, c.n_pbits)) + all_E = np.zeros((self.Nt, n_shots)) all_scales = [] for run in range(self.Nt): - J, h, scales = self._compose(c, run) + scales = {name: self._schedule_value(name, run) for name in components} beta = float(annealing_func(self, run)) if self.block_size is None: - m = self._binary_sweep(m, J, h, beta) + m = self._binary_sweep(m, base_J, comp_J, base_h, comp_h, + fields, scales, beta) else: - m = self._block_sweep(m, J, h, beta) + m = self._block_sweep(m, base_J, comp_J, base_h, comp_h, + fields, scales, layouts, beta) - target_E = self._energy(m, target_J, target_h) + target_E = self._target_energy(m, fields, base_h, comp_h, target) improved = target_E < best_E - if np.any(improved): - best_E[improved] = target_E[improved] - best_m[improved] = m[improved] + best_E[improved] = target_E[improved] + best_m[improved] = m[improved] if not return_final and not return_best: - all_m[run] = m - all_E[run] = target_E + all_m[run], all_E[run] = m, target_E all_scales.append(scales) if return_best: return best_m, best_E if return_final: return m[0] if n_shots == 1 else m - if n_shots == 1: return all_m[:, 0, :], all_E[:, 0], all_scales return all_m, all_E, all_scales def copy(self): return CorrelationAnnealingSolver( - Nt=self.Nt, - dt=self.dt, - i0=self.i0, - expected_mean=self.expected_mean, - seed=self.seed, - backend=self.backend, - tau=self.tau, + Nt=self.Nt, dt=self.dt, i0=self.i0, + expected_mean=self.expected_mean, seed=self.seed, + backend=self.backend, tau=self.tau, component_schedules=self.component_schedules, block_size=self.block_size, - ) - + ) \ No newline at end of file From a4e1323b380573f4894f24b5f3e1a743708534d9 Mon Sep 17 00:00:00 2001 From: Anton Date: Thu, 3 Sep 2026 18:49:12 +0200 Subject: [PATCH 4/7] parallel image classification. --- examples/advanced/image_classification.py | 46 +++++++++++++++-------- 1 file changed, 30 insertions(+), 16 deletions(-) diff --git a/examples/advanced/image_classification.py b/examples/advanced/image_classification.py index ca7202e..102183e 100644 --- a/examples/advanced/image_classification.py +++ b/examples/advanced/image_classification.py @@ -23,6 +23,11 @@ from p_kit.psl import PCircuit from p_kit.solver.corr_ann_solver import CorrelationAnnealingSolver, staged_ramp +from joblib import Parallel, delayed +import os +import time + +N_WORKERS = min(4, os.cpu_count() or 1) # ---------------------------------------------------------------------- # Configuration # ---------------------------------------------------------------------- @@ -41,7 +46,7 @@ I0 = 2.0 SAMPLES = 10 NT = 60 -TEST_LIMIT = 1000 # max test images +TEST_LIMIT = 200 # max test images SEED = 1234 rng = np.random.default_rng(SEED) @@ -238,32 +243,41 @@ def classify(image): for c in range(10): circuits[c].h = h_image - _, best_E = solvers[c].solve( - circuits[c], - n_shots=SAMPLES, - initial_state=initial, - return_best=True, - target_scales={"class": 1.0} + circuits[c], n_shots=SAMPLES, initial_state=initial, + return_best=True, target_scales={"class": 1.0} ) - E[c] = best_E.mean() + offsets[c] return E.argmin() +def classify_chunk(indices): + return [(i, classify(x_test[i])) for i in indices] + # ---------------------------------------------------------------------- # Main code # ---------------------------------------------------------------------- -print(f"Running controlled-correlation demo ({N_PBITS} p-bits)...") +if __name__ == "__main__": + t0 = time.perf_counter() + + print(f"Running controlled-correlation demo ({N_PBITS} p-bits, {N_WORKERS} workers)...") + + n = min(TEST_LIMIT, len(x_test)) + chunks = np.array_split(np.arange(n), N_WORKERS) + + results = Parallel(n_jobs=N_WORKERS, verbose=10)( + delayed(classify_chunk)(chunk) for chunk in chunks + ) -correct = 0 -n = min(TEST_LIMIT, len(x_test)) + pred = np.empty(n, dtype=int) + for chunk in results: + for i, p in chunk: + pred[i] = p -for i in range(n): - correct += classify(x_test[i]) == y_test[i] - if (i + 1) % 10 == 0: - print(f"\r{i+1}/{n} accuracy: {100*correct/(i+1):.2f}%", end="", flush=True) + elapsed = time.perf_counter() - t0 + correct = np.sum(pred == y_test[:n]) -print(f"\nFinal accuracy: {100*correct/n:.2f}%") \ No newline at end of file + print(f"Final accuracy: {100*correct/n:.2f}%") + print(f"Elapsed: {elapsed:.1f} s ({elapsed/60:.1f} min)") \ No newline at end of file From ef18985205d5979f4d3a6c0b1fbc3fffdca2ffe0 Mon Sep 17 00:00:00 2001 From: Anton Date: Fri, 4 Sep 2026 09:05:39 +0200 Subject: [PATCH 5/7] small update --- examples/advanced/image_classification.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/advanced/image_classification.py b/examples/advanced/image_classification.py index 102183e..bff7336 100644 --- a/examples/advanced/image_classification.py +++ b/examples/advanced/image_classification.py @@ -8,8 +8,8 @@ Image evidence is applied through h, while each class is represented by a learned J/h correlation component that is gradually activated by CorrelationAnnealingSolver. -The current model uses 720 p-bits. The implementation is intentionally simple -and currently unoptimized for speed. +The current model uses 720 p-bits. The implementation is intentionally simple. +Calculation speed has been optimized, but it can be further improved. This is POC demo and classification accuracy can also be improved. Inference is slower/heavier than training in general. From a7f28b31f2e4c16d613edad72e6bb829022faad2 Mon Sep 17 00:00:00 2001 From: Anton Date: Fri, 4 Sep 2026 09:10:38 +0200 Subject: [PATCH 6/7] added results --- examples/advanced/image_classification.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/examples/advanced/image_classification.py b/examples/advanced/image_classification.py index bff7336..e92f8d2 100644 --- a/examples/advanced/image_classification.py +++ b/examples/advanced/image_classification.py @@ -13,6 +13,9 @@ This is POC demo and classification accuracy can also be improved. Inference is slower/heavier than training in general. + +Results: + 04/09/2026 On 1000 test images accuracy is 55%, takes 370s on a laptop """ from pathlib import Path from urllib.request import urlretrieve From 228e4402430720c2d01ac8a54d2bd636699112fc Mon Sep 17 00:00:00 2001 From: Anton Date: Mon, 7 Sep 2026 10:08:14 +0200 Subject: [PATCH 7/7] Fixed two points: aligned solve() with the standard p-kit solver return API so it works correctly with annealing.execute() aligned the reported energy convention with the other p-kit solvers and updated the best-state comparison accordingly I updated the Fashion-MNIST demo scoring to preserve the same classification semantics with the new energy convention. --- examples/advanced/image_classification.py | 4 ++-- p_kit/solver/corr_ann_solver.py | 24 +++++++++++++---------- 2 files changed, 16 insertions(+), 12 deletions(-) diff --git a/examples/advanced/image_classification.py b/examples/advanced/image_classification.py index e92f8d2..13c627d 100644 --- a/examples/advanced/image_classification.py +++ b/examples/advanced/image_classification.py @@ -250,9 +250,9 @@ def classify(image): circuits[c], n_shots=SAMPLES, initial_state=initial, return_best=True, target_scales={"class": 1.0} ) - E[c] = best_E.mean() + offsets[c] + E[c] = best_E.mean() - I0 * offsets[c] - return E.argmin() + return E.argmax() def classify_chunk(indices): return [(i, classify(x_test[i])) for i in indices] diff --git a/p_kit/solver/corr_ann_solver.py b/p_kit/solver/corr_ann_solver.py index 73a755a..712809a 100644 --- a/p_kit/solver/corr_ann_solver.py +++ b/p_kit/solver/corr_ann_solver.py @@ -151,14 +151,13 @@ def _validate_block_initial(self, initial_state, n_shots, n_pbits): raise ValueError("block initial_state must contain exactly one +1 per block") return m - @staticmethod - def _target_energy(m, fields, base_h, comp_h, scales): + def _target_energy(self, m, fields, base_h, comp_h, scales): field = fields["base"].copy() h = base_h.copy() for name, scale in scales.items(): field += scale * fields[name] h += scale * comp_h[name] - return -0.5 * np.sum(m * field, axis=1) - m @ h + return self.i0 * (0.5 * np.sum(m * field, axis=1) + m @ h) def _apply_block_delta(self, field, live, delta, blk, J, layout, scale=1.0): K = self.block_size @@ -254,13 +253,19 @@ def solve(self, c, annealing_func=constant, n_shots=1, if not return_final and not return_best: all_m = np.zeros((self.Nt, n_shots, c.n_pbits)) - all_E = np.zeros((self.Nt, n_shots)) - all_scales = [] + all_I = np.zeros((self.Nt, c.n_pbits)) + E = np.zeros(self.Nt) for run in range(self.Nt): scales = {name: self._schedule_value(name, run) for name in components} beta = float(annealing_func(self, run)) + if not return_final and not return_best: + live = fields["base"] + base_h + for name, scale in scales.items(): + live += scale * (fields[name] + comp_h[name]) + all_I[run] = (beta * live)[0] + if self.block_size is None: m = self._binary_sweep(m, base_J, comp_J, base_h, comp_h, fields, scales, beta) @@ -269,21 +274,20 @@ def solve(self, c, annealing_func=constant, n_shots=1, fields, scales, layouts, beta) target_E = self._target_energy(m, fields, base_h, comp_h, target) - improved = target_E < best_E + improved = target_E > best_E best_E[improved] = target_E[improved] best_m[improved] = m[improved] if not return_final and not return_best: - all_m[run], all_E[run] = m, target_E - all_scales.append(scales) + all_m[run], E[run] = m, target_E[0] if return_best: return best_m, best_E if return_final: return m[0] if n_shots == 1 else m if n_shots == 1: - return all_m[:, 0, :], all_E[:, 0], all_scales - return all_m, all_E, all_scales + return all_I, all_m[:, 0, :], E + return all_m def copy(self): return CorrelationAnnealingSolver(