diff --git a/CHANGELOG.md b/CHANGELOG.md index 17d4f138..f1494c42 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,10 @@ ### Added +- add unified ContractionAlgebra interface (`set_contraction_algebra(...)` + boundary + encode/decode) with tropical (max-plus) and complex pair-algebra + reference applications. + - Add efficient `expectation_pss` method for `U1Circuit`. - Support MVP mode for timeevol methods. diff --git a/applications/bcomplex32_algebra.py b/applications/bcomplex32_algebra.py new file mode 100644 index 00000000..96c04356 --- /dev/null +++ b/applications/bcomplex32_algebra.py @@ -0,0 +1,173 @@ +"""complex pair-algebra — a reference APPLICATION of ContractionAlgebra. + +Pair repr: complex tensor split into PairTensor(re, im) of bf16. Contraction = 4 real +bf16 matmuls (4M). Activated via ``cons.set_contraction_algebra(ComplexPairAlgebra())`` or +the ``bcomplex32()`` CM. PairTensor keeps the pair axis off tn.Node (dodges the axis==edge wall). +""" + +from typing import Any, Dict, Iterator, List, Tuple +import contextlib + +import numpy as np + +import tensorcircuit.cons as cons +from tensorcircuit.contraction_algebra import ( + ContractionAlgebra, + PairTensor, + Representation, +) + +Tensor = Any +Backend = Any + + +def _bf16_dtype() -> Any: + import ml_dtypes + + return ml_dtypes.bfloat16 + + +def _complex_to_pair(be: Backend, t: Tensor) -> PairTensor: + """complex tensor -> PairTensor(re, im) of bf16.""" + bf = _bf16_dtype() + re = be.cast(be.real(t), bf) + im = be.cast(be.imag(t), bf) + return PairTensor(re, im) + + +def _pair_to_complex(be: Backend, pair: PairTensor) -> Tensor: + """PairTensor of bf16 -> complex tensor (recombine; no copy risk via cast).""" + re, im = pair.unpack() + re = be.cast(re, cons.rdtypestr) + im = be.cast(im, cons.rdtypestr) + return be.cast(re + 1j * im, cons.dtypestr) + + +def _pair_tensordot(be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor: + """Complex tensordot = 4 real bf16 tensordots (4M). Uses be.tensordot (never patched).""" + ar, ai = PairTensor.unpack_pair(a) + br, bi = PairTensor.unpack_pair(b) + cr = be.tensordot(ar, br, axes) - be.tensordot(ai, bi, axes) + ci = be.tensordot(ar, bi, axes) + be.tensordot(ai, br, axes) + return PairTensor.pack_result(be, cr, ci, not isinstance(a, PairTensor)) + + +# Strategy note: bf16 single-operand einsum manually decomposes into diagonal → +# sum → transpose (staying in bf16 end-to-end) because numpy's ``np.einsum`` +# rejects bfloat16 dtypes. This is different from the tropical algebra's +# ``_tropical_einsum`` single-operand path, which delegates to ``be.einsum`` +# for repeated-index resolution — the tropical backend (float64) has no such +# dtype restriction. Both produce equivalent results under their respective +# semirings, but the implementation strategy is dictated by dtype constraints +# rather than algebraic differences. +def _einsum_single_operand_half( + be: Backend, x: Tensor, lhs: str, out_subs: str +) -> Tensor: + """Apply a 1-operand einsum to one bf16 half (decomposed into diagonal + sum + + transpose), staying in bf16 end-to-end. Handles reductions, transposes, + diagonals, and traces without float32 upcast. + """ + x_subs = list(lhs) # mutable subscript list we update in place + + # Step 1 — diagonalise every repeated index. + while True: + dup = next((c for c in set(x_subs) if x_subs.count(c) > 1), None) + if dup is None: + break + pos = [i for i, c in enumerate(x_subs) if c == dup] + x = np.diagonal(x, axis1=pos[0], axis2=pos[-1]) + x_subs = [c for i, c in enumerate(x_subs) if i != pos[-1]] + [dup] + + # Step 2 — sum over indices NOT wanted in the output. + out_set = set(out_subs) + sum_indices = [c for c in x_subs if c not in out_set] + if sum_indices: + x = be.sum(x, axis=tuple(x_subs.index(c) for c in sum_indices)) + x_subs = [c for c in x_subs if c in out_set] + + # Step 3 — transpose remaining indices into the requested output order. + if x_subs != list(out_subs): + perm = tuple(x_subs.index(c) for c in out_subs) + x = be.transpose(x, perm) + + return x + + +def _pair_einsum(be: Backend, eq: str, *operands: Tensor) -> PairTensor: + """Complex einsum = 4 real bf16 einsums (4M for 2 operands, 2 for 1).""" + if len(operands) == 1: + a = operands[0] + if "->" not in eq: + return a + lhs, out_subs = eq.split("->") + ar, ai = PairTensor.unpack_pair(a) + return PairTensor( + _einsum_single_operand_half(be, ar, lhs, out_subs), + _einsum_single_operand_half(be, ai, lhs, out_subs), + ) + + a, b = operands + ar, ai = PairTensor.unpack_pair(a) + br, bi = PairTensor.unpack_pair(b) + + if "->" not in eq: + raise ValueError( + f"implicit-mode einsum {eq!r} not supported for bf16; use explicit '->'" + ) + lhs, out_subs = eq.split("->") + a_subs, b_subs = lhs.split(",") + + a_set: set[str] = set(a_subs) + b_set: set[str] = set(b_subs) + contracted = [c for c in a_subs if c in b_set] + + a_free = [c for c in a_subs if c not in b_set] + b_free = [c for c in b_subs if c not in a_set] + out_order = list(out_subs) + + def _contract(x: Tensor, y: Tensor) -> Tensor: + if contracted: + a_axes = [a_subs.index(c) for c in contracted] + b_axes = [b_subs.index(c) for c in contracted] + result = be.tensordot(x, y, axes=(a_axes, b_axes)) + else: + result = be.tensordot(x, y, axes=0) + free_order = a_free + b_free + if free_order != out_order: + perm = [free_order.index(c) for c in out_order] + result = be.transpose(result, perm) + return result + + cr = _contract(ar, br) - _contract(ai, bi) + ci = _contract(ar, bi) + _contract(ai, br) + return PairTensor.pack_result(be, cr, ci, not isinstance(a, PairTensor)) + + +class PairBf16Representation(Representation): + name = "pair_bf16" + + def encode(self, be: Backend, tensors: List[Tensor]) -> List[Tensor]: + return [_complex_to_pair(be, t) for t in tensors] + + def decode(self, be: Backend, tensor: Tensor) -> Tuple[Tensor, Dict[str, Tensor]]: + return _pair_to_complex(be, tensor), {} + + +class ComplexPairAlgebra(ContractionAlgebra): + name = "bcomplex32_pair" + representation = PairBf16Representation() + + def get_contractor_kwargs(self) -> Dict[str, Any]: + return {"prefer_einsum": True} + + def tensordot(self, be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor: + return _pair_tensordot(be, a, b, axes) + + def einsum(self, be: Backend, eq: str, *operands: Tensor) -> Tensor: + return _pair_einsum(be, eq, *operands) + + +@contextlib.contextmanager +def bcomplex32() -> Iterator[None]: + with cons.runtime_contraction_algebra(ComplexPairAlgebra()): + yield diff --git a/applications/tropical_algebra.py b/applications/tropical_algebra.py new file mode 100644 index 00000000..cd7efb4d --- /dev/null +++ b/applications/tropical_algebra.py @@ -0,0 +1,766 @@ +"""Tropical (max-plus) contraction algebra -- a reference APPLICATION of the generic +``ContractionAlgebra`` ABCs shipped in ``tensorcircuit.contraction_algebra``. + +The feature is ``tensorcircuit.contraction_algebra`` -- the ``ContractionAlgebra`` and +``Representation`` ABCs -- wired in-source into ``cons._base``, which routes any +non-standard algebra through ``cons._algebraic_base_contraction`` (no monkey-patching). +A custom algebra is activated via ``tc.set_contraction_algebra(alg)`` or by the +``tropical()`` / ``counting_tropical()`` context managers below. This file is one +complete algebra built on top of the ABCs -- an importable reference module, not a +runnable demo (for that, see ``examples/tropical_ising.py``). Kept under +``applications/`` as a reference application; promote it to its own package/location if +it needs active development, independent distribution, or outgrows ``applications/`` +conventions. + +Implements the three standard tropical-tensor-network outputs (energy / +configuration / degeneracy) from Liu, Wang, Zhang PRL 126, 090506 (2021) +(arXiv:2008.06888). Contract under a tropical algebra via:: + + from applications.tropical_algebra import tropical + with tropical(): + ... # max-plus contraction + +Sections (consolidated from the original package modules): + 1. max-plus primitives + MaxPlusAlgebra (was tensorcircuit/contraction_algebra/tropical.py) + 2. counting (energy, degeneracy) + split_energy_count (was .../counting.py) + 3. tracking + configuration recovery (was .../config.py) + 4. context managers: tropical(track=...) / counting_tropical() +""" + +import contextlib +import logging +from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple + +import numpy as np + +import tensorcircuit.cons as cons +from tensorcircuit.contraction_algebra import ( + ContractionAlgebra, + PairTensor, + Representation, +) + +Tensor = Any +Backend = Any +_EPS = 1e-9 + +logger = logging.getLogger(__name__) + + +# ===== Section 1: max-plus primitives + MaxPlusAlgebra ===== + + +def _pair_layout( + be: Backend, a: Tensor, b: Tensor, axes: Any +) -> "tuple[Tensor, Tensor, tuple[int, ...]]": + """Transpose+reshape ``a``, ``b`` into the broadcast pair layout. + + Returns ``(a3, b3, out_shape)`` where ``a3`` is shape ``(m, k, 1)``, ``b3`` is + ``(1, k, n)``, and ``out_shape = a_free_shape + b_free_shape``. Here ``m`` is the + product of ``a``'s free axes, ``n`` of ``b``'s free axes, and ``k`` of the + contracted axes. The native broadcast of ``a3 + b3`` yields the ``(m, k, n)`` + pair-sum tensor required by max-plus (and counting) pairwise contraction. + """ + ashape = tuple(int(x) for x in be.shape_tuple(a)) + bshape = tuple(int(x) for x in be.shape_tuple(b)) + if isinstance(axes, int): + a_axes = list(range(len(ashape) - axes, len(ashape))) + b_axes = list(range(0, axes)) + else: + a_axes, b_axes = list(axes[0]), list(axes[1]) + a_free = [i for i in range(len(ashape)) if i not in a_axes] + b_free = [i for i in range(len(bshape)) if i not in b_axes] + a_t = be.transpose(a, tuple(a_free + a_axes)) # contracted axes to the tail + b_t = be.transpose(b, tuple(b_axes + b_free)) # contracted axes to the front + a_fs = [ashape[i] for i in a_free] + b_fs = [bshape[i] for i in b_free] + m = int(np.prod(a_fs)) + k = int(np.prod([ashape[i] for i in a_axes])) + n = int(np.prod(b_fs)) + a3 = be.reshape(a_t, (m, k, 1)) + b3 = be.reshape(b_t, (1, k, n)) + return a3, b3, (tuple(a_fs) + tuple(b_fs)) + + +def _tropical_tensordot(be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor: + """max-plus tensordot: Y[i,j] = max_k (A[i,k] + B[k,j]).""" + a3, b3, out_shape = _pair_layout(be, a, b, axes) + s = a3 + b3 # native broadcast add -> (m, k, n) + red = be.max(s, axis=1) # max over contracted axis (tropical addition) + return be.reshape(red, out_shape) + + +def _expand_to_layout( + be: Backend, t: Tensor, idxs: Sequence[str], full: Sequence[str] +) -> Tensor: + """Reshape/transpose ``t`` (axes == ``idxs``) into the ``full`` index layout, + inserting size-1 axes for indices not in ``idxs`` (enables broadcasting).""" + shape = tuple(int(x) for x in be.shape_tuple(t)) + present = {c: shape[i] for i, c in enumerate(idxs)} + sub = [c for c in full if c in present] # present indices in full order + tt = be.transpose(t, tuple(idxs.index(c) for c in sub)) + newshape = tuple(present[c] if c in present else 1 for c in full) + return be.reshape(tt, newshape) + + +# Strategy note: single-operand tropical einsum uses ``be.einsum`` for +# repeated-index resolution (diagonal gather). This is safe because tropical +# tensors are float64, which numpy's einsum accepts natively. For bf16 +# (which numpy rejects), see ``_einsum_single_operand_half`` in +# ``applications/bcomplex32_algebra.py`` for the manual decomposition approach. +def _tropical_einsum(be: Backend, eq: str, *operands: Tensor) -> Tensor: + """max-plus einsum: product -> +, sum -> max. Handles 1- or 2-operand forms.""" + if len(operands) == 1: + a = operands[0] + in_str, _sep, out_str = eq.partition("->") + lhs = in_str.split(",")[0] + rhs = out_str + if len(set(lhs)) != len(lhs): + resolved = "".join(dict.fromkeys(lhs)) + a = be.einsum(lhs + "->" + resolved, a) + lhs = resolved + contract = [c for c in lhs if c not in rhs] + for ax in sorted([lhs.index(c) for c in contract], reverse=True): + a = be.max(a, axis=ax) + remaining = [c for c in lhs if c not in contract] + return be.transpose(a, tuple(remaining.index(c) for c in rhs)) + a, b = operands + lhs, rhs = eq.split("->") + ia_s, ib_s = lhs.split(",") + ia, ib = list(ia_s), list(ib_s) + if len(set(ia)) != len(ia): + resolved = "".join(dict.fromkeys(ia)) + a = be.einsum("".join(ia) + "->" + resolved, a) + ia = list(resolved) + if len(set(ib)) != len(ib): + resolved = "".join(dict.fromkeys(ib)) + b = be.einsum("".join(ib) + "->" + resolved, b) + ib = list(resolved) + all_idx = list(dict.fromkeys(ia + ib)) + out_idx = list(rhs) + contract = [c for c in all_idx if c not in out_idx] + s = _expand_to_layout(be, a, ia, all_idx) + _expand_to_layout(be, b, ib, all_idx) + for ax in sorted([all_idx.index(c) for c in contract], reverse=True): + s = be.max(s, axis=ax) + remaining = [c for c in all_idx if c not in contract] + return be.transpose(s, tuple(remaining.index(c) for c in out_idx)) + + +class MaxPlusAlgebra(ContractionAlgebra): + """Tropical (max, +) semiring: addition -> max, multiplication -> +.""" + + name = "maxplus" + # representation defaults to IdentityRepresentation (real tensors, no codec): + # leaves enter/exit the contraction unchanged; only the kernels differ. + + def tensordot(self, be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor: + return _tropical_tensordot(be, a, b, axes) + + def einsum(self, be: Backend, eq: str, *operands: Tensor) -> Tensor: + return _tropical_einsum(be, eq, *operands) + + +# ===== Section 2: counting (energy, degeneracy) ===== + + +def split_energy_count(stacked: Tensor) -> "tuple[Tensor, Tensor]": + """Split a stacked ``[..., 2]`` tensor into ``(energy, count)`` numpy arrays. + + The last axis is interpreted as ``[..., 0] = energy, [..., 1] = count``. + """ + arr = stacked if isinstance(stacked, np.ndarray) else np.asarray(stacked) + return arr[..., 0], arr[..., 1] + + +def _counting_tensordot(be: Backend, a: Tensor, b: Tensor, axes: Any) -> PairTensor: + """(max-plus energy, degeneracy count) pairwise contraction. + + ``a`` and ``b`` are ``PairTensor`` (energy, count) or legacy stacked tensors. + """ + ae, an = PairTensor.unpack_pair(a) + b_e, b_n = PairTensor.unpack_pair(b) + a3x, b3x, out_shape = _pair_layout(be, ae, b_e, axes) + s = a3x + b3x # (m, k, n) energy pair sums + y = be.max(s, axis=1) # (m, n) max energy per output slot + m_n = tuple(int(d) for d in be.shape_tuple(y)) + y_b = be.reshape(y, (m_n[0], 1, m_n[1])) + mask = be.abs(s - y_b) < _EPS + a3n, b3n, _ = _pair_layout(be, an, b_n, axes) + pn = a3n * b3n # (m, k, n) pairwise count products + cy = be.reshape(be.sum(pn * mask, axis=1), out_shape) + y_shaped = be.reshape(y, out_shape) + return PairTensor.pack_result(be, y_shaped, cy, not isinstance(a, PairTensor)) + + +def _insert_axis_of( + be: Backend, reduced: Tensor, axis: int, full_shape: Sequence[int] +) -> Tensor: + """Reshape an axis-reduced tensor so its ``axis`` is size 1, restoring the + pre-reduce rank for broadcasting against the original tensor. + + ``be.max`` lacks ``keepdims``; this reshape re-inserts the size-1 axis so + ``reduced`` broadcasts against a tensor of shape ``full_shape``. + """ + new_shape = tuple( + 1 if i == axis else int(full_shape[i]) for i in range(len(full_shape)) + ) + return be.reshape(reduced, new_shape) + + +def _resolve_repeats( + be: Backend, pair: Tensor, idxs: Sequence[str] +) -> "tuple[Tensor, Tensor, list[str]]": + """Resolve intra-operand repeated indices via per-stream diagonal gather. + + Pure per-stream indexing -- the energy and count streams are gathered + independently via ``.unpack()``; the trailing pair axis stays invisible to + ``be.einsum``. The max/tie degeneracy logic lives entirely in the later + reduction; this helper only rewrites each stream's index layout. + + Returns ``(energy, count, resolved_idxs)`` where ``resolved_idxs`` has no + repeats. If ``idxs`` has no repeats, the streams are returned unchanged + and ``idxs`` is returned as-is. + """ + if len(set(idxs)) != len(idxs): + resolved = "".join(dict.fromkeys(idxs)) + e, n = PairTensor.unpack_pair(pair) + e = be.einsum("".join(idxs) + "->" + resolved, e) + n = be.einsum("".join(idxs) + "->" + resolved, n) + return e, n, list(resolved) + return PairTensor.unpack_pair(pair) + (list(idxs),) + + +def _counting_einsum(be: Backend, eq: str, *operands: Tensor) -> Tensor: + """(energy, degeneracy) einsum over max-plus. + + Handles the 1-operand and 2-operand forms that cotengra can emit. + + Mirrors ``_tropical_einsum``: each operand is broadcast to the full index + layout via ``_expand_to_layout``; the energy stream sums then takes the max + over each contracted axis; the count stream sums ``a_count * b_count`` only + at positions within ``_EPS`` of the running energy max (ties contribute). + + Intra-operand repeated indices (diagonal/trace) are resolved per-stream via + ``_resolve_repeats`` (pure diagonal gather -- no max/tie logic, which lives + in the reduction). This is the per-stream analogue of ``_tropical_einsum``'s + ``be.einsum`` gather, split across the two trailing-axis streams so the + ``[..., 2]`` stack axis is preserved rather than treated as an index. + """ + if len(operands) == 1: + a = operands[0] + in_str, _sep, out_str = eq.partition("->") + lhs = in_str.split(",")[0] + rhs = out_str + e, n, idxs = _resolve_repeats(be, a, list(lhs)) + contract = [c for c in idxs if c not in rhs] + for ax in sorted([idxs.index(c) for c in contract], reverse=True): + shp = tuple(int(d) for d in be.shape_tuple(e)) + y_ax = be.max(e, axis=ax) + y_b = _insert_axis_of(be, y_ax, ax, shp) + mask = be.abs(e - y_b) < _EPS + e = y_ax + n = be.sum(n * mask, axis=ax) + remaining = [c for c in idxs if c not in contract] + out_e = be.transpose(e, tuple(remaining.index(c) for c in rhs)) + out_n = be.transpose(n, tuple(remaining.index(c) for c in rhs)) + return PairTensor.pack_result( + be, out_e, out_n, not isinstance(operands[0], PairTensor) + ) + a, b = operands + lhs, rhs = eq.split("->") + ia_s, ib_s = lhs.split(",") + a_e, a_n, ia = _resolve_repeats(be, a, list(ia_s)) + b_e, b_n, ib = _resolve_repeats(be, b, list(ib_s)) + all_idx = list(dict.fromkeys(ia + ib)) + out_idx = list(rhs) + contract = [c for c in all_idx if c not in out_idx] + # Energy pair-sum and count product broadcast over the full index layout. + sx = _expand_to_layout(be, a_e, ia, all_idx) + _expand_to_layout( + be, b_e, ib, all_idx + ) + pn = _expand_to_layout(be, a_n, ia, all_idx) * _expand_to_layout( + be, b_n, ib, all_idx + ) + # Reduce over contracted axes (highest index first so earlier indices stay valid). + for ax in sorted([all_idx.index(c) for c in contract], reverse=True): + sx_shape = tuple(int(d) for d in be.shape_tuple(sx)) + y_ax = be.max(sx, axis=ax) # axis removed + y_b = _insert_axis_of(be, y_ax, ax, sx_shape) # axis size 1 -> broadcasts + mask = be.abs(sx - y_b) < _EPS + sx = y_ax + pn = be.sum(pn * mask, axis=ax) + remaining = [c for c in all_idx if c not in contract] + out_x = be.transpose(sx, tuple(remaining.index(c) for c in out_idx)) + out_n = be.transpose(pn, tuple(remaining.index(c) for c in out_idx)) + return PairTensor.pack_result(be, out_x, out_n, not isinstance(a, PairTensor)) + + +class CountingRepresentation(Representation): + """Attach count=1 (the counting-semiring multiplicative identity) to each + leaf via ``PairTensor(t, ones)``. Decode unpacks into energy (primary) + + count (aux, stashed in ``_last_aux`` for ``degeneracy()``). + """ + + name = "counting" + + def encode(self, be: Backend, tensors: List[Tensor]) -> List[Tensor]: + out = [] + for t in tensors: + ones = be.ones_like(t, dtype="float64") + out.append(PairTensor(t, ones)) + return out + + def decode(self, be: Backend, tensor: Tensor) -> Tuple[Tensor, Dict[str, Tensor]]: + e, n = PairTensor.unpack_pair(tensor) + self._last_aux = {"count": n} + return e, {} + + +class CountingTropicalAlgebra(ContractionAlgebra): + """Counting tropical algebra: (energy, degeneracy) over max-plus. + + Carries ``CountingRepresentation``: encode attaches count=1 to each leaf; + decode splits the final pair into energy (primary) + count (aux, stashed + in ``CountingRepresentation._last_aux`` for ``degeneracy()`` to read). + """ + + name = "counting_maxplus" + representation = CountingRepresentation() + + def tensordot(self, be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor: + return _counting_tensordot(be, a, b, axes) + + def einsum(self, be: Backend, eq: str, *operands: Tensor) -> Tensor: + return _counting_einsum(be, eq, *operands) + + +def degeneracy() -> Optional[Any]: + """Degeneracy (number of optimal configs) of the most recent counting + contraction. Call inside ``counting_tropical()`` after a contraction has run. + Returns ``None`` if no counting algebra is active.""" + alg = cons.get_contraction_algebra() + if alg is None: + return None + rep = alg.representation + if hasattr(rep, "_last_aux"): + return rep._last_aux.get("count") + return None + + +# ===== Section 3: tracking + configuration recovery ===== + +# Per-call argmax records, appended in call-order (== tree.traverse() order). +# Reset by on_contraction_start (fired by cons._algebraic_base_contraction) at +# the start of each tracked contraction so it always reflects exactly the most +# recent contraction. +_trace: List[Dict[str, Any]] = [] + +# Context stashed by the tracking hooks (fired by cons._algebraic_base_contraction): +# the cotengra tree plus the extracted topology (raw leaf arrays + index terms). +_ctx: Dict[str, Any] = { + "tree": None, + "input_sets": None, + "raw_tensors": None, +} + + +def _reset_trace() -> None: + """Clear the argmax trace (called by on_contraction_start, fired by cons._algebraic_base_contraction). + + Also drops any stashed tree so a failed/short-circuited contraction cannot + be confused with the previous one by ``recover_configuration``. + """ + _trace.clear() + _ctx["tree"] = None + + +def _set_tracking_context( + tree: Any, + input_sets: Optional[Sequence[Sequence[Any]]] = None, + raw_tensors: Optional[Sequence[Any]] = None, +) -> None: + """Stash the cotengra tree (and optionally the leaf topology) for backtracking.""" + _ctx["tree"] = tree + _ctx["input_sets"] = input_sets + _ctx["raw_tensors"] = raw_tensors + + +def get_recorded_topology() -> Tuple[Any, Any, Any]: + """Return ``(tree, input_sets, raw_tensors)`` stashed by the last tracked + contraction. The test harness uses this to verify a recovered config's + energy without needing to know cotengra's internal symbol->spin mapping.""" + return _ctx["tree"], _ctx["input_sets"], _ctx["raw_tensors"] + + +def _axes_to_lists(ashape: Sequence[int], axes: Any) -> Tuple[List[int], List[int]]: + """Normalize a tensordot ``axes`` arg to ``(a_axes, b_axes)`` (mirror of + ``_pair_layout``).""" + if isinstance(axes, int): + a_axes = list(range(len(ashape) - axes, len(ashape))) + b_axes = list(range(0, axes)) + else: + a_axes, b_axes = list(axes[0]), list(axes[1]) + return a_axes, b_axes + + +def _tracking_tensordot(be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor: + """max-plus tensordot (identical result to ``_tropical_tensordot``) that + additionally records the argmax over the contracted axis per output pos. + + Record shape: ``argmax`` has the pairwise output shape in the algebra's + natural order (``a_free + b_free``); values are the *flattened* contracted + index in row-major over ``a_axes`` (so ``np.unravel_index`` with + ``contract_dims`` recovers per-label values). ``contract_dims`` follows the + order of ``a_axes`` (== ``tree.get_tensordot_axes(p)[0]``). + + Argmax uses ``be.argmax`` (portable across numpy/jax/tensorflow); the + result is materialised to numpy for the backtracking index walk. + """ + a3, b3, out_shape = _pair_layout(be, a, b, axes) # (m,k,1),(1,k,n) + s = a3 + b3 # (m,k,n) + red = be.max(s, axis=1) # (m,n) max-plus reduction + out = be.reshape(red, out_shape) + + ashape = tuple(int(x) for x in be.shape_tuple(a)) + a_axes, _b_axes = _axes_to_lists(ashape, axes) + contract_dims = tuple(ashape[ax] for ax in a_axes) + + am = be.argmax(s, axis=1) # (m,n) -> flattened contracted index + am = be.reshape(am, out_shape) + _trace.append( + { + "kind": "td", + "argmax": np.asarray(am), + "contract_dims": contract_dims, + } + ) + return out + + +def _tracking_einsum(be: Backend, eq: str, *operands: Tensor) -> Tensor: + """max-plus einsum (identical result to ``_tropical_einsum``) that records + the argmax over the contracted index subspace per output position. + + For 2-operand forms the pair-sum is built over the full index layout (as in + ``_tropical_einsum``), transposed to ``(out_labels + contract_labels)``, + the contracted axes flattened into one trailing axis, and ``be.argmax`` + taken over it -- vectorised and portable (numpy/jax/tensorflow). Single- + operand forms record a no-op entry (no pairwise choice to backtrack). + """ + if len(operands) == 1: + _trace.append( + { + "kind": "ein1", + "argmax": None, + "contract_dims": (), + "contract_labels": (), + "eq": eq, + } + ) + else: + a, b = operands + lhs, rhs = eq.split("->") + ia_s, ib_s = lhs.split(",") + ia, ib = list(ia_s), list(ib_s) + out_labels = list(rhs) + all_idx = list(dict.fromkeys(ia + ib)) + contract_labels = [c for c in all_idx if c not in out_labels] + + # Full-layout pair sum (same construction as _tropical_einsum). + s = _expand_to_layout(be, a, ia, all_idx) + _expand_to_layout( + be, b, ib, all_idx + ) + + if not contract_labels: + # Pure outer product (hyperedge with no contraction): no argmax to take. + _trace.append( + { + "kind": "ein2", + "argmax": None, + "contract_dims": (), + "contract_labels": (), + "eq": eq, + } + ) + else: + # Reorder to (out_labels + contract_labels) and flatten the contracted axes. + perm = tuple(all_idx.index(c) for c in out_labels + contract_labels) + s = be.transpose(s, perm) + full_shape = tuple(int(d) for d in be.shape_tuple(s)) + n_out = len(out_labels) + out_shape = full_shape[:n_out] + contract_dims = full_shape[n_out:] + flat = int(np.prod(contract_dims)) + s_flat = be.reshape(s, out_shape + (flat,)) + am = be.argmax( + s_flat, axis=-1 + ) # shape == out_shape, row-major over contract + _trace.append( + { + "kind": "ein2", + "argmax": np.asarray(am), + "contract_dims": tuple(contract_dims), + "contract_labels": tuple(contract_labels), + "eq": eq, + } + ) + return _tropical_einsum(be, eq, *operands) + + +class MaxPlusTrackingAlgebra(MaxPlusAlgebra): + """Max-plus algebra that records the per-step argmax for config recovery. + + Produces the same contraction value as ``MaxPlusAlgebra``; the only + side-effect is appending an argmax record to ``_trace`` per pairwise call. + The trace is reset by ``on_contraction_start`` (fired by + ``cons._algebraic_base_contraction``) at the start of each tracked + contraction, so ``recover_configuration`` always reflects the most recent + contraction. + """ + + name = "maxplus_tracking" + + def tensordot(self, be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor: + return _tracking_tensordot(be, a, b, axes) + + def einsum(self, be: Backend, eq: str, *operands: Tensor) -> Tensor: + return _tracking_einsum(be, eq, *operands) + + def on_contraction_start(self, nodes: Any) -> None: + _reset_trace() + try: + _raw, _inputs, _output, _sizes = cons._extract_topology(nodes) + _set_tracking_context(tree=None, input_sets=_inputs, raw_tensors=_raw) + except Exception: + logger.debug( + "tracking topology stash failed; recover_configuration may raise " + "a trace/tree mismatch later", + exc_info=True, + ) + _set_tracking_context(tree=None) + + def on_contractor_ready(self, tree: Any) -> None: + _ctx["tree"] = tree + + +def _td_backtrack_step( + tree: Any, + p: Any, + l: Any, + r: Any, + rec: Dict[str, Any], + p_pos: Tuple[int, ...], + assignment: Dict[Any, int], +) -> Tuple[Tuple[int, ...], Tuple[int, ...]]: + """Propagate one tensordot contraction's optimal position to its children. + + Returns ``(l_pos, r_pos)`` (positions in ``get_inds(l)`` / ``get_inds(r)`` + order). Side-effect: records the contracted index-label values into + ``assignment``. + + ``get_tensordot_perm(p)`` is the transpose cotengra applies AFTER the + algebra call, mapping the algebra's natural output order (``l_free + + r_free``) to ``get_inds(p)``. It is therefore INVERTED when indexing the + recorded argmax (which lives in the algebra's natural order) by the + canonical (``get_inds(p)``) position: ``td_pos[perm[i]] = p_pos[i]``. + """ + l_inds = list(tree.get_inds(l)) + r_inds = list(tree.get_inds(r)) + l_axes, r_axes = tree.get_tensordot_axes(p) + perm = tree.get_tensordot_perm(p) + + # Convert canonical (get_inds(p)) position -> algebra-natural (td) order. + if perm is None: + td_pos = list(p_pos) + else: + td_pos = [0] * len(p_pos) + for i, v in enumerate(p_pos): + td_pos[perm[i]] = v + td_pos_t = tuple(td_pos) + + argmax = rec["argmax"] + if argmax is None: + # No contracted axis (should not happen for a tensordot step); nothing + # to unravel -- pass the position through unchanged. + return tuple(p_pos[: len(l_inds)]), tuple(p_pos[len(l_inds) :]) + + k = int(argmax[td_pos_t]) + contract_dims = rec["contract_dims"] + if contract_dims: + unraveled = np.unravel_index(k, contract_dims) # per contracted axis + else: + unraveled = () + + # Contracted labels follow l_axes order (== the algebra's a_axes order). + contract_labels = [l_inds[ax] for ax in l_axes] + for i, lbl in enumerate(contract_labels): + assignment[lbl] = int(unraveled[i]) + + # Split td_pos into the (l_free, r_free) halves (td-natural order). + l_free_labels = [c for i, c in enumerate(l_inds) if i not in set(l_axes)] + r_free_labels = [c for i, c in enumerate(r_inds) if i not in set(r_axes)] + n_lf = len(l_free_labels) + l_free_vals = td_pos[:n_lf] + r_free_vals = td_pos[n_lf:] + + l_pos = [0] * len(l_inds) + for j, lbl in enumerate(l_free_labels): + l_pos[l_inds.index(lbl)] = l_free_vals[j] + for i, ax in enumerate(l_axes): + l_pos[ax] = int(unraveled[i]) + + r_pos = [0] * len(r_inds) + for j, lbl in enumerate(r_free_labels): + r_pos[r_inds.index(lbl)] = r_free_vals[j] + for i, ax in enumerate(r_axes): + r_pos[ax] = int(unraveled[i]) # r_axes[i] pairs with l_axes[i] (same label) + + return tuple(l_pos), tuple(r_pos) + + +def _ein_backtrack_step( + p_pos: Tuple[int, ...], + rec: Dict[str, Any], +) -> Dict[Any, int]: + """Recover the contracted-label assignment for one einsum step (given the + output position). Used when ``tree.get_can_dot(p)`` is False.""" + argmax = rec["argmax"] + contract_dims = rec["contract_dims"] + if argmax is None or not contract_dims: + return {} + k = int(argmax[p_pos]) + unraveled = np.unravel_index(k, contract_dims) + return {lbl: int(v) for lbl, v in zip(rec["contract_labels"], unraveled)} + + +def _validate_trace(tree: Any) -> None: + """Validate that a stashed contraction tree exists and that ``_trace`` + length matches its number of contractions; raise RuntimeError otherwise. + + Called at the top of ``recover_configuration`` so the trace/tree invariant + holds before backtracking begins. + """ + if tree is None: + raise RuntimeError( + "recover_configuration: no stashed tree -- contract under " + "tropical(track=True) first." + ) + n_nodes = len(list(tree.traverse())) + if len(_trace) != n_nodes: + raise RuntimeError( + f"recover_configuration: trace length ({len(_trace)}) != tree " + f"contractions ({n_nodes}); call after exactly one tracked " + "contraction." + ) + + +def _finalize_from_leaves( + tree: Any, + opt_pos: Dict[Any, Tuple[int, ...]], + assignment: Dict[Any, int], +) -> None: + """Finalize ``assignment`` from leaf positions, catching labels that appear + only on a leaf axis (never seen on a contracted/tensordot step). + Mutates ``assignment`` in place. + """ + for leaf in tree.gen_leaves(): + pos = opt_pos.get(leaf) + if pos is None: + continue + for ax, lbl in enumerate(list(tree.get_inds(leaf))): + assignment[lbl] = int(pos[ax]) + + +def recover_configuration() -> Dict[Any, int]: + """Walk the stashed tree top-down and recover each index label's optimal + value. Returns ``{index_label: value}``. + + Requires that the contraction was performed under ``MaxPlusTrackingAlgebra`` + (i.e. ``tropical(track=True)``) so that ``_trace`` and ``_ctx['tree']`` are + populated, and called after exactly one tracked contraction (the trace is + reset per contraction by ``on_contraction_start``). + + Only scalar roots are supported (a full contraction to an energy, the + Ising use case). If the contraction has dangling/free output indices + (``len(tree.output) != 0``) a ``NotImplementedError`` is raised: the + non-scalar backtracking path has known ordering bugs (``tree.output`` order + vs result-shape order; output-label values lost in the einsum branch) that + would produce wrong configs, so it is gated rather than shipping wrong answers. + Contract to a scalar first. + + Tie-breaking is first-argument-wins (lowest flattened contracted index on + ties), so the returned configuration is *an* optimum; degenerate optima are + not enumerated (that is Task B's remit). + """ + tree = _ctx["tree"] + _validate_trace(tree) + + trav = list(tree.traverse()) # call-order == algebra call order + node_rec: Dict[Any, Dict[str, Any]] = {} + for (p, _l, _r), rec in zip(trav, _trace): + node_rec[p] = rec + + # Scalar-only: gate the non-scalar (free/output-index) root, whose + # backtracking wiring has known ordering bugs (tree.output order vs + # result-shape order; output-label values lost in the einsum branch). The + # tested scalar path below stays intact; ``get_tensordot_perm`` inversion + # logic (correct) is still exercised by the scalar canary + synthetic test. + n_out = len(tree.output) + if n_out != 0: + raise NotImplementedError( + "non-scalar configuration recovery not supported; " + "contract to a scalar first" + ) + root_pos: Tuple[int, ...] = () + + assignment: Dict[Any, int] = {} + opt_pos: Dict[Any, Tuple[int, ...]] = {} + opt_pos[tree.root] = root_pos + + for p, l, r in tree.descend(): + rec = node_rec[p] + p_pos = opt_pos[p] + if rec["kind"] == "td": + l_pos, r_pos = _td_backtrack_step(tree, p, l, r, rec, p_pos, assignment) + opt_pos[l] = l_pos + opt_pos[r] = r_pos + else: + # einsum (hyperedge): recover contracted labels; the child + # positions in canonical order are derived from the labels. + cont = _ein_backtrack_step(p_pos, rec) + assignment.update(cont) + # Reconstruct child positions from the label->value map so deeper + # tensordot steps (which read opt_pos) stay consistent. Each child's + # surviving labels are exactly its get_inds restricted to known vals. + l_inds = list(tree.get_inds(l)) + r_inds = list(tree.get_inds(r)) + opt_pos[l] = tuple(assignment[c] if c in assignment else 0 for c in l_inds) + opt_pos[r] = tuple(assignment[c] if c in assignment else 0 for c in r_inds) + + _finalize_from_leaves(tree, opt_pos, assignment) + return assignment + + +# ===== Section 4: context managers ===== + + +@contextlib.contextmanager +def tropical(track: bool = False) -> Iterator[None]: + """Contract under the max-plus (tropical) algebra within the block. + + ``track=True`` switches to ``MaxPlusTrackingAlgebra`` so + ``recover_configuration()`` can recover the optimal configuration after the + contraction. Off by default -> zero behaviour change relative to plain + max-plus. + """ + algebra = MaxPlusTrackingAlgebra() if track else MaxPlusAlgebra() + with cons.runtime_contraction_algebra(algebra): + yield + + +@contextlib.contextmanager +def counting_tropical() -> Iterator[None]: + """Contract under the counting (energy, degeneracy) max-plus algebra within + the block.""" + with cons.runtime_contraction_algebra(CountingTropicalAlgebra()): + yield diff --git a/examples/tropical_ising.py b/examples/tropical_ising.py new file mode 100644 index 00000000..7324cad8 --- /dev/null +++ b/examples/tropical_ising.py @@ -0,0 +1,63 @@ +"""Tropical (max-plus) Ising ground-state energy via TensorCircuit-NG contraction.""" + +import itertools +import numpy as np +import tensornetwork as tn +import tensorcircuit as tc +import tensorcircuit.cons as cons +from applications.tropical_algebra import tropical + + +def brute_force_ground_energy(n, edges, j_vals, h): + best = np.inf + for cfg in itertools.product([-1, 1], repeat=n): + e = -sum(jij * cfg[i] * cfg[j] for (i, j), jij in zip(edges, j_vals)) - sum( + hi * cfg[i] for i, hi in enumerate(h) + ) + best = min(best, e) + return best + + +def build_tn(n, edges, j_vals, h): + be = tc.backend + degree = [0] * n + for i, j in edges: + degree[i] += 1 + degree[j] += 1 + nodes, copy_nodes = [], {} + for i in range(n): + cn = tn.CopyNode(degree[i] + 1, 2) # rank (degree+1) delta, dim 2 -> hyperedge + copy_nodes[i] = cn + nodes.append(cn) + tv = np.array([h[i], -h[i]], dtype=np.float64) + tvn = tn.Node(be.cast(be.convert_to_tensor(tv), "float64")) + tn.connect(cn[0], tvn[0]) + nodes.append(tvn) + leg = [1] * n + for (i, j), jij in zip(edges, j_vals): + te = np.array([[jij, -jij], [-jij, jij]], dtype=np.float64) + ten = tn.Node(be.cast(be.convert_to_tensor(te), "float64")) + tn.connect(ten[0], copy_nodes[i][leg[i]]) + tn.connect(ten[1], copy_nodes[j][leg[j]]) + leg[i] += 1 + leg[j] += 1 + nodes.append(ten) + return nodes + + +def main(): + n = 5 + edges = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 0)] + J = [1.0, -1.0, 1.0, 1.0, -1.0] + h = [0.5, -0.3, 0.2, 0.0, 0.4] + nodes = build_tn(n, edges, J, h) + with tropical(): + val = float(np.array(cons.contractor(nodes, output_edge_order=[]).tensor)) + e_ground = brute_force_ground_energy(n, edges, J, h) + print(f"tropical contraction = {val:.6f}") + print(f"-E_ground (brute) = {-e_ground:.6f}") + print(f"match: {np.isclose(val, -e_ground)}") + + +if __name__ == "__main__": + main() diff --git a/tensorcircuit/cons.py b/tensorcircuit/cons.py index 6cb13b77..bc12507e 100644 --- a/tensorcircuit/cons.py +++ b/tensorcircuit/cons.py @@ -4,6 +4,7 @@ # pylint: disable=invalid-name +import functools import logging import sys import time @@ -21,6 +22,10 @@ from .backends.numpy_backend import NumpyBackend from .backends import get_backend from .simplify import _multi_remove +from .contraction_algebra import ( + ContractionAlgebra as _ContractionAlgebra, + IdentityRepresentation, +) logger = logging.getLogger(__name__) @@ -138,6 +143,19 @@ def set_tensornetwork_backend( set_tensornetwork_backend() +# --- ContractionAlgebra state (mirrors dtypestr/set_dtype pattern) --- + +_contraction_algebra: Optional[_ContractionAlgebra] = None + + +def get_contraction_algebra() -> Optional[_ContractionAlgebra]: + return _contraction_algebra + + +def set_contraction_algebra(alg: Optional[_ContractionAlgebra]) -> None: + global _contraction_algebra + _contraction_algebra = alg + def set_function_backend(backend: Optional[str] = None) -> Callable[..., Any]: """ @@ -297,6 +315,12 @@ def _sizen(node: tn.Node, is_log: bool = False) -> int: def _merge_single_gates( nodes: List[Any], total_size: Optional[int] = None ) -> Tuple[List[Any], int]: + if _contraction_algebra is not None: + # _merge_single_gates contracts via tn.contract_parallel (native sum-product), + # bypassing the algebra kernel — skip it under a non-standard algebra. + if total_size is None: + total_size = sum([_sizen(t) for t in nodes]) + return nodes, total_size # TODO(@refraction-ray): investigate whether too much copy here so that staging is slow for large circuit nodes = list(nodes) if total_size is None: @@ -697,48 +721,82 @@ def _wrap_omeco_optimizer(optimizer: Any) -> Any: return optimizer -def _algebraic_base_contraction( - nodes: List[tn.Node], +def _run_contraction( + raw_tensors: Any, + input_sets: Any, + output_set: Any, + size_dict: Any, algorithm: Any, - output_edge_order: Optional[Sequence[tn.Edge]] = None, - ignore_edge_order: bool = False, - **kws: Any, + ns: bool, + algebra: Any, + be: Any, + strip_exponent: bool, + ctg: Any, ) -> Any: + """Run the (possibly single-tensor) contraction, returning ``(final, exponent)``. + + Parameters: + raw_tensors: leaf tensors extracted from the tn.Node topology. + input_sets: cotengra input index sets (one per leaf). + output_set: cotengra output index set. + size_dict: cotengra size dictionary (index → dimension). + algorithm: cotengra path-finder (e.g. ``opt_einsum.paths.dynamic_programming``). + ns: ``True`` when a non-standard ``ContractionAlgebra`` is active. + algebra: the active ``ContractionAlgebra`` instance (carries kernels + hooks). + be: backend for all tensor operations (algebra kernels + tn.Node wrap). + strip_exponent: if ``True``, return the exponent separately via cotengra. + ctg: the ``cotengra`` module (imported once at the call site). """ - Execute contraction using cotengra and autoray for bare tensors. - """ - import cotengra as ctg - - raw_tensors, input_sets, output_set, size_dict = _extract_topology(nodes) - # Use the backend of the first node - be = nodes[0].backend - + exponent = 0.0 if len(raw_tensors) == 1: # Avoid cotengra bug for empty contraction paths - final_raw_tensor = be.einsum(input_sets[0] + "->" + output_set, *raw_tensors) - exponent = 0.0 + eq = input_sets[0] + "->" + output_set + final = ( + algebra.einsum(be, eq, *raw_tensors) if ns else be.einsum(eq, *raw_tensors) + ) else: path = algorithm(input_sets, output_set, size_dict) logger.info("the contraction path is given as %s" % str(path)) - tree = ctg.ContractionTree.from_path( input_sets, output_set, size_dict, path=path ) - - # Use autoray to keep AD and JIT support across backends - # Note: cotengra's make_contractor handles the orchestration - if not kws.get("strip_exponent", False): + if ns: + algebra.on_contractor_ready(tree) + impl = ( + functools.partial(algebra.einsum, be), + functools.partial(algebra.tensordot, be), + ) + contractor = ctg.core.make_contractor( + tree, implementation=impl, **algebra.get_contractor_kwargs() + ) + final = contractor(*raw_tensors) + elif not strip_exponent: + # autoray keeps AD and JIT support across backends contractor = ctg.core.make_contractor(tree, implementation="autoray") - final_raw_tensor = contractor(*raw_tensors) + final = contractor(*raw_tensors) else: - final_raw_tensor, exponent = tree.contract(raw_tensors, strip_exponent=True) - - final_node = tn.Node(final_raw_tensor, backend=be) + final, exponent = tree.contract(raw_tensors, strip_exponent=True) + return final, exponent + + +def _decode(ns: bool, rep: Any, be: Any, final: Any, output_set: Any) -> Any: + """Decode the contraction output under a non-standard algebra. + + ``rep.decode`` is responsible for stashing any aux internally.""" + if not ns: + return final + primary, _ = rep.decode(be, final) + if primary.ndim != len(output_set): + raise ValueError( + "representation.decode primary rank %d != len(output_set) %d; " + "decode must strip non-physical storage axes before tn.Node wraps it" + % (primary.ndim, len(output_set)) + ) + return primary - # Resolve dangling edges in the same order as in _extract_topology - dangling_edges = sorted_edges(tn.get_subgraph_dangling(nodes)) - # Update the edges to point to the new final_node +def _rewire_dangling_edges(final_node: Any, dangling_edges: Any, nodes: Any) -> None: + """Point every dangling edge at final_node, preserving topology order.""" for i, edge in enumerate(dangling_edges): if edge.node1 in nodes: edge.node1 = final_node @@ -748,12 +806,72 @@ def _algebraic_base_contraction( edge.axis2 = i final_node.edges = list(dangling_edges) + +def _algebraic_base_contraction( + nodes: List[tn.Node], + algorithm: Any, + output_edge_order: Optional[Sequence[tn.Edge]] = None, + ignore_edge_order: bool = False, + **kws: Any, +) -> Any: + """ + Execute contraction using cotengra and autoray for bare tensors. + """ + import cotengra as ctg + + raw_tensors, input_sets, output_set, size_dict = _extract_topology(nodes) + + alg = get_contraction_algebra() + + if alg is not None: + rep = alg.representation + else: + rep = IdentityRepresentation() # no-op; _decode skips it when ns=False + if alg is not None: + if kws.get("strip_exponent", False): + raise ValueError( + "strip_exponent is incompatible with a non-standard ContractionAlgebra" + ) + alg.on_contraction_start(nodes) + raw_tensors = rep.encode(backend, raw_tensors) + + final, exponent = _run_contraction( + raw_tensors, + input_sets, + output_set, + size_dict, + algorithm, + alg is not None, + algebra=alg, + be=backend, + strip_exponent=kws.get("strip_exponent", False), + ctg=ctg, + ) + + final = _decode(alg is not None, rep, backend, final, output_set) + + final_node = tn.Node(final, backend=backend) + + # Resolve dangling edges in the same order as in _extract_topology + dangling_edges = sorted_edges(tn.get_subgraph_dangling(nodes)) + + # Update the edges to point to the new final_node + _rewire_dangling_edges(final_node, dangling_edges, nodes) + if not ignore_edge_order: if output_edge_order is None: output_edge_order = dangling_edges final_node.reorder_edges(list(output_edge_order)) - if kws.get("strip_exponent", False): + # Apply the same output_edge_order permutation to aux (count co-indexed with energy) + if alg is not None and hasattr(rep, "_last_aux") and rep._last_aux: + order = output_edge_order if output_edge_order is not None else dangling_edges + perm = [dangling_edges.index(e) for e in order] + rep._last_aux = { + k: backend.transpose(v, tuple(perm)) for k, v in rep._last_aux.items() + } + + if kws.get("strip_exponent", False) and alg is None: return final_node, exponent return final_node @@ -902,10 +1020,11 @@ def _base( # 1. Resolve topology and check for hyperedges has_hyperedges = any(isinstance(n, tn.CopyNode) for n in nodes) + if _contraction_algebra is not None: + return _algebraic_base_contraction( + nodes, algorithm, output_edge_order, ignore_edge_order, **kws + ) if use_primitives is True or (use_primitives is None and has_hyperedges): - # ========================================== - # NEW ALGEBRAIC EXECUTION PATH (Opt-in) - # ========================================== return _algebraic_base_contraction( nodes, algorithm, output_edge_order, ignore_edge_order, **kws ) @@ -1148,6 +1267,11 @@ def set_contractor( To set runtime contractor of the tensornetwork for a better contraction path. For more information on the usage of contractor, please refer to independent tutorial. + To change the contraction algebra, use ``cons.set_contraction_algebra(alg)`` + separately (the algebra is orthogonal to the contractor configuration). The + ``tropical()`` / ``bcomplex32()`` / ``counting_tropical()`` context managers + are the recommended way to switch algebras for a block of code. + :param method: "auto", "greedy", "branch", "plain", "tng", "custom", "custom_stateful". Also supports shortcuts like "cotengra", "cotengra-30-64", "omeco", and "omeco-16-32". defaults to None ("auto") @@ -1318,6 +1442,21 @@ def runtime_contractor(*confargs: Any, **confkws: Any) -> Iterator[Any]: _set_global_contractor(old_contractor) +@contextmanager +def runtime_contraction_algebra(alg: _ContractionAlgebra) -> Iterator[None]: + """Context manager to temporarily set a non-standard contraction algebra. + + Mirrors ``runtime_backend`` / ``runtime_dtype``: saves the current algebra, + sets ``alg`` for the block, restores on exit. + """ + prev = get_contraction_algebra() + set_contraction_algebra(alg) + try: + yield + finally: + set_contraction_algebra(prev) + + def split_rules( max_singular_values: Optional[int] = None, max_truncation_err: Optional[float] = None, diff --git a/tensorcircuit/contraction_algebra.py b/tensorcircuit/contraction_algebra.py new file mode 100644 index 00000000..82ed2a2f --- /dev/null +++ b/tensorcircuit/contraction_algebra.py @@ -0,0 +1,156 @@ +"""ContractionAlgebra: a generic interface for swapping contraction primitives + +boundary representation, consulted by ``cons._algebraic_base_contraction``. + +Implement ``ContractionAlgebra`` (with a ``Representation``) and activate via +``cons.set_contraction_algebra(...)`` -- the single entry-point for switching +contraction algebras. The in-source ``cons._base`` +routes any non-standard algebra to ``_algebraic_base_contraction``, which runs +encode -> algebra kernels -> decode; no monkey-patching is required. + +Reference applications: ``applications/tropical_algebra.py`` +(max-plus / counting / tracking) and ``applications/bcomplex32_algebra.py`` +(bf16 pair). +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any, Dict, List, Tuple + +Tensor = Any +Backend = Any + + +class Representation(ABC): + """Boundary codec between logical elements and physical storage.""" + + name: str = "abstract" + + @abstractmethod + def encode(self, be: Backend, tensors: List[Tensor]) -> List[Tensor]: + """Transform the leaf raw_tensors (one pass, per-tensor, topology-agnostic).""" + + @abstractmethod + def decode(self, be: Backend, tensor: Tensor) -> Tuple[Tensor, Dict[str, Tensor]]: + """Transform the final contracted tensor into ``(primary, aux)``. + ``primary`` must have rank == len(output_set) so ``tn.Node`` wraps it + consistently. ``aux`` carries side outputs (e.g., degeneracy), sharing + the primary's physical axis order (sliced from the same tensor).""" + + +class IdentityRepresentation(Representation): + """No-op codec: standard storage (real/complex scalars), no aux.""" + + name = "identity" + + def encode(self, be: Backend, tensors: List[Tensor]) -> List[Tensor]: + return tensors + + def decode(self, be: Backend, tensor: Tensor) -> Tuple[Tensor, Dict[str, Tensor]]: + return tensor, {} + + +class ContractionAlgebra(ABC): + """How to perform the two atomic ops of a contraction, plus observation hooks. + + cotengra decomposes any contraction into pairwise steps; each is a + ``tensordot`` (ordinary pair) or ``einsum`` (hyperedge / copy-node). An + algebra supplies both. The kernel must be closed over the layout its + ``representation`` produces. + """ + + name: str = "abstract" + representation: Representation = IdentityRepresentation() + + def get_contractor_kwargs(self) -> Dict[str, Any]: + """Extra kwargs forwarded to cotengra's ``make_contractor``. + + Override to return ``{'prefer_einsum': True}`` when your algebra's + ``tensordot`` kernel carries non-physical storage axes (e.g. the + complex pair axis) that cotengra's post-tensordot autoray + transpose would mishandle (ValueError: axes don't match array). + ``prefer_einsum=True`` forces einsum-only execution, which skips + the transpose entirely. + + Default ``{}`` keeps the standard tensordot+einsum mix — required + by tropical config-recovery backtracking, which depends on the + tensordot intermediate layout. + """ + return {} + + @abstractmethod + def tensordot(self, be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor: + """Pairwise contraction over ``axes`` (np.tensordot convention).""" + + @abstractmethod + def einsum(self, be: Backend, eq: str, *operands: Tensor) -> Tensor: + """Pairwise einsum (cotengra may also call with 1 operand).""" + + def on_contraction_start(self, nodes: Any) -> None: + """Called before each non-standard contraction (default no-op).""" + + def on_contractor_ready(self, tree: Any) -> None: + """Called when the cotengra tree is built (default no-op).""" + + +class StandardAlgebra(ContractionAlgebra): + """The usual (sum, product) ring — identical to native backend behaviour.""" + + name = "standard" + + def tensordot(self, be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor: + return be.tensordot(a, b, axes) + + def einsum(self, be: Backend, eq: str, *operands: Tensor) -> Tensor: + return be.einsum(eq, *operands) + + +class PairTensor: + """Virtual tensor wrapping two halves (e.g., re/im, energy/count). + + Algebra kernels unpack via ``.unpack()`` and operate on the halves + directly, then return new ``PairTensor`` results. ``encode()`` and + ``decode()`` are the only places that create / consume the pair + representation. The pair axis is kept off ``tn.Node`` (cotengra sees only + the primary half's topology via the symbolic tree), so no ``.shape`` / + ``.ndim`` array-protocol is exposed. + """ + + __slots__ = ("_p", "_s") + + def __init__(self, primary: Tensor, secondary: Tensor): + self._p = primary + self._s = secondary + + def unpack(self) -> Tuple[Tensor, Tensor]: + return self._p, self._s + + @staticmethod + def unpack_pair(pair: "PairTensor | Tensor") -> Tuple[Tensor, Tensor]: + """Unpack a PairTensor or legacy stack [..., 2] — backward compatible.""" + if isinstance(pair, PairTensor): + return pair.unpack() + return pair[..., 0], pair[..., 1] + + @staticmethod + def pack_result( + be: Any, primary: Tensor, secondary: Tensor, legacy: bool + ) -> "PairTensor | Tensor": + """Return PairTensor or legacy stack matching input format.""" + # Kernels intentionally accept both PairTensor and legacy stacked + # [..., 2] operands and echo the input format on output, so direct + # callers and the counting tests that build be.stack([..., -1]) keep + # working. The contraction flow itself is PairTensor end-to-end (encode + # wraps every leaf); only direct kernel entry points see the legacy form. + if legacy: + return be.stack([primary, secondary], axis=-1) + return PairTensor(primary, secondary) + + +__all__ = [ + "ContractionAlgebra", + "StandardAlgebra", + "Representation", + "IdentityRepresentation", + "PairTensor", +] diff --git a/tests/test_bcomplex32_algebra.py b/tests/test_bcomplex32_algebra.py new file mode 100644 index 00000000..d262dc27 --- /dev/null +++ b/tests/test_bcomplex32_algebra.py @@ -0,0 +1,129 @@ +import numpy as np +from applications.bcomplex32_algebra import ( + _pair_tensordot, + _pair_einsum, + _complex_to_pair, + _pair_to_complex, +) +import tensorcircuit.backends.numpy_backend as nb + +be = nb.NumpyBackend() + + +def test_complex_to_pair_roundtrip_within_bf16_quant(): + t = np.array( + [[1.0 + 2.0j, 3.0 - 1.0j], [0.25 + 0.5j, -2.0 + 7.0j]], dtype=np.complex64 + ) + pair = _complex_to_pair(be, t) + back = _pair_to_complex(be, pair) + np.testing.assert_allclose(np.asarray(back), t, rtol=2e-2) + + +def test_pair_tensordot_matches_complex_tensordot(): + rng = np.random.default_rng(0) + a = rng.standard_normal((2, 3)).astype(np.complex64) + b = rng.standard_normal((3, 4)).astype(np.complex64) + pa = _complex_to_pair(be, a) + pb = _complex_to_pair(be, b) + out = _pair_tensordot(be, pa, pb, axes=([1], [0])) + ref = np.tensordot(a, b, axes=([1], [0])) + np.testing.assert_allclose( + np.asarray(_pair_to_complex(be, out)), ref, rtol=2e-2, atol=5e-3 + ) + + +def test_pair_einsum_single_operand(): + rng = np.random.default_rng(0) + a = rng.standard_normal((2, 2)).astype(np.complex64) + pa = _complex_to_pair(be, a) + out = _pair_einsum( + be, "ab->a", pa + ) # reduce over b (sum in complex -> here identity semantics) + ref = np.einsum("ab->a", a) + np.testing.assert_allclose( + np.asarray(_pair_to_complex(be, out)), ref, rtol=2e-2, atol=5e-3 + ) + + +# --- Task 12: end-to-end through real tc.Circuit + wall-avoidance canary --- + +import tensorcircuit as tc +from applications.bcomplex32_algebra import bcomplex32 + + +def test_bf16_end_to_end_matches_complex64(): + import numpy as np + + def build(): + c = tc.Circuit(3) + c.H(0) + c.cnot(0, 1) + c.cnot(1, 2) + return np.asarray(c.state()) + + ref = build() # default complex64 + with bcomplex32(): + got = build() # bf16 pair path + np.testing.assert_allclose(got, ref, rtol=2e-2) + + +def test_bf16_wall_avoidance_canary(): + import numpy as np + + with bcomplex32(): + c = tc.Circuit(4) + for i in range(4): + c.H(i) + for i in range(3): + c.cnot(i, i + 1) + st = np.asarray(c.state()) + assert st.shape == (16,) # ran cleanly, no axis==edge crash + import tensorcircuit.cons as cons + + assert cons.get_contraction_algebra() is None # CM restored, no leak + c2 = tc.Circuit(2) + c2.H(0) + c2.cnot(0, 1) # subsequent native contraction + assert np.asarray(c2.state()).shape == (4,) # algebra restored, no leak + + +def test_pair_einsum_keeps_bfloat16_dtype(): + """T1: _pair_einsum must compute in bf16, not upcast to float32. + + numpy's np.einsum rejects bf16; the old _pair_einsum worked around it by + upcasting to float32 (so its output pair was float32, not bf16). The rewrite + uses manual tensordot decomposition, which stays bf16. This test locks that. + """ + import ml_dtypes + + bf = ml_dtypes.bfloat16 + a = np.array([[1.0 + 2.0j, 3.0j], [-1.0j, 2.0 - 1.0j]], dtype=np.complex64) + b = np.array([[0.5 + 0.5j, 1.0j], [2.0j, -1.0 + 1.0j]], dtype=np.complex64) + pa = _complex_to_pair(be, a) + pb = _complex_to_pair(be, b) + result = _pair_einsum(be, "ij,jk->ik", pa, pb) + re, _ = result.unpack() + assert re.dtype == bf, f"_pair_einsum upcast to {re.dtype}; expected bfloat16" + + +def test_bf16_ghz8_runs_and_matches_native(): + """T3: the 8-qubit GHZ that previously crashed (cotengra autoray transpose on + a pair result) now runs under bcomplex32 and matches native within bf16 + tolerance. get_contractor_kwargs returns prefer_einsum=True, which avoids the transpose path; the genuine-bf16 + kernel keeps intermediates bf16 end-to-end. + """ + + def ghz(n): + c = tc.Circuit(n) + c.H(0) + for i in range(n - 1): + c.cnot(i, i + 1) + return np.asarray(c.state()) + + ref = ghz(8) + with bcomplex32(): + got = ghz(8) + assert got.shape == ref.shape + assert np.allclose( + got, ref, rtol=1.5e-2 + ), f"max abs diff = {np.abs(got - ref).max()}" diff --git a/tests/test_contraction_algebra.py b/tests/test_contraction_algebra.py new file mode 100644 index 00000000..ec4a19b1 --- /dev/null +++ b/tests/test_contraction_algebra.py @@ -0,0 +1,470 @@ +import numpy as np +import tensorcircuit as tc +import tensorcircuit.cons as cons +from tensorcircuit.contraction_algebra import ContractionAlgebra, StandardAlgebra + + +def test_standard_tensordot_matches_backend(): + be = tc.backend + a = be.cast( + be.convert_to_tensor(np.arange(12, dtype=np.float64).reshape(3, 4)), "float64" + ) + b = be.cast( + be.convert_to_tensor(np.arange(20, dtype=np.float64).reshape(4, 5)), "float64" + ) + alg = StandardAlgebra() + got = np.array(alg.tensordot(be, a, b, axes=1)) + ref = np.array(be.tensordot(a, b, axes=1)) + assert got.shape == (3, 5) + np.testing.assert_allclose(got, ref) + + +def test_standard_einsum_matches_backend(): + be = tc.backend + a = be.cast( + be.convert_to_tensor(np.arange(6, dtype=np.float64).reshape(2, 3)), "float64" + ) + b = be.cast( + be.convert_to_tensor(np.arange(12, dtype=np.float64).reshape(3, 4)), "float64" + ) + alg = StandardAlgebra() + got = np.array(alg.einsum(be, "ab,bc->ac", a, b)) + ref = np.array(be.einsum("ab,bc->ac", a, b)) + np.testing.assert_allclose(got, ref) + + +def test_standard_name(): + assert StandardAlgebra().name == "standard" + + +def test_public_api_surface(): + from tensorcircuit import contraction_algebra as tca + + # After Task 14 the package exports only the 4 base names; activation lives + # in-source via cons.set_contraction_algebra. + for name in [ + "ContractionAlgebra", + "StandardAlgebra", + "Representation", + "IdentityRepresentation", + ]: + assert hasattr(tca, name), name + # The old monkey-patch API names are intentionally gone: + for gone in [ + "activate", + "deactivate", + "standard", + "runtime_contraction_algebra", + "set_contraction_algebra", + "get_contraction_algebra", + "injection", + ]: + assert not hasattr(tca, gone), gone + + +def test_algebra_hooks_default_noop(): + alg = StandardAlgebra() + assert alg.on_contraction_start(["dummy_nodes"]) is None + assert alg.on_contractor_ready(["dummy_tree"]) is None + + +from tensorcircuit.contraction_algebra import ( + ContractionAlgebra, + StandardAlgebra, + Representation, + IdentityRepresentation, +) + +# --- Task 5: non-standard path (encode -> kernels via implementation= -> decode, +# hooks fire, primary.ndim == len(output_set), aux stashed) --- + + +def test_nonstandard_path_encodes_kernel_decodes_in_order(): + import numpy as np + import tensornetwork as tn + import opt_einsum + + log = [] + + class LogRep(Representation): + name = "log" + + def encode(self, be, tensors): + log.append("encode") + return tensors + + def decode(self, be, tensor): + log.append("decode") + return tensor, {} + + class LogAlg(ContractionAlgebra): + name = "log" + representation = LogRep() + + def tensordot(self, be, a, b, axes): + log.append("td") + return be.tensordot(a, b, axes) + + def einsum(self, be, eq, *ops): + log.append("ein") + return be.einsum(eq, *ops) + + def on_contraction_start(self, nodes): + log.append("start") + + def on_contractor_ready(self, tree): + log.append("ready") + + rng = np.random.default_rng(0) + a = tn.Node(rng.standard_normal((2, 3)).astype(np.complex64)) + b = tn.Node(rng.standard_normal((3, 4)).astype(np.complex64)) + tn.connect(a[1], b[0]) + prev = cons.get_contraction_algebra() + cons.set_contraction_algebra(LogAlg()) + try: + cons._algebraic_base_contraction( + [a, b], + algorithm=opt_einsum.paths.dynamic_programming, + output_edge_order=[a[0], b[1]], + ) + assert log[0] == "start" + assert "encode" in log and "ready" in log and "decode" in log + assert log.index("encode") < log.index("decode") + # hooks fire around encode/ready in the prescribed order: + assert log.index("start") < log.index("encode") + assert log.index("ready") < log.index("decode") + finally: + cons.set_contraction_algebra(prev) + + +def test_decode_wrong_rank_raises_loudly(): + import numpy as np + import tensornetwork as tn + import opt_einsum + import pytest + + class BadRep(Representation): + name = "bad" + + def encode(self, be, tensors): + return tensors + + def decode(self, be, tensor): + # Add a trailing storage axis that the representation forgot to strip. + return ( + be.reshape(tensor, tuple(be.shape_tuple(tensor)) + (1,)), + {}, + ) + + class BadAlg(ContractionAlgebra): + name = "bad" + representation = BadRep() + + def tensordot(self, be, a, b, axes): + return be.tensordot(a, b, axes) + + def einsum(self, be, eq, *ops): + return be.einsum(eq, *ops) + + # 2-node fully contracted scalar: output_set="", len 0. BadAlg.decode returns + # rank 1, so primary.ndim != len(output_set) trips the assert. + rng = np.random.default_rng(0) + a = tn.Node(rng.standard_normal(2).astype(np.complex64)) + b = tn.Node(rng.standard_normal(2).astype(np.complex64)) + tn.connect(a[0], b[0]) + prev = cons.get_contraction_algebra() + cons.set_contraction_algebra(BadAlg()) + try: + with pytest.raises(ValueError): + cons._algebraic_base_contraction( + [a, b], + algorithm=opt_einsum.paths.dynamic_programming, + output_edge_order=[], + ) + finally: + cons.set_contraction_algebra(prev) + + +def test_counting_representation_stashes_aux(): + import numpy as np + from applications.tropical_algebra import CountingRepresentation + + rep = CountingRepresentation() + t = np.array([[3.0, 7.0], [1.0, 4.0]]) + primary, _ = rep.decode(None, t) + assert np.array_equal(primary, np.array([3.0, 1.0])) + assert np.array_equal(rep._last_aux["count"], np.array([7.0, 4.0])) + + +def test_representation_identity_roundtrip(): + rep = IdentityRepresentation() + import numpy as np + + t = [np.ones((2, 3))] + assert rep.encode(None, t) is t # identity returns same + primary, aux = rep.decode(None, t[0]) + assert primary is t[0] and aux == {} + + +def test_standard_algebra_carries_identity_representation(): + alg = StandardAlgebra() + assert isinstance(alg.representation, IdentityRepresentation) + assert alg.name == "standard" + + +def test_contraction_algebra_representation_default_is_identity(): + # a minimal concrete algebra that does NOT override representation + class BareAlgebra(ContractionAlgebra): + name = "bare" + + def tensordot(self, be, a, b, axes): + return be.tensordot(a, b, axes) + + def einsum(self, be, eq, *ops): + return be.einsum(eq, *ops) + + assert isinstance(BareAlgebra().representation, IdentityRepresentation) + + +# --- Task 2: cons.py algebra state (set_contraction_algebra) --- + + +class _NS(ContractionAlgebra): # non-standard stub for constraint tests + name = "ns" + representation = None # noqa — not used in these tests + + def tensordot(self, be, a, b, axes): + return be.tensordot(a, b, axes) + + def einsum(self, be, eq, *ops): + return be.einsum(eq, *ops) + + +def test_set_contraction_algebra(): + prev = cons.get_contraction_algebra() + try: + cons.set_contraction_algebra(_NS()) + assert isinstance(cons.get_contraction_algebra(), _NS) + finally: + cons.set_contraction_algebra(prev) + + +def test_runtime_contractor_restores_algebra(): + prev = cons.get_contraction_algebra() + cons.set_contraction_algebra(_NS()) + try: + with cons.runtime_contractor("greedy"): + assert isinstance(cons.get_contraction_algebra(), _NS) + # runtime_contractor saves/restores algebra independently of the contractor, + # so algebra set before entering survives unchanged. + assert isinstance(cons.get_contraction_algebra(), _NS) + finally: + cons.set_contraction_algebra(prev) + + +def test_set_function_contractor_restores_algebra(): + prev = cons.get_contraction_algebra() + cons.set_contraction_algebra(_NS()) + try: + + @cons.set_function_contractor("greedy") + def f(): + return cons.get_contraction_algebra() + + inside = f() + assert isinstance(inside, _NS) # algebra active during the call + # after the call, algebra survives unchanged (it was not set by the decorator) + assert isinstance(cons.get_contraction_algebra(), _NS) + finally: + cons.set_contraction_algebra(prev) + + +# --- Task 3: cons.py guards (_merge_single_gates skip + _base routing) --- + + +def test_merge_single_gates_skipped_under_nonstandard_algebra(monkeypatch): + import numpy as np + import tensornetwork as tn + + prev = cons.get_contraction_algebra() + cons.set_contraction_algebra(_NS()) + try: + + def boom(*a, **k): + raise AssertionError( + "merge body ran — should be skipped under non-standard algebra" + ) + + monkeypatch.setattr(tn, "contract_parallel", boom) + node = tn.Node(np.zeros(2)) + out = cons._merge_single_gates([node], 7) + assert out[0] == [node] and out[1] == 7 + finally: + cons.set_contraction_algebra(prev) + + +def test_base_routes_to_algebraic_under_nonstandard(monkeypatch): + import numpy as np + import tensornetwork as tn + import opt_einsum + + rng = np.random.default_rng(0) + a = tn.Node(rng.standard_normal((2, 3)).astype(np.complex64)) + b = tn.Node(rng.standard_normal((3, 4)).astype(np.complex64)) + tn.connect(a[1], b[0]) + routed = {} + + def fake_alg(nodes, algorithm, *args, **kw): + routed["called"] = True + return "FINAL" + + monkeypatch.setattr(cons, "_algebraic_base_contraction", fake_alg) + prev = cons.get_contraction_algebra() + cons.set_contraction_algebra(_NS()) + try: + cons._base( + [a, b], + algorithm=opt_einsum.paths.dynamic_programming, + output_edge_order=[a[0], b[1]], + ) + assert routed.get("called") is True + finally: + cons.set_contraction_algebra(prev) + + +# --- Task 4: lock the StandardAlgebra backward-compat baseline --- +# A real 2-node contraction under the default StandardAlgebra must produce +# results bit-identical to a direct np.tensordot reference. This is the +# safety net before Task 5 adds the non-standard encode/decode path to +# _algebraic_base_contraction. If this test passes on today's code (which +# ignores the algebra), that is exactly the point: it locks the baseline. + + +def test_standard_algebra_contraction_matches_native(): + import numpy as np + import tensornetwork as tn + + rng = np.random.default_rng(0) + a = tn.Node(rng.standard_normal((2, 3)).astype(np.complex64)) + b = tn.Node(rng.standard_normal((3, 4)).astype(np.complex64)) + tn.connect(a[1], b[0]) + prev = cons.get_contraction_algebra() + assert prev is None + cons.set_contraction_algebra(StandardAlgebra()) + try: + import opt_einsum + + n = cons._algebraic_base_contraction( + [a, b], + algorithm=opt_einsum.paths.dynamic_programming, + output_edge_order=[a[0], b[1]], + ) + ref = np.tensordot(a.tensor, b.tensor, axes=([1], [0])) + np.testing.assert_allclose(n.tensor, ref, rtol=1e-6) + finally: + cons.set_contraction_algebra(prev) + + +def test_aux_stash_handles_ignore_edge_order_with_none_order(monkeypatch): + # A non-standard algebra returning aux under ignore_edge_order=True + output_edge_order=None + # must NOT crash (counting-scalar path). Uses the LogAlg/LogRep from test_nonstandard_path_*. + import numpy as np + import tensornetwork as tn + import opt_einsum + + class AuxRep(Representation): + name = "aux" + _last_aux = {} + + def encode(self, be, tensors): + return tensors + + def decode(self, be, tensor): + self._last_aux = {"count": np.ones_like(tensor)} + return tensor, {} + + class AuxAlg(ContractionAlgebra): + name = "aux" + representation = AuxRep() + + def tensordot(self, be, a, b, axes): + return be.tensordot(a, b, axes) + + def einsum(self, be, eq, *ops): + return be.einsum(eq, *ops) + + rng = np.random.default_rng(0) + a = tn.Node(rng.standard_normal(2).astype(np.complex64)) + b = tn.Node(rng.standard_normal(2).astype(np.complex64)) + tn.connect(a[0], b[0]) + prev = cons.get_contraction_algebra() + cons.set_contraction_algebra(AuxAlg()) + try: + # scalar contraction, ignore_edge_order=True, output_edge_order=None — must not raise + cons._algebraic_base_contraction( + [a, b], + algorithm=opt_einsum.paths.dynamic_programming, + output_edge_order=None, + ignore_edge_order=True, + ) + assert "count" in cons.get_contraction_algebra().representation._last_aux + finally: + cons.set_contraction_algebra(prev) + + +def test_get_contractor_kwargs_default(): + assert StandardAlgebra().get_contractor_kwargs() == {} + # Concrete algebra without override inherits the default + assert _NS().get_contractor_kwargs() == {} + + +# --- Coverage: single-tensor path + the strip_exponent x algebra guard. +# Both hit lines in cons._algebraic_base_contraction / _run_contraction that no +# other test exercises (len(raw_tensors)==1 cotengra empty-path workaround, and +# the strip_exponent incompatibility guard). + + +def test_single_tensor_contraction_uses_algebra_einsum(): + # A one-node contraction takes the len(raw_tensors)==1 branch, bypassing + # cotengra; it must still dispatch through algebra.einsum under a + # non-standard algebra (here StandardAlgebra), not just be.einsum. + import tensornetwork as tn + import opt_einsum + + t = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + a = tn.Node(t) + prev = cons.get_contraction_algebra() + cons.set_contraction_algebra(StandardAlgebra()) + try: + node = cons._algebraic_base_contraction( + [a], + algorithm=opt_einsum.paths.dynamic_programming, + output_edge_order=[a[0], a[1]], + ) + # "ab->ab" identity contraction: the result equals the leaf tensor. + np.testing.assert_allclose(np.array(node.tensor), t) + finally: + cons.set_contraction_algebra(prev) + + +def test_strip_exponent_incompatible_with_algebra_raises(): + # strip_exponent=True under any non-standard algebra is rejected up front. + import tensornetwork as tn + import opt_einsum + import pytest + + a = tn.Node(np.array([1.0, 2.0])) + b = tn.Node(np.array([3.0, 4.0])) + tn.connect(a[0], b[0]) + prev = cons.get_contraction_algebra() + cons.set_contraction_algebra(StandardAlgebra()) + try: + with pytest.raises(ValueError, match="strip_exponent is incompatible"): + cons._algebraic_base_contraction( + [a, b], + algorithm=opt_einsum.paths.dynamic_programming, + output_edge_order=[], + strip_exponent=True, + ) + finally: + cons.set_contraction_algebra(prev) diff --git a/tests/test_tropical.py b/tests/test_tropical.py new file mode 100644 index 00000000..33a6b681 --- /dev/null +++ b/tests/test_tropical.py @@ -0,0 +1,1138 @@ +"""Tropical (max-plus) and counting tropical algebra tests. + +Covers: max-plus primitives, einsum/tensordot, configuration recovery, +counting (energy + degeneracy), Ising end-to-end, and the non-scalar +aux-reorder canaries. +""" + +import itertools + +import numpy as np +import pytest +import tensornetwork as tn +import opt_einsum + +import tensorcircuit as tc +import tensorcircuit.cons as cons +import tensorcircuit.backends.numpy_backend as nb +import applications.tropical_algebra as tr +from applications.tropical_algebra import ( + MaxPlusAlgebra, + MaxPlusTrackingAlgebra, + CountingTropicalAlgebra, + tropical, + counting_tropical, + recover_configuration, + get_recorded_topology, + degeneracy, + split_energy_count, + _tropical_tensordot, + _tropical_einsum, + _counting_tensordot, + _counting_einsum, + _td_backtrack_step, +) + +# ═══════════════════════════════════════════════════════════════════════════════ +# Shared utilities (was tests/_tropical_test_utils.py) +# ═══════════════════════════════════════════════════════════════════════════════ + + +def build_ising_tn(spins, edges, j_vals, h): + """Build the tropical (max-plus) Ising tensor network as tensornetwork nodes. + + Per spin: a ``tn.CopyNode`` of rank (degree+1), dimension 2. A CopyNode is a true + hyperedge hub -- tensornetwork/cotengra identify all its legs as one shared index + (the delta/copy constraint is structural, not stored as data). This is what makes + the contraction route hyperedge-containing pairs through the tropical **einsum** + branch and ordinary pairs through the tropical **tensordot** branch, covering both. + Per edge (i,j): Te[si,sj] = J*si*sj -> [[J,-J],[-J,J]]. + Per spin: field Tv[s] = h*s -> [h,-h]. + Contracting over max-plus yields max_cfg(-E) = -E_ground. + """ + be = tc.backend + degree = dict.fromkeys(spins, 0) + for i, j in edges: + degree[i] += 1 + degree[j] += 1 + + nodes = [] + copy_nodes = {} + for i in spins: + cn = tn.CopyNode(degree[i] + 1, 2) + copy_nodes[i] = cn + nodes.append(cn) + tv = np.array([h[i], -h[i]], dtype=np.float64) + tvn = tn.Node(be.cast(be.convert_to_tensor(tv), "float64")) + tn.connect(cn[0], tvn[0]) + nodes.append(tvn) + leg_idx = dict.fromkeys(spins, 1) + for (i, j), jij in zip(edges, j_vals): + te = np.array([[jij, -jij], [-jij, jij]], dtype=np.float64) + ten = tn.Node(be.cast(be.convert_to_tensor(te), "float64")) + tn.connect(ten[0], copy_nodes[i][leg_idx[i]]) + tn.connect(ten[1], copy_nodes[j][leg_idx[j]]) + leg_idx[i] += 1 + leg_idx[j] += 1 + nodes.append(ten) + return nodes + + +def brute_force_energy(spins, edges, j_vals, h): + """Brute-force min energy: E = -sum J s_i s_j - sum h s_i.""" + n = len(spins) + best = np.inf + for cfg in itertools.product([-1, 1], repeat=n): + e = 0.0 + for (i, j), jij in zip(edges, j_vals): + e -= jij * cfg[i] * cfg[j] + for i, hi in zip(spins, h): + e -= hi * cfg[i] + best = min(best, e) + return best + + +def brute_force_energy_and_degeneracy(n, edges, j_vals, h): + """Brute-force ground energy + degeneracy over spin configs. + + cfg in {0,1}^n; spin s = 1 - 2*cfg (cfg=0 -> s=+1). + E = -sum J s_i s_j - sum h s_i. Returns (best_e, degeneracy). + """ + best_e = None + deg = 0 + for cfg in itertools.product([0, 1], repeat=n): + e = 0.0 + for (i, j), jij in zip(edges, j_vals): + si, sj = 1 - 2 * cfg[i], 1 - 2 * cfg[j] + e -= jij * si * sj + for i, hi in zip(range(n), h): + e -= hi * (1 - 2 * cfg[i]) + if best_e is None or e < best_e - 1e-9: + best_e, deg = e, 1 + elif abs(e - best_e) < 1e-9: + deg += 1 + return best_e, deg + + +# --- Non-scalar (free-spin) helpers --- + +_RING_SPINS = (0, 1, 2, 3) +_RING_EDGES = ((0, 1), (1, 2), (2, 3), (3, 0)) +_RING_J = (-2, 1, 1, -2) +_RING_H = (2, 0, 2, -1) +_FREE_SPIN = 0 +_FREE_SPIN_A = 0 +_FREE_SPIN_B = 1 + + +def build_ring_with_free_spin( + spins=_RING_SPINS, + edges=_RING_EDGES, + j_vals=_RING_J, + h=_RING_H, + free_spin=_FREE_SPIN, +): + """Build a small ring Ising with ONE spin's CopyNode given an extra dangling + leg (the free output index).""" + be = tc.backend + degree = dict.fromkeys(spins, 0) + for i, j in edges: + degree[i] += 1 + degree[j] += 1 + + nodes = [] + copy_nodes = {} + for i in spins: + rank = degree[i] + 1 + if i == free_spin: + rank += 1 + cn = tn.CopyNode(rank, 2) + copy_nodes[i] = cn + nodes.append(cn) + tv = np.array([h[i], -h[i]], dtype=np.float64) + tvn = tn.Node(be.cast(be.convert_to_tensor(tv), "float64")) + tn.connect(cn[0], tvn[0]) + nodes.append(tvn) + + leg_idx = dict.fromkeys(spins, 1) + for (i, j), jij in zip(edges, j_vals): + te = np.array([[jij, -jij], [-jij, jij]], dtype=np.float64) + ten = tn.Node(be.cast(be.convert_to_tensor(te), "float64")) + tn.connect(ten[0], copy_nodes[i][leg_idx[i]]) + tn.connect(ten[1], copy_nodes[j][leg_idx[j]]) + leg_idx[i] += 1 + leg_idx[j] += 1 + nodes.append(ten) + + free_edge = copy_nodes[free_spin][degree[free_spin] + 1] + return nodes, free_edge, j_vals, h + + +def build_ring_with_two_free_spins( + spins=_RING_SPINS, + edges=_RING_EDGES, + j_vals=_RING_J, + h=_RING_H, + free_spin_a=_FREE_SPIN_A, + free_spin_b=_FREE_SPIN_B, +): + """Build a small ring Ising with TWO spins' CopyNodes given an extra dangling + leg each (two free output indices).""" + be = tc.backend + free_spins = {free_spin_a, free_spin_b} + degree = dict.fromkeys(spins, 0) + for i, j in edges: + degree[i] += 1 + degree[j] += 1 + + nodes = [] + copy_nodes = {} + for i in spins: + rank = degree[i] + 1 + if i in free_spins: + rank += 1 + cn = tn.CopyNode(rank, 2) + copy_nodes[i] = cn + nodes.append(cn) + tv = np.array([h[i], -h[i]], dtype=np.float64) + tvn = tn.Node(be.cast(be.convert_to_tensor(tv), "float64")) + tn.connect(cn[0], tvn[0]) + nodes.append(tvn) + + leg_idx = dict.fromkeys(spins, 1) + for (i, j), jij in zip(edges, j_vals): + te = np.array([[jij, -jij], [-jij, jij]], dtype=np.float64) + ten = tn.Node(be.cast(be.convert_to_tensor(te), "float64")) + tn.connect(ten[0], copy_nodes[i][leg_idx[i]]) + tn.connect(ten[1], copy_nodes[j][leg_idx[j]]) + leg_idx[i] += 1 + leg_idx[j] += 1 + nodes.append(ten) + + edge_a = copy_nodes[free_spin_a][degree[free_spin_a] + 1] + edge_b = copy_nodes[free_spin_b][degree[free_spin_b] + 1] + return nodes, edge_a, edge_b, j_vals, h + + +def _ising_config_energy(cfg, edges, j_vals, spins, h): + """Ising energy E = -sum J s_i s_j - sum h s_i for a config dict (spin idx -> 0/1).""" + energy = 0.0 + for (i, j), jij in zip(edges, j_vals): + energy -= jij * (1 - 2 * cfg[i]) * (1 - 2 * cfg[j]) + for i, hi in zip(spins, h): + energy -= hi * (1 - 2 * cfg[i]) + return energy + + +def _min_energy_and_degeneracy(energies, eps=1e-9): + """Return ``(min energy, degeneracy)`` -- degeneracy counts configs within eps of min.""" + best = None + deg = 0 + for e in energies: + if best is None or e < best - eps: + best, deg = e, 1 + elif abs(e - best) < eps: + deg += 1 + return best, deg + + +def brute_nonscalar_counting( + nodes, + free_edge, + j_vals, + h, + spins=_RING_SPINS, + edges=_RING_EDGES, + free_spin=_FREE_SPIN, +): + """Brute-force per-output (energy, degeneracy) for the non-scalar ring with + one free spin.""" + del nodes, free_edge + others = [s for s in spins if s != free_spin] + + expected_e = np.zeros(2, dtype=np.float64) + expected_n = np.zeros(2, dtype=np.float64) + for v_free in (0, 1): + fixed = {free_spin: v_free} + energies = [ + _ising_config_energy( + {**fixed, **dict(zip(others, cfg_rest))}, edges, j_vals, spins, h + ) + for cfg_rest in itertools.product([0, 1], repeat=len(others)) + ] + best_e, deg = _min_energy_and_degeneracy(energies) + expected_e[v_free] = -best_e + expected_n[v_free] = deg + return expected_e, expected_n + + +def brute_nonscalar_counting_2d( + nodes, + edge_a, + edge_b, + j_vals, + h, + spins=_RING_SPINS, + edges=_RING_EDGES, + free_spin_a=_FREE_SPIN_A, + free_spin_b=_FREE_SPIN_B, +): + """Brute-force per-output (energy, degeneracy) for the non-scalar ring with + TWO free spins.""" + del nodes, edge_a, edge_b + others = [s for s in spins if s not in (free_spin_a, free_spin_b)] + + expected_e = np.zeros((2, 2), dtype=np.float64) + expected_n = np.zeros((2, 2), dtype=np.float64) + for v_a in (0, 1): + for v_b in (0, 1): + fixed = {free_spin_a: v_a, free_spin_b: v_b} + energies = [ + _ising_config_energy( + {**fixed, **dict(zip(others, cfg_rest))}, edges, j_vals, spins, h + ) + for cfg_rest in itertools.product([0, 1], repeat=len(others)) + ] + best_e, deg = _min_energy_and_degeneracy(energies) + expected_e[v_a, v_b] = -best_e + expected_n[v_a, v_b] = deg + return expected_e, expected_n + + +# ═══════════════════════════════════════════════════════════════════════════════ +# Section 1 — Core algebra: tensordot + einsum (was test_tropical_algebra.py) +# ═══════════════════════════════════════════════════════════════════════════════ + + +def _be(): + return tc.backend + + +def _ref_tropical_tensordot(anp, bnp, axes): + if isinstance(axes, int): + a_axes = list(range(anp.ndim - axes, anp.ndim)) + b_axes = list(range(0, axes)) + else: + a_axes, b_axes = list(axes[0]), list(axes[1]) + a_free = [i for i in range(anp.ndim) if i not in a_axes] + b_free = [i for i in range(bnp.ndim) if i not in b_axes] + a_t = np.transpose(anp, a_free + a_axes) + b_t = np.transpose(bnp, b_axes + b_free) + a_fs = [anp.shape[i] for i in a_free] + b_fs = [bnp.shape[i] for i in b_free] + a2 = a_t.reshape(-1, np.prod([anp.shape[i] for i in a_axes], dtype=int) or 1) + b2 = b_t.reshape(np.prod([bnp.shape[i] for i in b_axes], dtype=int) or 1, -1) + A, b_n = a2.shape[0], b2.shape[1] + res = np.full((A, b_n), -np.inf) + for i in range(A): + for j in range(b_n): + res[i, j] = np.max(a2[i, :] + b2[:, j]) + return res.reshape(tuple(a_fs) + tuple(b_fs)) + + +def _ref_tropical_einsum(eq, a, b): + lhs, rhs = eq.split("->") + ia, ib = lhs.split(",") + sizes = {} + for s, t in zip([ia, ib], [a, b]): + for c, dim in zip(s, t.shape): + sizes[c] = dim + out = np.full([sizes[c] for c in rhs], -np.inf) + allc = list(dict.fromkeys(list(ia) + list(ib))) + for combo in itertools.product(*[range(sizes[c]) for c in allc]): + env = dict(zip(allc, combo)) + ia_idx = tuple(env[c] for c in ia) + ib_idx = tuple(env[c] for c in ib) + val = a[ia_idx] + b[ib_idx] + oidx = tuple(env[c] for c in rhs) + if val > out[oidx]: + out[oidx] = val + return out + + +# --- tensordot tests --- + + +def test_tropical_tensordot_matrix(): + be = _be() + rng = np.random.default_rng(0) + anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(4, 5)) + a = be.cast(be.convert_to_tensor(anp), "float64") + b = be.cast(be.convert_to_tensor(bnp), "float64") + got = np.array(_tropical_tensordot(be, a, b, axes=1)) + np.testing.assert_allclose(got, _ref_tropical_tensordot(anp, bnp, 1)) + + +def test_tropical_tensordot_multi_axis(): + be = _be() + rng = np.random.default_rng(1) + anp, bnp = rng.normal(size=(2, 3, 4)), rng.normal(size=(3, 4, 5)) + a = be.cast(be.convert_to_tensor(anp), "float64") + b = be.cast(be.convert_to_tensor(bnp), "float64") + got = np.array(_tropical_tensordot(be, a, b, axes=([1, 2], [0, 1]))) + np.testing.assert_allclose(got, _ref_tropical_tensordot(anp, bnp, ([1, 2], [0, 1]))) + assert got.shape == (2, 5) + + +# --- einsum tests --- + + +def test_tropical_einsum_pair(): + be = _be() + rng = np.random.default_rng(2) + anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(4, 5)) + a = be.cast(be.convert_to_tensor(anp), "float64") + b = be.cast(be.convert_to_tensor(bnp), "float64") + got = np.array(_tropical_einsum(be, "ab,bc->ac", a, b)) + np.testing.assert_allclose(got, _ref_tropical_einsum("ab,bc->ac", anp, bnp)) + + +def test_tropical_einsum_hyperedge_shared_index(): + be = _be() + rng = np.random.default_rng(3) + anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(3, 5)) + a = be.cast(be.convert_to_tensor(anp), "float64") + b = be.cast(be.convert_to_tensor(bnp), "float64") + got = np.array(_tropical_einsum(be, "ab,ac->abc", a, b)) + ref = _ref_tropical_einsum("ab,ac->abc", anp, bnp) + np.testing.assert_allclose(got, ref) + + +def test_maxplus_algebra_uses_tropical_ops(): + be = _be() + rng = np.random.default_rng(4) + anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(4, 5)) + a = be.cast(be.convert_to_tensor(anp), "float64") + b = be.cast(be.convert_to_tensor(bnp), "float64") + alg = MaxPlusAlgebra() + assert alg.name == "maxplus" + np.testing.assert_allclose( + np.array(alg.tensordot(be, a, b, 1)), + _ref_tropical_tensordot(anp, bnp, 1), + ) + np.testing.assert_allclose( + np.array(alg.einsum(be, "ab,bc->ac", a, b)), + _ref_tropical_einsum("ab,bc->ac", anp, bnp), + ) + + +def test_tropical_einsum_single_tensor_transpose_ok(): + be = _be() + rng = np.random.default_rng(5) + anp = rng.normal(size=(3, 4)) + a = be.cast(be.convert_to_tensor(anp), "float64") + got = np.array(_tropical_einsum(be, "ab->ba", a)) + np.testing.assert_allclose(got, np.array(be.einsum("ab->ba", a))) + + +def test_tropical_einsum_single_tensor_trace(): + be = _be() + rng = np.random.default_rng(6) + square_a = rng.normal(size=(3, 3)) + a = be.cast(be.convert_to_tensor(square_a), "float64") + got = float(np.array(_tropical_einsum(be, "ii->", a))) + assert np.isclose(got, np.max(np.diag(square_a))) + + +def test_tropical_einsum_single_tensor_diagonal(): + be = _be() + rng = np.random.default_rng(7) + square_a = rng.normal(size=(4, 4)) + a = be.cast(be.convert_to_tensor(square_a), "float64") + got = np.array(_tropical_einsum(be, "ii->i", a)) + np.testing.assert_allclose(got, np.diag(square_a)) + + +def test_tropical_einsum_single_tensor_reduce(): + be = _be() + rng = np.random.default_rng(8) + anp = rng.normal(size=(3, 5)) + a = be.cast(be.convert_to_tensor(anp), "float64") + got = np.array(_tropical_einsum(be, "ab->a", a)) + np.testing.assert_allclose(got, np.max(anp, axis=1)) + + +def test_tropical_einsum_single_tensor_partial_trace(): + be = _be() + rng = np.random.default_rng(9) + anp = rng.normal(size=(4, 4, 3)) + a = be.cast(be.convert_to_tensor(anp), "float64") + got = np.array(_tropical_einsum(be, "iij->j", a)) + ref = np.max(np.array([anp[i, i, :] for i in range(4)]), axis=0) + np.testing.assert_allclose(got, ref) + + +def test_tropical_einsum_intra_operand_repeat_first(): + be = _be() + rng = np.random.default_rng(20) + anp, bnp = rng.normal(size=(3, 3, 4)), rng.normal(size=(4, 5)) + a = be.cast(be.convert_to_tensor(anp), "float64") + b = be.cast(be.convert_to_tensor(bnp), "float64") + got = np.array(_tropical_einsum(be, "iij,jk->ik", a, b)) + np.testing.assert_allclose(got, _ref_tropical_einsum("iij,jk->ik", anp, bnp)) + + +def test_tropical_einsum_intra_operand_repeat_contracted(): + be = _be() + rng = np.random.default_rng(21) + anp, bnp = rng.normal(size=(3, 3, 4)), rng.normal(size=(4, 5)) + a = be.cast(be.convert_to_tensor(anp), "float64") + b = be.cast(be.convert_to_tensor(bnp), "float64") + got = np.array(_tropical_einsum(be, "iij,jk->k", a, b)) + np.testing.assert_allclose(got, _ref_tropical_einsum("iij,jk->k", anp, bnp)) + + +def test_tropical_einsum_both_operands_repeat(): + be = _be() + rng = np.random.default_rng(22) + anp, bnp = rng.normal(size=(3, 3, 4)), rng.normal(size=(4, 4)) + a = be.cast(be.convert_to_tensor(anp), "float64") + b = be.cast(be.convert_to_tensor(bnp), "float64") + got = np.array(_tropical_einsum(be, "iij,jj->ij", a, b)) + np.testing.assert_allclose(got, _ref_tropical_einsum("iij,jj->ij", anp, bnp)) + + +def test_tropical_context_changes_result(): + be = _be() + anp = np.array([[1.0, 5.0], [3.0, 2.0]]) + bnp = np.array([10.0, 0.0]) + a = be.cast(be.convert_to_tensor(anp), "float64") + b = be.cast(be.convert_to_tensor(bnp), "float64") + na, nb = tn.Node(a), tn.Node(b) + tn.connect(na[1], nb[0]) + expected = np.array([max(1 + 10, 5 + 0), max(3 + 10, 2 + 0)]) + with tropical(): + got = np.array(cons.contractor([na, nb], output_edge_order=[na[0]]).tensor) + np.testing.assert_allclose(got, expected) + + +def test_preprocessing_does_not_corrupt_tropical(): + be = _be() + rng = np.random.default_rng(0) + n0 = rng.normal(size=(2,)) + n1 = rng.normal(size=(2, 2)) + n2 = rng.normal(size=(2, 2)) + n3 = rng.normal(size=(2, 2)) + n4 = rng.normal(size=(2,)) + na = tn.Node(be.cast(be.convert_to_tensor(n0), "float64")) + nb = tn.Node(be.cast(be.convert_to_tensor(n1), "float64")) + nc = tn.Node(be.cast(be.convert_to_tensor(n2), "float64")) + nd = tn.Node(be.cast(be.convert_to_tensor(n3), "float64")) + ne = tn.Node(be.cast(be.convert_to_tensor(n4), "float64")) + tn.connect(na[0], nb[0]) + tn.connect(nb[1], nc[0]) + tn.connect(nc[1], nd[0]) + tn.connect(nd[1], ne[0]) + nodes = [na, nb, nc, nd, ne] + assert len(nodes) >= 5 + + ref = -np.inf + for s0, s1, s2, s3 in itertools.product(range(2), repeat=4): + val = n0[s0] + n1[s0, s1] + n2[s1, s2] + n3[s2, s3] + n4[s3] + ref = max(ref, val) + + with tropical(): + got = np.array(cons.contractor(nodes, output_edge_order=[]).tensor).item() + + np.testing.assert_allclose(got, ref, atol=1e-6) + + +def test_tropical_public_api_surface(): + assert MaxPlusAlgebra is not None + assert MaxPlusTrackingAlgebra is not None + assert CountingTropicalAlgebra is not None + for name in [ + "tropical", + "counting_tropical", + "recover_configuration", + "split_energy_count", + ]: + assert hasattr(tr, name), name + + +# ═══════════════════════════════════════════════════════════════════════════════ +# Section 2 — Configuration recovery (was test_tropical_config.py) +# ═══════════════════════════════════════════════════════════════════════════════ + + +def _brute_cfg(n, edges, j_vals, h): + """Brute-force min energy over spin configs; returns (best_e, set_of_optimal_cfgs).""" + best_e = None + best_cfgs = set() + for cfg in itertools.product([0, 1], repeat=n): + e = 0.0 + for (i, j), jij in zip(edges, j_vals): + e -= jij * (1 - 2 * cfg[i]) * (1 - 2 * cfg[j]) + for i, hi in zip(range(n), h): + e -= hi * (1 - 2 * cfg[i]) + if best_e is None or e < best_e - 1e-9: + best_e = e + best_cfgs = {cfg} + elif abs(e - best_e) < 1e-9: + best_cfgs.add(cfg) + return best_e, best_cfgs + + +def _energy_of(assignment, input_sets, raw_tensors): + """Max-plus total for a symbol->value assignment = -E for that config.""" + total = 0.0 + for term, t in zip(input_sets, raw_tensors): + idx = tuple(int(assignment[c]) for c in term) + total += float(np.asarray(t)[idx]) + return total + + +def _contract_tropical_track(nodes): + """Contract under tracking tropical; return (value, config).""" + with tropical(track=True): + val = np.array(cons.contractor(nodes, output_edge_order=[]).tensor).item() + cfg = recover_configuration() + return val, cfg + + +def test_config_recovery_tiny_ising_unique(): + spins = [0, 1] + edges = [(0, 1)] + J = [0.7] + h = [0.3, -0.2] + best_e, best_cfgs = _brute_cfg(len(spins), edges, J, h) + + nodes = build_ising_tn(spins, edges, J, h) + val, cfg = _contract_tropical_track(nodes) + + np.testing.assert_allclose(val, -best_e, atol=1e-6) + + _tree, input_sets, raw_tensors = get_recorded_topology() + assert input_sets is not None and raw_tensors is not None + all_labels = set().union(*[set(t) for t in input_sets]) + assert set(cfg.keys()) == all_labels + + recovered_total = _energy_of(cfg, input_sets, raw_tensors) + np.testing.assert_allclose(recovered_total, val, atol=1e-6) + np.testing.assert_allclose(recovered_total, -best_e, atol=1e-6) + + spin_cfg = {} + for term, _t in zip(input_sets, raw_tensors): + if len(term) == 1: + spin_cfg[len(spin_cfg)] = int(cfg[term[0]]) + assert tuple(spin_cfg[i] for i in range(len(spins))) in best_cfgs + + +def test_config_recovery_tiny_ising_degenerate(): + spins = [0, 1] + edges = [(0, 1)] + J = [1.0] + h = [0.0, 0.0] + best_e, best_cfgs = _brute_cfg(len(spins), edges, J, h) + assert len(best_cfgs) == 2 + + nodes = build_ising_tn(spins, edges, J, h) + val, cfg = _contract_tropical_track(nodes) + + np.testing.assert_allclose(val, -best_e, atol=1e-6) + _tree, input_sets, raw_tensors = get_recorded_topology() + recovered_total = _energy_of(cfg, input_sets, raw_tensors) + np.testing.assert_allclose(recovered_total, -best_e, atol=1e-6) + + spin_cfg = {} + for term, _t in zip(input_sets, raw_tensors): + if len(term) == 1: + spin_cfg[len(spin_cfg)] = int(cfg[term[0]]) + assert tuple(spin_cfg[i] for i in range(len(spins))) in best_cfgs + + +def test_config_recovery_three_ring(): + spins = [0, 1, 2] + edges = [(0, 1), (1, 2)] + J = [0.5, -0.8] + h = [0.1, 0.2, -0.3] + best_e, _best_cfgs = _brute_cfg(len(spins), edges, J, h) + + nodes = build_ising_tn(spins, edges, J, h) + val, cfg = _contract_tropical_track(nodes) + np.testing.assert_allclose(val, -best_e, atol=1e-6) + + _tree, input_sets, raw_tensors = get_recorded_topology() + recovered_total = _energy_of(cfg, input_sets, raw_tensors) + np.testing.assert_allclose(recovered_total, -best_e, atol=1e-6) + + +def test_config_recovery_non_scalar_raises(): + be = tc.backend + a_np = np.array([[1.0, 4.0, 2.0], [3.0, 1.0, 5.0], [0.0, 2.0, 1.0]]) + b_np = np.array([0.5, 1.0, -0.5]) + + a_node = tn.Node(be.cast(be.convert_to_tensor(a_np), "float64")) + b_node = tn.Node(be.cast(be.convert_to_tensor(b_np), "float64")) + tn.connect(a_node[1], b_node[0]) + nodes = [a_node, b_node] + + with tropical(track=True): + res = cons.contractor(nodes, output_edge_order=[a_node[0]]) + assert np.array(res.tensor).shape == (3,) + with pytest.raises(NotImplementedError): + recover_configuration() + + +class _StubTree: + """Minimal stand-in for a cotengra ``ContractionTree``.""" + + def __init__(self, inds, axes, perm): + self._inds = inds + self._axes = axes + self._perm = perm + + def get_inds(self, node): + return self._inds[node] + + def get_tensordot_axes(self, node): + return self._axes[node] + + def get_tensordot_perm(self, node): + return self._perm[node] + + +def test_td_backtrack_step_perm_inversion(): + l_node, r_node, p_node = "l", "r", "p" + tree = _StubTree( + inds={l_node: "ab", r_node: "bcd", p_node: "dac"}, + axes={p_node: ([1], [0])}, + perm={p_node: (2, 0, 1)}, + ) + argmax = np.zeros((2, 3, 2), dtype=np.int64) + argmax[1, 2, 0] = 3 + rec = {"kind": "td", "argmax": argmax, "contract_dims": (4,)} + p_pos = (0, 1, 2) + assignment = {} + l_pos, r_pos = _td_backtrack_step( + tree, p_node, l_node, r_node, rec, p_pos, assignment + ) + + assert assignment["b"] == 3 + assert l_pos == (1, 3) + assert r_pos == (3, 2, 0) + + +# ═══════════════════════════════════════════════════════════════════════════════ +# Section 3 — Counting (was test_tropical_counting.py) +# ═══════════════════════════════════════════════════════════════════════════════ + + +def _ref_counting_tensordot(anp, bnp, axes, eps=1e-9): + if isinstance(axes, int): + a_axes = list(range(anp.ndim - axes, anp.ndim)) + b_axes = list(range(0, axes)) + else: + a_axes, b_axes = list(axes[0]), list(axes[1]) + a_free = [i for i in range(anp.ndim) if i not in a_axes] + b_free = [i for i in range(bnp.ndim) if i not in b_axes] + at = np.transpose(anp, a_free + a_axes) + bt = np.transpose(bnp, b_axes + b_free) + a_fs = [anp.shape[i] for i in a_free] + b_fs = [bnp.shape[i] for i in b_free] + k = int(np.prod([anp.shape[i] for i in a_axes]) or 1) + A = int(np.prod(a_fs) or 1) + b_n = int(np.prod(b_fs) or 1) + a2 = at.reshape(A, k) + b2 = bt.reshape(k, b_n) + out_e = np.full((A, b_n), -np.inf) + out_n = np.zeros((A, b_n)) + for i in range(A): + for j in range(b_n): + s = a2[i, :] + b2[:, j] + mx = np.max(s) + out_e[i, j] = mx + out_n[i, j] = int(np.sum(np.abs(s - mx) < eps)) + return out_e.reshape(a_fs + b_fs), out_n.reshape(a_fs + b_fs) + + +def _ref_counting_einsum(eq, a, b, eps=1e-9): + lhs, rhs = eq.split("->") + ia, ib = lhs.split(",") + sizes = {} + for s, t in zip([ia, ib], [a, b]): + for c, dim in zip(s, t.shape): + sizes[c] = dim + out_e = np.full([sizes[c] for c in rhs], -np.inf) + out_n = np.zeros([sizes[c] for c in rhs]) + allc = list(dict.fromkeys(list(ia) + list(ib))) + for combo in itertools.product(*[range(sizes[c]) for c in allc]): + env = dict(zip(allc, combo)) + e = a[tuple(env[c] for c in ia)] + b[tuple(env[c] for c in ib)] + oidx = tuple(env[c] for c in rhs) + if e > out_e[oidx] + eps: + out_e[oidx] = e + out_n[oidx] = 1 + elif abs(e - out_e[oidx]) < eps: + out_n[oidx] += 1 + return out_e, out_n + + +def test_counting_tensordot_matrix(): + be = _be() + rng = np.random.default_rng(10) + anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(4, 5)) + a = be.cast( + be.convert_to_tensor(np.stack([anp, np.ones_like(anp)], axis=-1)), "float64" + ) + b = be.cast( + be.convert_to_tensor(np.stack([bnp, np.ones_like(bnp)], axis=-1)), "float64" + ) + got = np.array(_counting_tensordot(be, a, b, axes=1)) + ref_e, ref_n = _ref_counting_tensordot(anp, bnp, 1) + np.testing.assert_allclose(got[..., 0], ref_e) + np.testing.assert_allclose(got[..., 1], ref_n) + + +def test_split_energy_count(): + be = _be() + arr = np.array([[1.0, 2.0], [3.0, 4.0]]) + stacked = be.cast( + be.convert_to_tensor(np.stack([arr, arr * 2], axis=-1)), "float64" + ) + e, n = split_energy_count(stacked) + np.testing.assert_allclose(np.array(e), arr) + np.testing.assert_allclose(np.array(n), arr * 2) + + +def test_counting_einsum_pair(): + be = _be() + rng = np.random.default_rng(11) + anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(4, 5)) + a = be.cast(be.convert_to_tensor(np.stack([anp, np.ones_like(anp)], -1)), "float64") + b = be.cast(be.convert_to_tensor(np.stack([bnp, np.ones_like(bnp)], -1)), "float64") + got = np.array(_counting_einsum(be, "ab,bc->ac", a, b)) + ref_e, ref_n = _ref_counting_einsum("ab,bc->ac", anp, bnp) + np.testing.assert_allclose(got[..., 0], ref_e) + np.testing.assert_allclose(got[..., 1], ref_n) + + +def test_counting_einsum_hyperedge(): + be = _be() + rng = np.random.default_rng(12) + anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(3, 5)) + a = be.cast(be.convert_to_tensor(np.stack([anp, np.ones_like(anp)], -1)), "float64") + b = be.cast(be.convert_to_tensor(np.stack([bnp, np.ones_like(bnp)], -1)), "float64") + got = np.array(_counting_einsum(be, "ab,ac->abc", a, b)) + ref_e, ref_n = _ref_counting_einsum("ab,ac->abc", anp, bnp) + np.testing.assert_allclose(got[..., 0], ref_e) + np.testing.assert_allclose(got[..., 1], ref_n) + + +def test_counting_einsum_tie_degeneracy(): + be = _be() + anp = np.array([[1.0, 1.0]]) + bnp = np.array([[2.0], [2.0]]) + a = be.cast(be.convert_to_tensor(np.stack([anp, np.ones_like(anp)], -1)), "float64") + b = be.cast(be.convert_to_tensor(np.stack([bnp, np.ones_like(bnp)], -1)), "float64") + got = np.array(_counting_einsum(be, "ab,bc->ac", a, b)) + ref_e, ref_n = _ref_counting_einsum("ab,bc->ac", anp, bnp) + np.testing.assert_allclose(got[..., 0], ref_e) + np.testing.assert_allclose(got[..., 1], ref_n) + assert ref_n[0, 0] == 2 + np.testing.assert_allclose(got[0, 0, 0], 3.0) + np.testing.assert_allclose(got[0, 0, 1], 2.0) + + +def test_counting_einsum_single_operand_diagonal(): + be = nb.NumpyBackend() + e = np.array([[1.0, 5.0], [3.0, 2.0]]) + n = np.array([[1.0, 2.0], [1.0, 1.0]]) + a = np.stack([e, n], axis=-1) + out = _counting_einsum(be, "aa->a", a) + np.testing.assert_allclose(out[..., 0], np.diag(e)) + np.testing.assert_allclose(out[..., 1], np.diag(n)) + + +def test_counting_einsum_single_operand_reduce(): + be = nb.NumpyBackend() + e = np.array([[1.0, 5.0], [3.0, 3.0]]) + n = np.array([[1.0, 2.0], [1.0, 4.0]]) + a = np.stack([e, n], axis=-1) + out = _counting_einsum(be, "ab->a", a) + np.testing.assert_allclose(out[..., 0], [5.0, 3.0]) + np.testing.assert_allclose(out[..., 1], [2.0, 5.0]) + + +def test_counting_einsum_two_operand_intra_repeat(): + be = nb.NumpyBackend() + e_a = np.array([[2.0, 1.0], [1.0, 2.0]]) + n_a = np.ones((2, 2)) + a = np.stack([e_a, n_a], axis=-1) + e_b = np.array([0.0, 0.0]) + n_b = np.array([1.0, 1.0]) + b = np.stack([e_b, n_b], axis=-1) + out = _counting_einsum(be, "aa,b->ab", a, b) + np.testing.assert_allclose(out[..., 0], [[2, 2], [2, 2]]) + np.testing.assert_allclose(out[..., 1], [[1, 1], [1, 1]]) + + +def test_counting_einsum_multi_axis(): + be = _be() + rng = np.random.default_rng(14) + anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(3, 5)) + a = be.cast(be.convert_to_tensor(np.stack([anp, np.ones_like(anp)], -1)), "float64") + b = be.cast(be.convert_to_tensor(np.stack([bnp, np.ones_like(bnp)], -1)), "float64") + got = np.array(_counting_einsum(be, "ab,ac->c", a, b)) + ref_e, ref_n = _ref_counting_einsum("ab,ac->c", anp, bnp) + np.testing.assert_allclose(got[..., 0], ref_e) + np.testing.assert_allclose(got[..., 1], ref_n) + + +def test_degeneracy_none_after_standard_contraction(): + spins = [0, 1, 2, 3, 4] + edges = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 0)] + J = [1.0, 1.0, 1.0, 1.0, 1.0] + h = [0.0, 0.0, 0.0, 0.0, 0.0] + expected_e, expected_n = brute_force_energy_and_degeneracy(len(spins), edges, J, h) + assert expected_n == 2 + nodes = build_ising_tn(spins, edges, J, h) + + with counting_tropical(): + node = cons._algebraic_base_contraction( + nodes, + algorithm=opt_einsum.paths.dynamic_programming, + output_edge_order=[], + ignore_edge_order=True, + ) + np.testing.assert_allclose(np.array(degeneracy()), expected_n) + np.testing.assert_allclose(np.array(node.tensor), -expected_e, atol=1e-6) + + rng = np.random.default_rng(0) + a = tn.Node(rng.standard_normal((2, 3)).astype(np.complex64)) + b = tn.Node(rng.standard_normal((3, 4)).astype(np.complex64)) + tn.connect(a[1], b[0]) + cons._algebraic_base_contraction( + [a, b], + algorithm=opt_einsum.paths.dynamic_programming, + output_edge_order=[a[0], b[1]], + ) + assert degeneracy() is None + + +def test_nonscalar_counting_energy_and_degeneracy_per_output(): + nodes, free_edge, J, h = build_ring_with_free_spin() + expected_e, expected_n = brute_nonscalar_counting(nodes, free_edge, J, h) + + with counting_tropical(): + node = cons._algebraic_base_contraction( + nodes, + algorithm=opt_einsum.paths.dynamic_programming, + output_edge_order=[free_edge], + ) + E = np.asarray(node.tensor) + N = np.asarray(degeneracy()) + + np.testing.assert_allclose(E, expected_e) + np.testing.assert_allclose(N, expected_n) + assert E[0] != E[1] or N[0] != N[1] + + +def test_nonscalar_counting_two_free_spins_reversed_order(): + nodes, edge_a, edge_b, J, h = build_ring_with_two_free_spins() + expected_e_ab, expected_n_ab = brute_nonscalar_counting_2d( + nodes, edge_a, edge_b, J, h + ) + assert not np.allclose( + expected_e_ab, expected_e_ab.T + ), "weak canary: brute-force energy is symmetric under transpose" + assert not np.allclose( + expected_n_ab, expected_n_ab.T + ), "weak canary: brute-force degeneracy is symmetric under transpose" + + with counting_tropical(): + node = cons._algebraic_base_contraction( + nodes, + algorithm=opt_einsum.paths.dynamic_programming, + output_edge_order=[edge_b, edge_a], + ) + E = np.asarray(node.tensor) + N = np.asarray(degeneracy()) + + np.testing.assert_allclose(E, expected_e_ab.T) + np.testing.assert_allclose(N, expected_n_ab.T) + + +# ═══════════════════════════════════════════════════════════════════════════════ +# Section 4 — Example-level tests (was test_tropical_example.py) +# ═══════════════════════════════════════════════════════════════════════════════ + + +def test_example_maxplus_tensordot_matches_brute(): + be = tc.backend + rng = np.random.default_rng(0) + anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(4, 5)) + a = be.cast(be.convert_to_tensor(anp), "float64") + b = be.cast(be.convert_to_tensor(bnp), "float64") + got = np.array(MaxPlusAlgebra().tensordot(be, a, b, 1)) + ref = np.max(anp[:, :, None] + bnp[None, :, :], axis=1) + np.testing.assert_allclose(got, ref) + + +def test_example_recover_configuration_on_tiny_ising(): + be = tc.backend + n = tn.Node(be.cast(be.convert_to_tensor(np.array([0.0, 0.0])), "float64")) + e = tn.Node( + be.cast(be.convert_to_tensor(np.array([[1.0, -1.0], [-1.0, 1.0]])), "float64") + ) + m = tn.Node(be.cast(be.convert_to_tensor(np.array([0.5, -0.5])), "float64")) + tn.connect(n[0], e[0]) + tn.connect(e[1], m[0]) + nodes = [n, e, m] + with tropical(track=True): + float(cons.contractor(nodes, output_edge_order=[]).tensor) + cfg = recover_configuration() + assert isinstance(cfg, dict) + assert len(cfg) >= 1 + + +# ═══════════════════════════════════════════════════════════════════════════════ +# Section 5 — Ising end-to-end (was test_tropical_ising.py) +# ═══════════════════════════════════════════════════════════════════════════════ + + +def _ising_energy_of_cfg(cfg, edges, j_vals, h): + """Ising energy of a {0,1}-config: E = -sum J s_i s_j - sum h s_i (s=1-2cfg).""" + e = 0.0 + for (i, j), jij in zip(edges, j_vals): + si, sj = 1 - 2 * cfg[i], 1 - 2 * cfg[j] + e -= jij * si * sj + for i, hi in zip(range(len(h)), h): + e -= hi * (1 - 2 * cfg[i]) + return e + + +def _maxplus_total_of_assignment(cfg_sym, input_sets, raw_tensors): + """Max-plus total of a symbol->value assignment = sum of tensor entries = + -E for that configuration.""" + total = 0.0 + for term, t in zip(input_sets, raw_tensors): + idx = tuple(int(cfg_sym[c]) for c in term) + total += float(np.asarray(t)[idx]) + return total + + +def test_ising_ring_ground_state_matches_bruteforce(): + spins = [0, 1, 2, 3, 4] + edges = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 0)] + J = [1.0, -1.0, 1.0, 1.0, -1.0] + h = [0.5, -0.3, 0.2, 0.0, 0.4] + nodes = build_ising_tn(spins, edges, J, h) + + e_ground = brute_force_energy(spins, edges, J, h) + + with tropical(): + val = np.array(cons.contractor(nodes, output_edge_order=[]).tensor).item() + + np.testing.assert_allclose(val, -e_ground, atol=1e-6) + + +def test_ising_config_recovery_five_ring(): + spins = [0, 1, 2, 3, 4] + edges = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 0)] + J = [1.0, -1.0, 1.0, 1.0, -1.0] + h = [0.5, -0.3, 0.2, 0.0, 0.4] + + best_e = None + for cfg in itertools.product([0, 1], repeat=len(spins)): + e = _ising_energy_of_cfg(cfg, edges, J, h) + if best_e is None or e < best_e - 1e-9: + best_e = e + + nodes = build_ising_tn(spins, edges, J, h) + with tropical(track=True): + val = np.array(cons.contractor(nodes, output_edge_order=[]).tensor).item() + cfg_sym = recover_configuration() + + np.testing.assert_allclose(val, -best_e, atol=1e-6) + + _tree, input_sets, raw_tensors = get_recorded_topology() + assert input_sets is not None and raw_tensors is not None + all_labels = set().union(*[set(t) for t in input_sets]) + assert set(cfg_sym.keys()) == all_labels + + recovered_total = _maxplus_total_of_assignment(cfg_sym, input_sets, raw_tensors) + np.testing.assert_allclose(recovered_total, val, atol=1e-6) + np.testing.assert_allclose(recovered_total, -best_e, atol=1e-6) + + spin_cfg = {} + for term, _t in zip(input_sets, raw_tensors): + if len(term) == 1: + spin_cfg[len(spin_cfg)] = int(cfg_sym[term[0]]) + assert len(spin_cfg) == len(spins) + recovered_e = _ising_energy_of_cfg( + [spin_cfg[i] for i in range(len(spins))], edges, J, h + ) + np.testing.assert_allclose(recovered_e, best_e, atol=1e-6) + + +def test_ising_config_recovery_degenerate_ring(): + spins = [0, 1, 2, 3, 4] + edges = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 0)] + J = [1.0, 1.0, 1.0, 1.0, 1.0] + h = [0.0, 0.0, 0.0, 0.0, 0.0] + + best_e = None + for cfg in itertools.product([0, 1], repeat=len(spins)): + e = _ising_energy_of_cfg(cfg, edges, J, h) + if best_e is None or e < best_e - 1e-9: + best_e = e + assert best_e == -5.0 + + nodes = build_ising_tn(spins, edges, J, h) + with tropical(track=True): + val = np.array(cons.contractor(nodes, output_edge_order=[]).tensor).item() + cfg_sym = recover_configuration() + + np.testing.assert_allclose(val, -best_e, atol=1e-6) + _tree, input_sets, raw_tensors = get_recorded_topology() + recovered_total = _maxplus_total_of_assignment(cfg_sym, input_sets, raw_tensors) + np.testing.assert_allclose(recovered_total, -best_e, atol=1e-6) + + spin_cfg = {} + for term, _t in zip(input_sets, raw_tensors): + if len(term) == 1: + spin_cfg[len(spin_cfg)] = int(cfg_sym[term[0]]) + recovered_e = _ising_energy_of_cfg( + [spin_cfg[i] for i in range(len(spins))], edges, J, h + ) + np.testing.assert_allclose(recovered_e, best_e, atol=1e-6) + + +def test_ising_counting_ground_state_and_degeneracy(): + spins = [0, 1, 2, 3, 4] + edges = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 0)] + J = [1.0, -1.0, 1.0, 1.0, -1.0] + h = [0.5, -0.3, 0.2, 0.0, 0.4] + best_e, deg = brute_force_energy_and_degeneracy(len(spins), edges, J, h) + + nodes = build_ising_tn(spins, edges, J, h) + + with counting_tropical(): + res = cons.contractor(nodes, output_edge_order=[], ignore_edge_order=True) + count = degeneracy() + + energy = np.array(res.tensor) + np.testing.assert_allclose(energy, -best_e, atol=1e-6) + np.testing.assert_allclose(np.array(count), deg, atol=1e-6) + + +def test_ising_counting_degenerate_ring(): + spins = [0, 1, 2, 3, 4] + edges = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 0)] + J = [1.0, 1.0, 1.0, 1.0, 1.0] + h = [0.0, 0.0, 0.0, 0.0, 0.0] + best_e, deg = brute_force_energy_and_degeneracy(len(spins), edges, J, h) + + assert best_e == -5.0 + assert deg == 2 + + nodes = build_ising_tn(spins, edges, J, h) + + with counting_tropical(): + res = cons.contractor(nodes, output_edge_order=[], ignore_edge_order=True) + count = degeneracy() + + energy = np.array(res.tensor) + np.testing.assert_allclose(energy, -best_e, atol=1e-6) + np.testing.assert_allclose(np.array(count), deg, atol=1e-6)