From a1bb7ff7d7d038e3b80171b60bf95838cfb792e8 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Mon, 27 Jul 2026 13:15:09 -0600 Subject: [PATCH 01/70] Fix fermion operator gate cache keys --- src/pepsy/tensors/symmetric.py | 44 +++++++++++- tests/test_fermion_gate_cache.py | 114 +++++++++++++++++++++++++++++++ 2 files changed, 157 insertions(+), 1 deletion(-) create mode 100644 tests/test_fermion_gate_cache.py diff --git a/src/pepsy/tensors/symmetric.py b/src/pepsy/tensors/symmetric.py index b5dbf7f..ecdce1b 100644 --- a/src/pepsy/tensors/symmetric.py +++ b/src/pepsy/tensors/symmetric.py @@ -2,6 +2,7 @@ from __future__ import annotations +import hashlib import warnings from collections.abc import Mapping from dataclasses import dataclass, field @@ -4865,6 +4866,41 @@ def _apply_to_array_blocks(value, to_backend): return to_backend(value) +def _operator_content_fingerprint(operator): + """Return a stable content digest for a native Symmray operator. + + ``Fermion.operator_gate`` caches gate exponentials, so the cache key must + depend on the operator's *contents* rather than its Python ``id``. A + freshly built operator can be garbage collected and have its memory + address reused, which would otherwise let the cache return a stale gate + for an unrelated operator. Returns ``None`` when a stable fingerprint + cannot be formed -- an unrecognised operator type or blocks that carry an + autodiff graph -- signalling that the resulting gate must not be cached. + """ + blocks = getattr(operator, "blocks", None) + if blocks is None: + return None + hasher = hashlib.blake2b(digest_size=16) + header = ( + getattr(operator, "symmetry", None), + getattr(operator, "charge", None), + tuple(getattr(operator, "duals", ()) or ()), + tuple(getattr(operator, "shape", ()) or ()), + ) + hasher.update(repr(header).encode()) + for sector in sorted(blocks, key=repr): + block = blocks[sector] + if getattr(block, "requires_grad", False): + return None + try: + array = np.ascontiguousarray(ar.to_numpy(block)) + except Exception: # pragma: no cover - defensive backend guard + return None + hasher.update(repr((sector, array.shape, array.dtype.str)).encode()) + hasher.update(array.tobytes()) + return hasher.hexdigest() + + def _apply_to_tensor_network_arrays(tn, to_backend): if to_backend is None: return tn @@ -8224,7 +8260,7 @@ def operator_gate(self, operator, theta, *, imaginary=False): operator_key = ("name", name) else: factory = lambda: operator - operator_key = ("term", id(operator)) + operator_key = ("term", _operator_content_fingerprint(operator)) def build(): term = factory() @@ -8237,6 +8273,12 @@ def build(): gate = _gate_from_term(term, theta, imaginary=imaginary) return _apply_to_array_blocks(gate, self.to_backend) + if operator_key[1] is None: + # No stable content fingerprint (an unrecognised operator type or + # autodiff tensors): build without caching so that a recycled + # Python ``id`` can never alias an unrelated operator's gate. + return build() + return self._cached_gate( ("operator", operator_key, theta, imaginary), build, diff --git a/tests/test_fermion_gate_cache.py b/tests/test_fermion_gate_cache.py new file mode 100644 index 0000000..3d6e143 --- /dev/null +++ b/tests/test_fermion_gate_cache.py @@ -0,0 +1,114 @@ +"""Regression tests for content-addressed ``Fermion.operator_gate`` caching. + +``operator_gate`` memoises gate exponentials. Historically the cache key for a +raw (already-built) operator used ``id(operator)``. A freshly built operator +can be garbage collected and have its memory address reused, so a later, +unrelated operator could alias a stale cache entry and return the wrong gate. +These tests pin the content-addressed behaviour that fixes that. +""" + +from __future__ import annotations + +import gc + +import numpy as np +import pytest + +from pepsy.tensors.symmetric import ( + Fermion, + _gate_from_term, + _operator_content_fingerprint, +) + + +def _dense(gate): + return np.asarray(gate.to_dense()) + + +def test_fingerprint_is_content_addressed_not_identity(): + """Equal contents share a fingerprint; different contents do not.""" + fermion = Fermion(spinful=True, symmetry="U1", t=1.0, U=8.0) + + up_a = fermion.hopping_operator(spin="up") + up_b = fermion.hopping_operator(spin="up") + down = fermion.hopping_operator(spin="down") + inter = fermion.interaction_operator() + + # Freshly built objects (distinct ids) with identical contents must agree. + assert up_a is not up_b + assert _operator_content_fingerprint(up_a) == _operator_content_fingerprint(up_b) + + # Different operators must produce different fingerprints. + assert _operator_content_fingerprint(up_a) != _operator_content_fingerprint(down) + assert _operator_content_fingerprint(up_a) != _operator_content_fingerprint(inter) + + # A non-operator object cannot be fingerprinted and must not be cached. + assert _operator_content_fingerprint(object()) is None + + +def test_operator_gate_does_not_alias_distinct_operators_under_id_reuse(): + """Interleaving fresh operators at a fixed angle must stay correct. + + This reproduces the real failure mode: building many short-lived operators + of different kinds at the same ``theta`` used to let a recycled ``id`` return + a stale cached gate. Every gate must match an independent exponential. + """ + fermion = Fermion(spinful=True, symmetry="U1", t=1.0, U=8.0) + theta = 0.10667747 + + references = { + "up": _dense(_gate_from_term(fermion.hopping_operator(spin="up"), theta)), + "down": _dense(_gate_from_term(fermion.hopping_operator(spin="down"), theta)), + "inter": _dense(_gate_from_term(fermion.interaction_operator(), theta)), + } + + for _ in range(64): + for kind, ref in references.items(): + if kind == "up": + operator = fermion.hopping_operator(spin="up") + elif kind == "down": + operator = fermion.hopping_operator(spin="down") + else: + operator = fermion.interaction_operator() + gate = _dense(fermion.operator_gate(operator, theta)) + np.testing.assert_allclose(gate, ref, atol=1e-12) + del operator + gc.collect() # encourage id recycling between iterations + + +def test_operator_gate_cache_hits_return_equivalent_gate(): + """Repeated calls with equal contents reuse a single correct gate.""" + fermion = Fermion(spinful=True, symmetry="U1", t=1.0, U=8.0) + theta = 0.37 + + reference = _dense(_gate_from_term(fermion.hopping_operator(spin="up"), theta)) + first = fermion.operator_gate(fermion.hopping_operator(spin="up"), theta) + second = fermion.operator_gate(fermion.hopping_operator(spin="up"), theta) + + # Content-addressed cache: equal contents collapse to one stored gate. + assert first is second + np.testing.assert_allclose(_dense(first), reference, atol=1e-12) + + +def test_operator_gate_skips_cache_for_autodiff_operators(): + """Blocks that carry an autodiff graph must not be fingerprinted/cached.""" + torch = pytest.importorskip("torch") + + class _StubOperator: + symmetry = "U1" + charge = 0 + duals = (False, True) + shape = (1, 1) + + def __init__(self, blocks): + self.blocks = blocks + + grad_tracked = _StubOperator( + {(0, 0): torch.zeros((1, 1), dtype=torch.complex128, requires_grad=True)} + ) + assert _operator_content_fingerprint(grad_tracked) is None + + plain = _StubOperator( + {(0, 0): torch.zeros((1, 1), dtype=torch.complex128)} + ) + assert _operator_content_fingerprint(plain) is not None From e6009ec7c51b414866a34ebeec1c9e658f2f8a18 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Mon, 27 Jul 2026 13:35:38 -0600 Subject: [PATCH 02/70] Optimize SymDMRG compiled block fanout --- docs/api/optimizers/sym_dmrg.md | 18 +- .../plans/symdmrg_matvec_fanout.md | 33 ++- src/pepsy/optimizers/sym_dmrg.py | 191 +++++++++++++++++- tests/test_symmetric_tensors.py | 30 ++- 4 files changed, 259 insertions(+), 13 deletions(-) diff --git a/docs/api/optimizers/sym_dmrg.md b/docs/api/optimizers/sym_dmrg.md index d6c7696..0c144ce 100644 --- a/docs/api/optimizers/sym_dmrg.md +++ b/docs/api/optimizers/sym_dmrg.md @@ -82,13 +82,25 @@ and groups equal `(M, K, N)` products into batched `numpy.matmul` calls. This avoids rediscovering block routing and collapses repeated small dense products inside every hot-loop `H_eff` application. The normal Fermi-Hubbard DMRG path is already bosonized before these contractions, so it uses this fast path. -Compatible native fermionic arrays can also use it for unfused, shared-leg, -NumPy-only contractions: Pepsy caches Symmray's sector phases with the plan and +For unfused bosonic NumPy plans, output blocks with an identical dynamic right +source schedule can additionally use a source-fanout GEMM: Pepsy stacks their +static left maps even when their row count `M` differs, builds the shared right +matrix once, and scatters the GEMM rows back to their sector blocks. The +additional stacked maps are collectively capped at 32 MiB per compiled pair; +all unclaimed outputs retain the existing batched or single-matmul routes. +Compatible native fermionic arrays can still use the compiled route for +unfused, shared-leg, NumPy-only contractions, but retain their existing +per-output batching: Pepsy caches Symmray's sector phases with the plan and checks fermionic metadata before each reuse. Fused, outer-product, mixed, or non-NumPy fermionic contractions retain Symmray's exact `tensordot` path. `profile_summary()` aggregates the new batch timing phases, while sampled matvec diagnostics report the batch-plan shape and call counters so scale runs -can confirm the hot matvec path is reusing this setup work. +can confirm the hot matvec path is reusing this setup work. Fanout diagnostics +use the `*_compiled_block_plan_fanout_*` prefix and report eligible/enabled +groups, output coverage, static bytes, predicted output-product savings, and +actual fanout GEMM calls; timing totals use +`*_compiled_block_fanout_pack_elapsed` and +`*_compiled_block_fanout_matmul_elapsed`. For small, right-first, bosonic projected problems, Pepsy also considers a private dense block-sector effective-Hamiltonian cache. It is capped at 32 MiB and is retained only after its first result agrees with the streamed diff --git a/docs/development/plans/symdmrg_matvec_fanout.md b/docs/development/plans/symdmrg_matvec_fanout.md index 7690dd4..78f51df 100644 --- a/docs/development/plans/symdmrg_matvec_fanout.md +++ b/docs/development/plans/symdmrg_matvec_fanout.md @@ -2,8 +2,11 @@ ## Status -Planned. This is the next SymDMRG2 performance experiment after compiled -same-shape batching and the bounded sector-operator prototype. +Implemented and retained after a warmed 6x6 PBC chi=64 hot-loop A/B. The +compiled path now builds bounded source-fanout groups for eligible bosonic, +unfused NumPy plans and reports the census, static storage, predicted savings, +and actual GEMM-call diagnostics. Native-fermionic, fused, outer-product, and +metadata-mismatch paths retain their existing fallbacks. ## Evidence and decision @@ -119,6 +122,32 @@ do not introduce an unmeasured duplicate-matrix footprint. Otherwise retain the current batching and move the next investigation to the dominant SVD or local-solver phase. +### 2026-07-27 result + +On the actual 6x6 PBC U/t=8 construction, a fixed `chi=64` block-fill state +at the central projected window was warmed once and then evaluated 100 times +with single-threaded BLAS. The paired compiled right/left contractions fell +from 2.2492 s (pre-change) to 2.0618 s (fanout), an 8.3% reduction; total +measured application time fell from 2.3156 s to 2.1279 s. The active right +contraction formed 1,481 fanout groups covering 3,841 output blocks, with +1,271,840 bytes of added static maps and 2,360 predicted output-product +savings per application. + +A matching one-sweep `chi=64` solver control retained its energy to +`4.9e-13` and used the same 280 Lanczos matvecs. Its end-to-end wall time was +not used as a throughput claim because it includes cold projected-plan setup +and shared-host noise. + +The required full control used the same 6x6 PBC product ramp, density-matrix +mixer, `variational_sector_basis="off"`, seed, and single-threaded BLAS for 30 +sweeps. It produced indistinguishable energy (`-16.547808951298467` before, +`-16.547808951298432` after) and the same 7,533 Lanczos matvecs. Aggregated +compiled left/right contraction time fell from 43.3675 s to 41.4009 s (4.5%). +End-to-end wall time was effectively tied (198.6 s before, 199.3 s after), so +this is a retained hot-loop optimization rather than a claim of full-solver +speedup. Additional scale runs should report the fanout timing fields +separately from plan-build and total wall time. + ## Test commands ```bash diff --git a/src/pepsy/optimizers/sym_dmrg.py b/src/pepsy/optimizers/sym_dmrg.py index cdc37b9..bf19e21 100644 --- a/src/pepsy/optimizers/sym_dmrg.py +++ b/src/pepsy/optimizers/sym_dmrg.py @@ -27,6 +27,10 @@ # the streamed contraction path; this is a correctness-first cache, not an # unbounded dense-local solver. _SECTOR_OPERATOR_MAX_BYTES = 32 * 1024**2 +# A fanout plan retains a stacked copy of static left maps. Keep the extra +# plan-only storage bounded; the original per-output maps remain available for +# sector-operator diagnostics and for clear plan rebuilding semantics. +_FANOUT_MAX_STATIC_BYTES = 32 * 1024**2 def _is_symmray_array(value): @@ -807,6 +811,7 @@ def __init__(self, optimizer, left, right_inds, *, layout="unfused"): self.left_output_axes = tuple(range(len(self.left_inds))) self.right_output_axes = tuple(range(len(self.right_inds))) self.compiled_block_plan = None + self.compiled_block_plan_fanouts = () self.compiled_block_plan_batches = () self.compiled_block_plan_singles = () self.compiled_right_layout = None @@ -818,6 +823,15 @@ def __init__(self, optimizer, left, right_inds, *, layout="unfused"): self.compiled_block_plan_uses = 0 self.compiled_block_plan_terms = 0 self.compiled_block_plan_output_blocks = 0 + self.compiled_block_plan_fanout_eligible_groups = 0 + self.compiled_block_plan_fanout_eligible_output_blocks = 0 + self.compiled_block_plan_fanout_eligible_static_bytes = 0 + self.compiled_block_plan_fanout_groups = 0 + self.compiled_block_plan_fanout_output_blocks = 0 + self.compiled_block_plan_fanout_max_size = 0 + self.compiled_block_plan_fanout_static_bytes = 0 + self.compiled_block_plan_fanout_predicted_matmul_savings = 0 + self.compiled_block_plan_fanout_matmul_calls = 0 self.compiled_block_plan_batch_groups = 0 self.compiled_block_plan_batched_output_blocks = 0 self.compiled_block_plan_max_batch_size = 0 @@ -1005,6 +1019,7 @@ def _compile_block_plan(self, right, output): # Clear it before attempting a replacement so a failed compilation # cannot leave a stale plan eligible for the next matvec. self.compiled_block_plan = None + self.compiled_block_plan_fanouts = () self.compiled_block_plan_batches = () self.compiled_block_plan_singles = () self.compiled_right_layout = None @@ -1140,8 +1155,81 @@ def _compile_block_plan(self, right, output): ) ) + # A fanout group shares one dynamically assembled right matrix across + # static left maps of potentially different row counts. This is unlike + # the batch route below: a stacked 2D GEMM gives BLAS one ordinary + # matrix product rather than asking batched matmul to dispatch each + # small output independently. Native fermions, fused layouts, and + # outer products deliberately retain their existing paths. + fanout_candidates = {} + if self.shared and not fermionic and self.layout == "unfused": + for index, (sector, _, left_matrix, right_specs) in enumerate( + compiled_plan + ): + right_dtype = np.result_type( + *( + _block_dtype(right.data.blocks[right_sector]) + for right_sector, _, _, _ in right_specs + ) + ) + template_dtype = _block_dtype(output.data.blocks[sector]) + # ``right_specs`` contains the exact source schedule, + # including its offsets and reduced N. Deliberately omit M: + # it is represented by each group's row slices instead. + key = ( + right_specs, + left_matrix.shape[1], + left_matrix.dtype.str, + right_dtype.str, + template_dtype.str, + ) + fanout_candidates.setdefault(key, []).append(index) + + eligible_fanouts = tuple( + tuple(indices) + for indices in fanout_candidates.values() + if len(indices) > 1 + ) + self.compiled_block_plan_fanout_eligible_groups = len(eligible_fanouts) + self.compiled_block_plan_fanout_eligible_output_blocks = sum( + len(indices) for indices in eligible_fanouts + ) + self.compiled_block_plan_fanout_eligible_static_bytes = sum( + sum(compiled_plan[index][2].nbytes for index in indices) + for indices in eligible_fanouts + ) + + fanouts = [] + fanout_claimed = set() + fanout_static_bytes = 0 + for indices in eligible_fanouts: + static_bytes = sum(compiled_plan[index][2].nbytes for index in indices) + if fanout_static_bytes + static_bytes > _FANOUT_MAX_STATIC_BYTES: + continue + row_start = 0 + row_slices = [] + for index in indices: + rows = compiled_plan[index][2].shape[0] + row_slices.append((row_start, row_start + rows)) + row_start += rows + fanouts.append( + ( + indices, + np.concatenate( + tuple(compiled_plan[index][2] for index in indices), + axis=0, + ), + tuple(row_slices), + compiled_plan[indices[0]][3], + ) + ) + fanout_claimed.update(indices) + fanout_static_bytes += static_bytes + group_members = {} for index, (sector, _, left_matrix, right_specs) in enumerate(compiled_plan): + if index in fanout_claimed: + continue right_dtype = np.result_type( *( _block_dtype(right.data.blocks[right_sector]) @@ -1219,6 +1307,7 @@ def _compile_block_plan(self, right, output): self.compiled_output_template = output.data self.compiled_right_layout = _block_data_layout_map(right.data) self.compiled_block_plan = tuple(compiled_plan) + self.compiled_block_plan_fanouts = tuple(fanouts) self.compiled_block_plan_batches = tuple(batches) self.compiled_block_plan_singles = tuple(singles) self.compiled_block_plan_fermionic = fermionic @@ -1233,6 +1322,23 @@ def _compile_block_plan(self, right, output): self.compiled_block_plan_builds += 1 self.compiled_block_plan_terms = int(num_terms) self.compiled_block_plan_output_blocks = len(self.compiled_block_plan) + self.compiled_block_plan_fanout_groups = len(self.compiled_block_plan_fanouts) + self.compiled_block_plan_fanout_output_blocks = sum( + len(indices) + for indices, _, _, _ in self.compiled_block_plan_fanouts + ) + self.compiled_block_plan_fanout_max_size = max( + ( + len(indices) + for indices, _, _, _ in self.compiled_block_plan_fanouts + ), + default=0, + ) + self.compiled_block_plan_fanout_static_bytes = fanout_static_bytes + self.compiled_block_plan_fanout_predicted_matmul_savings = sum( + len(indices) - 1 + for indices, _, _, _ in self.compiled_block_plan_fanouts + ) self.compiled_block_plan_batch_groups = len(self.compiled_block_plan_batches) self.compiled_block_plan_batched_output_blocks = sum( len(indices) for indices, _, _, _, _ in self.compiled_block_plan_batches @@ -1245,9 +1351,13 @@ def _compile_block_plan(self, right, output): default=0, ) self.compiled_block_plan_mode = ( - "output_block_batched_matmul" - if self.compiled_block_plan_batches - else "output_block_matmul" + "output_block_fanout_gemm" + if self.compiled_block_plan_fanouts + else ( + "output_block_batched_matmul" + if self.compiled_block_plan_batches + else "output_block_matmul" + ) ) self.compiled_block_plan_disabled_reason = None @@ -1276,6 +1386,54 @@ def get_right_matrix(right_sector, shared_size, right_output_size): right_matrix_cache[cache_key] = right_matrix return right_matrix + for ( + indices, + left_stack, + row_slices, + right_specs, + ) in self.compiled_block_plan_fanouts: + fanout_start = time.perf_counter() if timings is not None else None + right_matrices = tuple( + get_right_matrix( + right_sector, + term_shared_size, + term_output_size, + ) + for ( + right_sector, + _, + term_shared_size, + term_output_size, + ) in right_specs + ) + right_matrix = ( + right_matrices[0] + if len(right_matrices) == 1 + else np.concatenate(right_matrices, axis=0) + ) + _add_elapsed( + timings, + f"{prefix}_compiled_block_fanout_pack_elapsed", + fanout_start, + ) + matmul_start = time.perf_counter() if timings is not None else None + out_stack = left_stack @ right_matrix + _add_elapsed( + timings, + f"{prefix}_compiled_block_fanout_matmul_elapsed", + matmul_start, + ) + for plan_index, (row_start, row_stop) in zip(indices, row_slices): + output_sector, output_shape, _, _ = self.compiled_block_plan[ + plan_index + ] + out = out_stack[row_start:row_stop] + template_dtype = _block_dtype(template_blocks[output_sector]) + if out.dtype != template_dtype: + out = np.asarray(out, dtype=template_dtype) + blocks[output_sector] = out.reshape(output_shape) + self.compiled_block_plan_fanout_matmul_calls += 1 + for ( indices, left_batch, @@ -1499,6 +1657,33 @@ def summary(self, prefix): f"{prefix}_compiled_block_plan_output_blocks": int( self.compiled_block_plan_output_blocks ), + f"{prefix}_compiled_block_plan_fanout_eligible_groups": int( + self.compiled_block_plan_fanout_eligible_groups + ), + f"{prefix}_compiled_block_plan_fanout_eligible_output_blocks": int( + self.compiled_block_plan_fanout_eligible_output_blocks + ), + f"{prefix}_compiled_block_plan_fanout_eligible_static_bytes": int( + self.compiled_block_plan_fanout_eligible_static_bytes + ), + f"{prefix}_compiled_block_plan_fanout_groups": int( + self.compiled_block_plan_fanout_groups + ), + f"{prefix}_compiled_block_plan_fanout_output_blocks": int( + self.compiled_block_plan_fanout_output_blocks + ), + f"{prefix}_compiled_block_plan_fanout_max_size": int( + self.compiled_block_plan_fanout_max_size + ), + f"{prefix}_compiled_block_plan_fanout_static_bytes": int( + self.compiled_block_plan_fanout_static_bytes + ), + f"{prefix}_compiled_block_plan_fanout_predicted_matmul_savings": int( + self.compiled_block_plan_fanout_predicted_matmul_savings + ), + f"{prefix}_compiled_block_plan_fanout_matmul_calls": int( + self.compiled_block_plan_fanout_matmul_calls + ), f"{prefix}_compiled_block_plan_batch_groups": int( self.compiled_block_plan_batch_groups ), diff --git a/tests/test_symmetric_tensors.py b/tests/test_symmetric_tensors.py index 7e59932..761a60b 100644 --- a/tests/test_symmetric_tensors.py +++ b/tests/test_symmetric_tensors.py @@ -407,8 +407,8 @@ def _randomized_block_tensor(tensor, seed): return _tensor_with_data(tensor, data) -def test_symdmrg_compiled_batched_plan_matches_blockwise_fh_matvec(): - """The actual FH DMRG path is bosonic and batches repeated output shapes.""" +def test_symdmrg_compiled_fanout_plan_matches_blockwise_fh_matvec(): + """The bosonized FH path reuses one right matrix through fanout GEMM.""" mapper = OneDMap(3, 2, mode="snake") edges = tuple(qtn.edges_2d_square(3, 2, cyclic=True)) mpo = SymHamiltonian.from_edges( @@ -456,14 +456,33 @@ def test_symdmrg_compiled_batched_plan_matches_blockwise_fh_matvec(): got = contraction.apply(random_theta) assert contraction.compiled_block_plan_uses == 1 - assert contraction.compiled_block_plan_batch_groups > 0 - assert contraction.compiled_block_plan_batched_output_blocks > 1 - assert contraction.compiled_block_plan_batched_matmul_calls > 0 + assert contraction.compiled_block_plan_fanout_eligible_groups > 0 + assert contraction.compiled_block_plan_fanout_groups > 0 + assert contraction.compiled_block_plan_fanout_output_blocks > 1 + assert contraction.compiled_block_plan_fanout_static_bytes > 0 + assert contraction.compiled_block_plan_fanout_predicted_matmul_savings > 0 + assert contraction.compiled_block_plan_fanout_matmul_calls > 0 + assert any( + len({row_stop - row_start for row_start, row_stop in row_slices}) > 1 + for _, _, row_slices, _ in contraction.compiled_block_plan_fanouts + ) for sector, expected in reference.data.blocks.items(): np.testing.assert_allclose( got.data.blocks[sector], expected, atol=1e-12, rtol=1e-12 ) + # A compatible layout must reuse the same compiled fanout plan. + fanout_calls_before = contraction.compiled_block_plan_fanout_matmul_calls + repeated = contraction.apply(random_theta) + assert contraction.compiled_block_plan_uses == 2 + assert contraction.compiled_block_plan_fanout_matmul_calls == ( + fanout_calls_before + contraction.compiled_block_plan_fanout_groups + ) + for sector, expected in got.data.blocks.items(): + np.testing.assert_allclose( + repeated.data.blocks[sector], expected, atol=1e-12, rtol=1e-12 + ) + # This window's equivalent composed map changes the accumulation order by # slightly more than the strict validation tolerance, so it must retain # the streamed plan. This is the guard against numerical regressions. @@ -536,6 +555,7 @@ def test_native_fermionic_compiled_plan_preserves_phases_and_dummy_modes(): contraction = _BlockPairContraction(None, left, right.inds) contraction.apply(right) # Direct reference call that compiles the plan. assert contraction.compiled_block_plan_fermionic + assert contraction.compiled_block_plan_fanout_groups == 0 random_right = _randomized_block_tensor(right, 17) reference = _BlockPairContraction(None, left, right.inds).apply(random_right) From fd5f383a824c141cfbb0478109db4d913e6729b7 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Mon, 27 Jul 2026 14:06:51 -0600 Subject: [PATCH 03/70] Avoid unamortized SymDMRG sector cache setup --- docs/api/optimizers/sym_dmrg.md | 14 +++++------ src/pepsy/optimizers/sym_dmrg.py | 13 ++++++---- tests/test_symmetric_tensors.py | 41 +++++++++++++++++++++++++++++++- 3 files changed, 54 insertions(+), 14 deletions(-) diff --git a/docs/api/optimizers/sym_dmrg.md b/docs/api/optimizers/sym_dmrg.md index 0c144ce..6997bee 100644 --- a/docs/api/optimizers/sym_dmrg.md +++ b/docs/api/optimizers/sym_dmrg.md @@ -101,14 +101,12 @@ groups, output coverage, static bytes, predicted output-product savings, and actual fanout GEMM calls; timing totals use `*_compiled_block_fanout_pack_elapsed` and `*_compiled_block_fanout_matmul_elapsed`. -For small, right-first, bosonic projected problems, Pepsy also considers a -private dense block-sector effective-Hamiltonian cache. It is capped at 32 MiB -and is retained only after its first result agrees with the streamed -contractions at `1e-12`; left-first, fermionic, oversized, and failed -validation cases continue to use the compiled streaming plans. The diagnostic -record exposes the cache state, size, block count, reuse count, validation -error, and any disabled reason. This is a bounded experimental optimization -rather than a replacement local solver. +The private dense block-sector effective-Hamiltonian cache is disabled by +default. Its composed matmuls change summation order and did not amortize their +setup cost within bounded local Krylov solves. The experimental path remains +available to focused benchmarks and still requires first-result agreement with +the streamed contractions at `1e-12`; its diagnostic record exposes the cache +state, size, block count, reuse count, validation error, and disabled reason. `matvec_layout="fused"` is available as an opt-in prototype for the block-native path. It attempts to fuse multiple shared contraction legs inside each cached projected problem, using Symmray's fused-index support when the diff --git a/src/pepsy/optimizers/sym_dmrg.py b/src/pepsy/optimizers/sym_dmrg.py index bf19e21..86d4fc5 100644 --- a/src/pepsy/optimizers/sym_dmrg.py +++ b/src/pepsy/optimizers/sym_dmrg.py @@ -22,11 +22,11 @@ from .energy import MpsEnergyOptimizer -# The block-sector effective-Hamiltonian prototype only materializes an -# operator when its complete dense block map is small. Larger windows retain -# the streamed contraction path; this is a correctness-first cache, not an -# unbounded dense-local solver. -_SECTOR_OPERATOR_MAX_BYTES = 32 * 1024**2 +# A composed dense sector operator changes the contraction's summation order +# and costs more to prepare than it saves for the bounded local Krylov solves +# used by SymDMRG2. Keep the experimental path disabled by default, with a +# module-level limit so focused experiments can explicitly re-enable it. +_SECTOR_OPERATOR_MAX_BYTES = 0 # A fanout plan retains a stacked copy of static left maps. Keep the extra # plan-only storage bounded; the original per-output maps remain available for # sector-operator diagnostics and for clear plan rebuilding semantics. @@ -1956,6 +1956,9 @@ def _maybe_compile_sector_operator(self, theta, reference): if self.sector_operator_build_attempted: return self.sector_operator_build_attempted = True + if _SECTOR_OPERATOR_MAX_BYTES <= 0: + self.sector_operator_disabled_reason = "disabled" + return layout = self._sector_operator_layout() if layout is None: if self.sector_operator_disabled_reason is None: diff --git a/tests/test_symmetric_tensors.py b/tests/test_symmetric_tensors.py index 761a60b..e36d2a5 100644 --- a/tests/test_symmetric_tensors.py +++ b/tests/test_symmetric_tensors.py @@ -5,6 +5,7 @@ import quimb.tensor as qtn import pepsy +from pepsy.optimizers import sym_dmrg as sym_dmrg_mod from pepsy.operators import gate, gate_simple from pepsy.optimizers.sym_dmrg import ( _BlockPairContraction, @@ -407,8 +408,9 @@ def _randomized_block_tensor(tensor, seed): return _tensor_with_data(tensor, data) -def test_symdmrg_compiled_fanout_plan_matches_blockwise_fh_matvec(): +def test_symdmrg_compiled_fanout_plan_matches_blockwise_fh_matvec(monkeypatch): """The bosonized FH path reuses one right matrix through fanout GEMM.""" + monkeypatch.setattr(sym_dmrg_mod, "_SECTOR_OPERATOR_MAX_BYTES", 32 * 1024**2) mapper = OneDMap(3, 2, mode="snake") edges = tuple(qtn.edges_2d_square(3, 2, cyclic=True)) mpo = SymHamiltonian.from_edges( @@ -536,6 +538,43 @@ def test_symdmrg_compiled_fanout_plan_matches_blockwise_fh_matvec(): ) +def test_symdmrg_sector_operator_default_bypasses_layout_build(monkeypatch): + """The default streamed matvec does not pay experimental cache setup.""" + mapper = OneDMap(3, 2, mode="snake") + edges = tuple(qtn.edges_2d_square(3, 2, cyclic=True)) + mpo = SymHamiltonian.from_edges( + "fermi_hubbard_u1u1", "U1U1", edges, t=1.0, U=8.0 + ).to_mpo(mapper=mapper, compress=True, cutoff=1e-12) + state = SymMPS.for_model( + "fermi_hubbard_u1u1", + 6, + bond_dim=4, + site_charge=site_charge_from_occupations([(1, 0), (0, 1)] * 3), + seed=3, + dtype="complex128", + ) + optimizer = pepsy.SymDMRG2( + mpo, + state, + bond_dims=[4], + cutoffs=[1e-10], + compute_initial_energy=False, + ) + theta = optimizer.two_site_theta(0) + problem, _ = optimizer._get_projected_problem(0, theta) + monkeypatch.setattr( + problem, + "_sector_operator_layout", + lambda: pytest.fail("default matvec must not inspect the sector cache layout"), + ) + + result = problem.apply(_randomized_block_tensor(theta, 31)) + + assert result.data.blocks + assert problem.sector_operator is None + assert problem.sector_operator_disabled_reason == "disabled" + + def test_native_fermionic_compiled_plan_preserves_phases_and_dummy_modes(): """Direct native contractions compile only with phase-stable metadata.""" state = SymMPS.for_model( From f3a545fa455f45b4b17270a7abac89a823633ce8 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Mon, 27 Jul 2026 14:18:44 -0600 Subject: [PATCH 04/70] Document SymDMRG performance closure --- .../plans/symdmrg_matvec_fanout.md | 32 +++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/docs/development/plans/symdmrg_matvec_fanout.md b/docs/development/plans/symdmrg_matvec_fanout.md index 78f51df..762e409 100644 --- a/docs/development/plans/symdmrg_matvec_fanout.md +++ b/docs/development/plans/symdmrg_matvec_fanout.md @@ -148,6 +148,38 @@ this is a retained hot-loop optimization rather than a claim of full-solver speedup. Additional scale runs should report the fanout timing fields separately from plan-build and total wall time. +### 2026-07-27 follow-up closure + +The subsequent setup/SVD/local-solver pass retained one additional change: +the private composed sector-operator cache is now disabled before it builds a +layout. In the fixed 30-sweep 6x6 replay this preserved the direct trajectory +exactly (`E=-16.52514449042574`, 7,495 Lanczos matvecs) while reducing its +measured setup time from 7.243 s to 0.0069 s and wall time from 185.90 s to +180.23 s. The projected-problem cache itself remains effective (it is reused +within every local Krylov solve); it cannot survive a tensor/environment update +without becoming stale. + +The remaining candidates were measured and rejected rather than expanded: + +- All 70 compiled contraction layouts in a representative chi=64 sweep were + distinct, so cross-window topology caching had zero reuse. +- `matvec_layout="fused"` agreed numerically but added fuse work and was 6.2% + slower on the same seeded sweep (7.83 s versus 7.38 s). +- Normalizing static left matrices to contiguous NumPy storage also agreed + numerically but cost more to prepare than BLAS recovered (7.55 s versus + 7.20 s). +- Reducing native Lanczos from `ncv=8` to `ncv=7` reduced the full-run matvec + count (7,495 to 6,633) but shifted the final energy by 5.03e-3. +- Disabling the density-matrix mixer after the chi ramp reduced wall time + (162.11 s) and SVD time (6.08 s), but shifted the final energy by 4.43e-4. + +Those changes exceed the solver-tolerance acceptance criterion. The retained +unfused fanout matvec, allocation-free sector-cache bypass, density-matrix +mixer duration, and `ncv=8` are therefore the current 6x6 configuration. Any +future performance work needs a new algorithmic approach (for example a +validated block-sparse density-matrix reduction kernel), not another caching or +parameter tweak. + ## Test commands ```bash From 0b22d90d170233ddfc7b4cb22a711acca5adbb0e Mon Sep 17 00:00:00 2001 From: rezaquant Date: Mon, 27 Jul 2026 14:33:06 -0600 Subject: [PATCH 05/70] Scope SymDMRG performance findings by bond dimension --- .../plans/symdmrg_matvec_fanout.md | 33 +++++++++++++++---- 1 file changed, 26 insertions(+), 7 deletions(-) diff --git a/docs/development/plans/symdmrg_matvec_fanout.md b/docs/development/plans/symdmrg_matvec_fanout.md index 762e409..5c996ee 100644 --- a/docs/development/plans/symdmrg_matvec_fanout.md +++ b/docs/development/plans/symdmrg_matvec_fanout.md @@ -159,7 +159,9 @@ measured setup time from 7.243 s to 0.0069 s and wall time from 185.90 s to within every local Krylov solve); it cannot survive a tensor/environment update without becoming stale. -The remaining candidates were measured and rejected rather than expanded: +The following throughput measurements are scoped to the warmed, single-threaded +CPU chi=64 control. They reject these variants for the overhead-bound regime, +not for arbitrarily large bond dimensions: - All 70 compiled contraction layouts in a representative chi=64 sweep were distinct, so cross-window topology caching had zero reuse. @@ -173,12 +175,29 @@ The remaining candidates were measured and rejected rather than expanded: - Disabling the density-matrix mixer after the chi ramp reduced wall time (162.11 s) and SVD time (6.08 s), but shifted the final energy by 4.43e-4. -Those changes exceed the solver-tolerance acceptance criterion. The retained -unfused fanout matvec, allocation-free sector-cache bypass, density-matrix -mixer duration, and `ncv=8` are therefore the current 6x6 configuration. Any -future performance work needs a new algorithmic approach (for example a -validated block-sparse density-matrix reduction kernel), not another caching or -parameter tweak. +The `ncv=7` and earlier-mixer rejections are numerical/basin decisions, rather +than chi=64 throughput judgments; retain `ncv=8` and the current mixer duration +unless a new full solver-tolerance study proves otherwise. The allocation-free +sector-cache bypass is likewise safe at any chi. Cross-window topology reuse +is structurally unhelpful in this control (zero hits) and becomes relatively +less important as block GEMM work grows. + +Do not extrapolate the fused-layout or contiguous-static-matrix losses to high +chi. Their extra packing/copy work lost at chi=64, where the measured kernels +are overhead-bound, but larger block products may become BLAS-bound and reverse +that tradeoff. Before declaring this route closed for chi >= 512, run held-sweep +A/Bs for unfused versus fused routing and native versus contiguous static +matrices, under both one-thread and appropriately sized multi-thread BLAS. +Require the direct-blockwise `1e-12` comparison, solver-tolerance agreement, +and separate compiled-matvec and wall timings. Audit fanout static-map storage +at each scale as well: its chi=64 footprint is not a safe predictor of the +block-sector distribution or memory pressure at chi=4096. + +The retained unfused fanout matvec, allocation-free sector-cache bypass, +density-matrix mixer duration, and `ncv=8` are therefore the current 6x6 +chi=64 configuration. At high chi, the next plausible work remains GEMM +throughput and a validated block-sparse density-matrix reduction kernel; the +6x6 control spent about 17.5 s of 180.2 s in SVD splits, not 0.1%. ## Test commands From ffe104f7ef25de773183ddcdc3fa28d7968e51c6 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Mon, 27 Jul 2026 18:17:01 -0600 Subject: [PATCH 06/70] feat(vmc): add reusable native measurement workflow --- .github/skills/pepsy-vmc/SKILL.md | 10 +- docs/api/vmc.md | 111 ++++- src/pepsy/vmc/__init__.py | 7 + src/pepsy/vmc/api.py | 10 +- src/pepsy/vmc/torch/__init__.py | 8 + src/pepsy/vmc/torch/_core.py | 8 + src/pepsy/vmc/torch/driver.py | 180 ++++++-- src/pepsy/vmc/torch/fermion.py | 665 +++++++++++++++++++++++++++--- src/pepsy/vmc/torch/importance.py | 78 +++- src/pepsy/vmc/torch/proposals.py | 2 +- src/pepsy/vmc/torch/results.py | 274 +++++++++++- src/pepsy/vmc/torch/sampler.py | 63 ++- tests/test_vmc_api.py | 433 +++++++++++++++++++ tests/test_vmc_importance.py | 45 +- 14 files changed, 1763 insertions(+), 131 deletions(-) diff --git a/.github/skills/pepsy-vmc/SKILL.md b/.github/skills/pepsy-vmc/SKILL.md index 92119e0..3210ab5 100644 --- a/.github/skills/pepsy-vmc/SKILL.md +++ b/.github/skills/pepsy-vmc/SKILL.md @@ -196,9 +196,13 @@ For sampled configurations `x`, evaluate `O_loc(x) = sum_{x'} O[x, x'] * psi(x') / psi(x)`. -Average local estimators over the Markov-chain samples. The canonical sampled -driver method is `estimate_observable(...)`; `estimate_energy(...)` remains a -compatibility wrapper. Return complex means +Average local estimators over the Markov-chain samples. The canonical native +Torch workflow is `samples = vmc.sample(sampling=...)`, followed by +`vmc.measure(samples, observables={...})`. One `measure` call shares stored +parent amplitudes and connected-target contractions across every named +observable. `run(observables=..., sampling=...)` is the convenience form that +warms up, samples, and measures once; `estimate_observables(...)` and +`estimate_energy(...)` remain compatibility wrappers. Return complex means when appropriate but report Hermitian observables using their real part only after checking the imaginary residual. Keep energy, diagonal observables, hopping, correlations, and arbitrary supported Fermion observables on the diff --git a/docs/api/vmc.md b/docs/api/vmc.md index 8fe5e46..b2fe985 100644 --- a/docs/api/vmc.md +++ b/docs/api/vmc.md @@ -638,10 +638,10 @@ The estimate result retains the legacy `energy_mean`, `energy_variance`, and `energy_stderr` field names; they contain statistics for the configured observable. `estimate_energy(...)` remains as a compatibility alias. -For a coordinate-labelled PEPS, use `TorchFermionVMC` to derive the lattice, -physical charge ordering, initial sector, and sampler rule in one place. Pass -`fermion` to generate the default Hamiltonian, or omit it when supplying -explicit `terms`: +For compatibility with an existing lower-level sweep loop, coordinate-labelled +PEPS can still initialize `TorchFermionVMC` in the constructor. New code +should prefer the first-run recipe below instead. Pass `fermion` to generate +the default Hamiltonian, or omit it when supplying explicit `terms`: ```python from pepsy import Fermion @@ -663,6 +663,109 @@ result = vmc.estimate_observable( ) ``` +For a native fermionic PEPS measurement, construct `TorchFermionVMC` from the +state and native Fermion terms only. The first `sample` (or `warmup`) owns both +the chain recipe and PEPS contraction recipe: + +1. `SamplingConfig` owns the number of chains, retained samples, burn-in, + thinning, and RNG seeds. In the native Torch sampler, `burn_in` counts + discarded thinning intervals, so the `Metropolis` total is + `(burn_in + n_samples_per_chain) * thin` batched sweeps; every batched + sweep advances all chains once. +2. `contraction_opts` is a single mapping with `method`, `chi`, `cutoff`, and + any backend options such as `mode`. It is consumed when the first operation builds + the amplitude model, then remains fixed with the Markov state. +3. `observables` is a name-to-term mapping. `measure(samples, observables=...)` + uses one retained batch for every entry. Include `"energy": terms` + explicitly when the Hamiltonian should be visible in the measurement recipe. + +```python +sampling = pvmc.SamplingConfig( + n_samples_per_chain=256, + n_chains=32, + burn_in=64, + thin=2, + seed=7, +) +contraction_opts = { + "method": "boundary", + "chi": 32, + "cutoff": 1e-10, + "mode": "mps", +} + +vmc = pvmc.TorchFermionVMC( + peps, + fermion=fermion, + terms=terms, # native Fermi-Hubbard terms; no JW conversion +) + +# Optional, non-MCMC warm-up: inspect one valid PEPS amplitude. +warmup = vmc.warmup( + sampling=sampling, + contraction_opts=contraction_opts, +) + +# Phase 1: exactly one Metropolis pass. `samples` retains psi(x) for every x. +samples = vmc.sample(sampling=sampling, progress=True) + +# Phase 2: no new Metropolis work. Reuse this batch for energy, eta, density, ... +estimates = vmc.measure( + samples, + observables={"energy": terms, "eta": eta_terms}, + progress=True, +) +print(warmup.amplitude) +print(estimates["energy"].energy_mean, estimates["eta"].energy_stderr) +``` + +`estimate_observables({...}, sampling=..., contraction_opts=...)` provides the +same one-batch behavior without retaining a separately named sample batch. +`run(...)` is the one-command convenience form: warm up, sample once, then +measure. `measure_samples(...)` remains the lower-level spelling of +`measure(samples, ...)`. Native sample batches carry the PEPS parameter +versions and contraction signature used to draw them, so `measure` rejects a +batch after either changes; draw fresh samples after an optimization update. +`progress=True` reports optional burn-in sweeps, MCMC +sampling, then the shared connection-building, amplitude-contraction, and +statistics phases. The `Metropolis` bar reports walkers (chains), retained +samples per walker, burn-in/thinning, proposal, contraction method/`chi`, +acceptance, and live boundary-environment cache reuse/build activity. Its +`phase` is `equilibrate` while discarded intervals run, then `retain i/n` as +each retained configuration per walker is recorded. The +`Evaluation` bar reports the shared sample shape, observables, whether parent +amplitudes were stored, connection count, and the diagonal/environment/direct +target-amplitude split. The warm-up amplitude is a representative PEPS +amplitude, not an energy estimate. The legacy constructor-level `n_walkers`, +`contraction`, `chi`, and `cutoff` options remain supported for existing +scripts, but new measurement code should keep them in `SamplingConfig` and +`contraction_opts` as above. + +External MPS/BP/tree proposal sampling uses the same explicit two-stage +shape, but has no Metropolis burn-in or thinning. Pass its independent count +as `n_samples`, rather than a `SamplingConfig`: + +```python +importance_samples = vmc.sample( + proposal=mps_sampler, + n_samples=512, + fermion=proposal_fermion, + one_d_to_two_d=mps_site_to_peps_coordinate, +) +importance_estimates = vmc.measure( + importance_samples, + observables={"energy": terms, "eta": eta_terms}, +) +``` + +`importance_samples` stores PEPS-code configurations, `log q(x)`, and target +parent amplitudes. The later `measure` call automatically forms the +self-normalized weights `|psi(x)|**2 / q(x)` once and shares the resulting +target-amplitude work across every observable. Unlike a target-Metropolis +batch, an external-proposal batch remains valid after a PEPS update: `measure` +refreshes its target amplitudes while retaining the fixed proposal density. +`measure_from_proposal(...)` remains the one-call compatibility shortcut. + By default, `pbc=None` reads the PEPS cyclic axes through Quimb's `is_cyclic_x()` and `is_cyclic_y()` metadata; pass `pbc=` or `edges=` to override that inference. When explicit two-site `terms` are supplied, their diff --git a/src/pepsy/vmc/__init__.py b/src/pepsy/vmc/__init__.py index d63ed8f..babbef6 100644 --- a/src/pepsy/vmc/__init__.py +++ b/src/pepsy/vmc/__init__.py @@ -26,8 +26,11 @@ "SymmetryFallbackWarning": ".api", "NetKetLocalConfigMap": ".netket", "NetKetChunkSettings": ".netket", + "NetKetBuildTiming": ".netket", + "NetKetAmplitudeTiming": ".netket", "NetKetPEPSVMC": ".netket", "NetKetVMCSetup": ".netket", + "NetKetVMCConfig": ".netket", "NetKetFermiHubbardVMC": ".netket", "NetKetSparseFermiHubbardVMC": ".netket", "NetKetVMCSettings": ".netket", @@ -38,7 +41,9 @@ "TorchFermionVMC": ".torch", "TorchFermionVMCMetadata": ".torch", "TorchChainDiagnostics": ".torch", + "TorchImportanceSamples": ".torch", "TorchMCMCSamples": ".torch", + "TorchSampleProvenance": ".torch", "TorchMetropolisResult": ".torch", "TorchMetropolisSampler": ".torch", "TorchBPMetropolisSampler": ".torch", @@ -48,7 +53,9 @@ "TorchVMCSetup": ".torch", "TorchVMCEnergyEstimate": ".torch", "TorchVMCImportanceEstimate": ".torch", + "TorchVMCMeasurementRun": ".torch", "TorchVMCStepResult": ".torch", + "TorchVMCWarmupResult": ".torch", "TorchSRResult": ".torch", "TorchSquareLattice": ".torch", "VMCMeasurement": ".api", diff --git a/src/pepsy/vmc/api.py b/src/pepsy/vmc/api.py index 18755d3..e1f5938 100644 --- a/src/pepsy/vmc/api.py +++ b/src/pepsy/vmc/api.py @@ -149,7 +149,15 @@ def _resolve_contraction_config(contraction, chi=None, cutoff=None, options=None @dataclass(frozen=True) class SamplingConfig: - """Shared chain-preserving sampling settings.""" + """Shared chain-preserving sampling settings. + + ``burn_in`` is the number of discarded *thinning intervals* per chain. + Thus the native Torch sampler advances each chain + ``(burn_in + n_samples_per_chain) * thin`` Metropolis sweeps: it discards + ``burn_in * thin`` sweeps, then retains one configuration after every + ``thin`` further sweeps. The returned batch has shape + ``(n_samples_per_chain, n_chains, n_sites)``. + """ n_samples_per_chain: int = 128 n_chains: int = 16 diff --git a/src/pepsy/vmc/torch/__init__.py b/src/pepsy/vmc/torch/__init__.py index c9470f8..5ddda1a 100644 --- a/src/pepsy/vmc/torch/__init__.py +++ b/src/pepsy/vmc/torch/__init__.py @@ -42,11 +42,15 @@ ) from .results import ( TorchMetropolisResult, + TorchImportanceSamples, TorchMCMCSamples, + TorchSampleProvenance, TorchChainDiagnostics, TorchVMCEnergyEstimate, TorchVMCImportanceEstimate, + TorchVMCMeasurementRun, TorchVMCStepResult, + TorchVMCWarmupResult, ) from .sampler import TorchBPMetropolisSampler, TorchMetropolisSampler, metropolis_local_sampler from .sr import TorchSRResult, apply_torch_sr_update, solve_torch_sr, torch_log_derivative_matrix @@ -59,7 +63,9 @@ "TorchPEPSBoundaryAmplitude", "TorchConnections", "TorchMetropolisResult", + "TorchImportanceSamples", "TorchMCMCSamples", + "TorchSampleProvenance", "TorchChainDiagnostics", "TorchMetropolisSampler", "TorchBPMetropolisSampler", @@ -68,7 +74,9 @@ "TorchVMCSetup", "TorchVMCEnergyEstimate", "TorchVMCImportanceEstimate", + "TorchVMCMeasurementRun", "TorchVMCStepResult", + "TorchVMCWarmupResult", "TorchSRResult", "TorchSquareLattice", "apply_torch_sr_update", diff --git a/src/pepsy/vmc/torch/_core.py b/src/pepsy/vmc/torch/_core.py index 4da3c71..74f9049 100644 --- a/src/pepsy/vmc/torch/_core.py +++ b/src/pepsy/vmc/torch/_core.py @@ -30,11 +30,15 @@ ) from .results import ( TorchChainDiagnostics, + TorchImportanceSamples, TorchMCMCSamples, TorchMetropolisResult, + TorchSampleProvenance, TorchVMCImportanceEstimate, TorchVMCEnergyEstimate, + TorchVMCMeasurementRun, TorchVMCStepResult, + TorchVMCWarmupResult, ) from .metadata import ( TorchFermionVMCMetadata, @@ -116,7 +120,9 @@ "TorchPEPSBoundaryAmplitude", "TorchConnections", "TorchMetropolisResult", + "TorchImportanceSamples", "TorchMCMCSamples", + "TorchSampleProvenance", "TorchChainDiagnostics", "TorchMetropolisSampler", "TorchBPMetropolisSampler", @@ -125,7 +131,9 @@ "TorchVMCSetup", "TorchVMCEnergyEstimate", "TorchVMCImportanceEstimate", + "TorchVMCMeasurementRun", "TorchVMCStepResult", + "TorchVMCWarmupResult", "TorchSRResult", "TorchSquareLattice", "apply_torch_sr_update", diff --git a/src/pepsy/vmc/torch/driver.py b/src/pepsy/vmc/torch/driver.py index 628ff3a..ce3d542 100644 --- a/src/pepsy/vmc/torch/driver.py +++ b/src/pepsy/vmc/torch/driver.py @@ -42,7 +42,9 @@ _accumulate_cache_profile, _cache_profile_snapshot, _make_progress, + _set_evaluation_progress_postfix, _set_vmc_progress_postfix, + _torch_sample_provenance, ) from .sampler import TorchBPMetropolisSampler, TorchMetropolisSampler @@ -600,9 +602,17 @@ def burn_in( bar = _make_progress( True, total=n_sweeps, - desc="Torch VMC burn-in", + desc="Metropolis warm-up", unit="sweep", ) + _set_vmc_progress_postfix( + bar, + n_sites=self.n_sites, + include_energy=False, + n_chains=self.n_walkers, + model=self.model, + proposal=self.proposal, + ) result = None n_proposed = 0 n_accepted = 0 @@ -623,6 +633,9 @@ def burn_in( result, n_sites=self.n_sites, include_energy=False, + n_chains=self.n_walkers, + model=self.model, + proposal=self.proposal, ) finally: bar.close() @@ -734,6 +747,7 @@ def measure_samples( proposal_log_probs=None, profile=False, deduplicate=True, + progress=False, ): """Measure saved chain samples without running another sampler. @@ -741,8 +755,14 @@ def measure_samples( tensor with shape ``(n_samples_per_chain, n_chains, n_sites)``. A two-dimensional tensor is interpreted as one retained sample per chain. Stored amplitudes from ``TorchMCMCSamples`` are reused unless - ``amplitudes=`` is supplied explicitly; pass an explicit amplitude - batch when the PEPS parameters have changed since sampling. + ``amplitudes=`` is supplied explicitly. Native Markov samples carry a + PEPS/contraction provenance record and are rejected after that state + has changed: draw a fresh Markov batch after an update rather than + mixing configurations from the old Born distribution with the new + local estimator. In contrast, :class:`TorchImportanceSamples` came + from a fixed external proposal ``q``; it remains valid after a PEPS + update and refreshes its target parent amplitudes before forming + ``|psi(x)|**2 / q(x)``. With ``observables=None`` the driver's configured connection function is measured and one :class:`TorchVMCEnergyEstimate` is returned. A @@ -763,13 +783,26 @@ def measure_samples( By default, repeated parent configurations and repeated connected targets are contracted once and scattered back to their original chain positions. Set ``deduplicate=False`` for compatibility - diagnostics or timing comparisons. + diagnostics or timing comparisons. Set ``progress=True`` to report + connection construction, amplitude contraction, and statistics. """ torch = _require_torch() start = time.perf_counter() model_device = _model_device(self.model) sample_object = samples if hasattr(samples, "configs") else None + provenance = getattr(sample_object, "provenance", None) + if provenance is not None and provenance != _torch_sample_provenance(self.model): + raise RuntimeError( + "Samples belong to a different PEPS/model state. Call " + "sample(...) again after modifying the model or its " + "contraction settings." + ) + target_provenance = getattr(sample_object, "target_provenance", None) + refresh_proposal_amplitudes = ( + target_provenance is not None + and target_provenance != _torch_sample_provenance(self.model) + ) raw_configs = ( getattr(sample_object, "configs", None) if sample_object is not None @@ -803,6 +836,12 @@ def measure_samples( if amplitudes is None and sample_object is not None: amplitudes = getattr(sample_object, "amplitudes", None) + if refresh_proposal_amplitudes: + amplitudes = None + parent_amplitude_source = ( + "stored" if amplitudes is not None else "refreshed" + if refresh_proposal_amplitudes else "contracted" + ) if amplitudes is None: with torch.no_grad(): if deduplicate and unique_parent_count < flat_configs.shape[0]: @@ -885,6 +924,26 @@ def measure_samples( raise ValueError("observables must contain at least one entry.") return_mapping = True + phase_bar = _make_progress( + progress, + total=3, + desc="Evaluation", + unit="stage", + ) + + def set_phase(stage, *, n_connections=None): + _set_evaluation_progress_postfix( + phase_bar, + model=self.model, + n_steps=n_steps, + n_chains=n_chains, + observables=(name for name, _ in observable_items), + parent_amplitudes=parent_amplitude_source, + stage=stage, + n_connections=n_connections, + ) + + set_phase("connections") connection_start = time.perf_counter() connection_map = { name: ( @@ -894,8 +953,15 @@ def measure_samples( ) for name, terms in observable_items } + n_connections = sum( + int(connections.configs.shape[0]) + for connections in connection_map.values() + ) connection_elapsed = time.perf_counter() - connection_start + if phase_bar is not None: + phase_bar.update(1) + set_phase("target amplitudes", n_connections=n_connections) local_start = time.perf_counter() with torch.no_grad(): flat_values = _local_energies_from_connection_map( @@ -909,6 +975,8 @@ def measure_samples( compile_kernels=self.compile_kernels, ) local_elapsed = time.perf_counter() - local_start + if phase_bar is not None: + phase_bar.update(1) elapsed = time.perf_counter() - start acceptance_rate = float( @@ -935,6 +1003,7 @@ def measure_samples( "weighted": importance_weights is not None, } + set_phase("statistics", n_connections=n_connections) results = {} for name, _ in observable_items: local_values = flat_values[name].reshape(n_steps, n_chains) @@ -994,6 +1063,9 @@ def measure_samples( ), ) + if phase_bar is not None: + phase_bar.update(1) + phase_bar.close() return results if return_mapping else results["observable"] def energy_estimate(self): @@ -1005,6 +1077,7 @@ def energy_estimate(self): def estimate_observable( self, *, + sampling=None, burn_in=0, n_measurements=1, sweeps_between=1, @@ -1050,8 +1123,11 @@ def estimate_observable( opt-in so normal short VMC loops keep their existing overhead. """ profile = bool(profile) + if sampling is not None and sampler is not None: + raise ValueError("Pass either sampling=... or sampler=..., not both.") modern_sampling = ( - sampler is not None + sampling is not None + or sampler is not None or n_samples is not None or n_chains is not None or n_discard_per_chain is not None @@ -1070,6 +1146,7 @@ def estimate_observable( ) if sampler is None: samples = self.sample( + sampling=sampling, n_samples=1024 if n_samples is None else n_samples, n_chains=n_chains, n_discard_per_chain=n_discard_per_chain, @@ -1344,6 +1421,7 @@ def estimate_observables( self, observables, *, + sampling=None, burn_in=0, n_measurements=1, sweeps_between=1, @@ -1383,6 +1461,8 @@ def estimate_observables( if not observable_items: raise ValueError("observables must contain at least one entry.") profile = bool(profile) + if sampling is not None and sampler is not None: + raise ValueError("Pass either sampling=... or sampler=..., not both.") def make_connection_map(configs): return { @@ -1444,7 +1524,8 @@ def make_results( return results modern_sampling = ( - sampler is not None + sampling is not None + or sampler is not None or n_samples is not None or n_chains is not None or n_discard_per_chain is not None @@ -1463,6 +1544,7 @@ def make_results( ) if sampler is None: samples = self.sample( + sampling=sampling, n_samples=1024 if n_samples is None else n_samples, n_chains=n_chains, n_discard_per_chain=n_discard_per_chain, @@ -1494,14 +1576,21 @@ def make_results( phase_bar = _make_progress( progress, total=3, - desc="Torch VMC evaluation", - unit="phase", + desc="Evaluation", + unit="stage", ) - observable_names = ", ".join(name for name, _ in observable_items) - def set_phase(stage): - if phase_bar is not None: - phase_bar.set_postfix({"stage": stage}) + def set_phase(stage, *, n_connections=None): + _set_evaluation_progress_postfix( + phase_bar, + model=self.model, + n_steps=samples.n_samples_per_chain, + n_chains=samples.n_chains, + observables=(name for name, _ in observable_items), + parent_amplitudes="stored", + stage=stage, + n_connections=n_connections, + ) try: estimator_start = time.perf_counter() @@ -1510,14 +1599,18 @@ def set_phase(stage): flat_configs = sample_configs.reshape(-1, self.n_sites) flat_amplitudes = sample_amplitudes.reshape(-1) - set_phase("building shared connections") + set_phase("connections") connection_start = time.perf_counter() connection_map = make_connection_map(flat_configs) + n_connections = sum( + int(connections.configs.shape[0]) + for connections in connection_map.values() + ) connection_elapsed = time.perf_counter() - connection_start if phase_bar is not None: phase_bar.update(1) - set_phase(f"contracting {observable_names}") + set_phase("target amplitudes", n_connections=n_connections) local_start = time.perf_counter() with _require_torch().no_grad(): flat_values = _local_energies_from_connection_map( @@ -1534,7 +1627,7 @@ def set_phase(stage): if phase_bar is not None: phase_bar.update(1) - set_phase("computing statistics") + set_phase("statistics", n_connections=n_connections) local_values = { name: values.reshape(sample_configs.shape[:-1]) for name, values in flat_values.items() @@ -1774,6 +1867,42 @@ def estimate_energy( auto_thin=auto_thin, ) + def sample_from_proposal( + self, + proposal, + *, + n_samples=128, + seed=None, + fermion=None, + one_d_to_two_d=None, + site_order=None, + occupation_map=None, + sample_kwargs=None, + progress=False, + amplitude_floor=0.0, + ): + """Draw reusable PEPS-code samples from an external proposal. + + The result retains ``log q(x)`` and target amplitudes. Pass it to + :meth:`measure_samples` to form the self-normalized importance + estimator for any observable map without drawing the proposal again. + """ + from .importance import sample_from_proposal + + return sample_from_proposal( + self, + proposal, + n_samples=n_samples, + seed=seed, + fermion=fermion, + one_d_to_two_d=one_d_to_two_d, + site_order=site_order, + occupation_map=occupation_map, + sample_kwargs=sample_kwargs, + progress=progress, + amplitude_floor=amplitude_floor, + ) + def measure_from_proposal( self, proposal, @@ -1793,16 +1922,13 @@ def measure_from_proposal( ): """Measure from an external MPS, BP, tree, or proposal batch. - The proposal is normalized at this boundary and the resulting batch - is delegated to :meth:`measure_samples`. ``one_d_to_two_d`` and - ``fermion`` are required only when a bare MPS must be wrapped in a - :class:`pepsy.sampling.MpsSampler`; sampled MPS batches carry their - own coordinate map and occupation decoder. + This is the compatibility one-shot form of + :meth:`sample_from_proposal` followed by :meth:`measure_samples`. + ``one_d_to_two_d`` and ``fermion`` are required only when a bare MPS + must be wrapped in a :class:`pepsy.sampling.MpsSampler`; sampled MPS + batches carry their own coordinate map and occupation decoder. """ - from .importance import measure_from_proposal - - return measure_from_proposal( - self, + samples = self.sample_from_proposal( proposal, n_samples=n_samples, seed=seed, @@ -1811,11 +1937,15 @@ def measure_from_proposal( site_order=site_order, occupation_map=occupation_map, sample_kwargs=sample_kwargs, - observables=observables, progress=progress, amplitude_floor=amplitude_floor, + ) + return self.measure_samples( + samples, + observables=observables, profile=profile, deduplicate=deduplicate, + progress=progress, ) def importance_energy_estimate( diff --git a/src/pepsy/vmc/torch/fermion.py b/src/pepsy/vmc/torch/fermion.py index 9a52578..5f03fc2 100644 --- a/src/pepsy/vmc/torch/fermion.py +++ b/src/pepsy/vmc/torch/fermion.py @@ -5,6 +5,7 @@ from dataclasses import dataclass, replace from itertools import product import numpy as np +import time from typing import Any from ..torch_types import FermionSiteEncoding, _check_positive_int, _require_torch @@ -15,12 +16,12 @@ ) from .amplitude import ( _call_amplitude_fn, - _validate_contraction, make_torch_peps_amplitude_model, ) from .connections import compile_operator_sum_torch, _normalize_terms_site_labels from .driver import TorchVMCDriver from .metadata import _infer_torch_fermion_metadata +from .results import TorchVMCMeasurementRun, TorchVMCWarmupResult __all__ = [ "TorchFermionVMC", @@ -236,6 +237,71 @@ def keep(candidate): return configs[choice], amplitudes[choice] +def _contraction_config( + contraction=None, + *, + chi=None, + cutoff=None, + contraction_opts=None, +): + """Normalize legacy or lazy-run contraction settings to one config.""" + from ..api import ContractionConfig + + if contraction is None: + if contraction_opts is None: + return None + try: + raw = dict(contraction_opts) + except (TypeError, ValueError) as exc: + raise TypeError("contraction_opts must be a mapping or None.") from exc + method = raw.pop("method", raw.pop("contraction", None)) + if method is None: + raise ValueError( + "contraction_opts must define 'method' (or 'contraction') when " + "passed without contraction=...." + ) + option_chi = raw.pop("chi", None) + option_cutoff = raw.pop("cutoff", 0.0) + options = raw.pop("options", raw.pop("backend_options", {})) + if raw: + options = {**dict(options), **raw} + if chi is not None and option_chi is not None and chi != option_chi: + raise ValueError("chi conflicts with contraction_opts.chi.") + if cutoff is not None and float(cutoff) != float(option_cutoff): + raise ValueError("cutoff conflicts with contraction_opts.cutoff.") + return ContractionConfig( + method=method, + chi=option_chi if chi is None else chi, + cutoff=option_cutoff if cutoff is None else cutoff, + options=options, + ) + + if isinstance(contraction, ContractionConfig): + if chi is not None and contraction.chi is not None and chi != contraction.chi: + raise ValueError(f"chi={chi} conflicts with contraction.chi={contraction.chi}.") + if cutoff is not None and float(cutoff) != contraction.cutoff: + raise ValueError( + f"cutoff={cutoff} conflicts with contraction.cutoff={contraction.cutoff}." + ) + if contraction_opts is not None and dict(contraction_opts) != dict(contraction.options): + raise ValueError("contraction_opts conflicts with contraction.options.") + return contraction + + return ContractionConfig( + method=contraction, + chi=chi, + cutoff=0.0 if cutoff is None else cutoff, + options={} if contraction_opts is None else contraction_opts, + ) + + +def _default_fermion_contraction(): + """Return the historical native-Torch default for legacy entry points.""" + from ..api import ContractionConfig + + return ContractionConfig(method="boundary", chi=4, cutoff=1.0e-10) + + class TorchFermionVMC(TorchVMCDriver): """Automatic native spinful Fermion VMC around a Quimb PEPS. @@ -247,6 +313,12 @@ class TorchFermionVMC(TorchVMCDriver): metadata. The lower-level :class:`TorchVMCDriver` remains available when callers need full manual control over configurations or connection functions. + + With the concise measurement API, omit the constructor-era chain and + contraction controls. The first :meth:`run` receives ``sampling=`` and + ``contraction_opts=`` and creates the matching sampler and PEPS amplitude + model. Constructor-level ``n_walkers`` and contraction keywords remain + supported for compatibility and initialize the driver immediately. """ def __init__( @@ -262,9 +334,9 @@ def __init__( site_order=None, sector=None, configs=None, - n_walkers=128, - contraction="boundary", - chi=4, + n_walkers=None, + contraction=None, + chi=None, cutoff=None, contraction_opts=None, dtype=None, @@ -287,8 +359,6 @@ def __init__( init_max_attempts=32, init_max_states=100_000, ): - torch = _require_torch() - from ..api import ContractionConfig if hamiltonian is not None and terms is not None: raise ValueError( "Pass either hamiltonian=... or terms=..., not both; " @@ -296,14 +366,12 @@ def __init__( ) if hamiltonian is not None: terms = hamiltonian - if isinstance(contraction, ContractionConfig): - if contraction.chi is not None: - chi = contraction.chi - if cutoff is None: - cutoff = contraction.cutoff - if contraction_opts is None: - contraction_opts = dict(contraction.options) - contraction = contraction.method + legacy_contraction = _contraction_config( + contraction, + chi=chi, + cutoff=cutoff, + contraction_opts=contraction_opts, + ) metadata = _infer_torch_fermion_metadata( peps, fermion, @@ -319,32 +387,8 @@ def __init__( "basis. Omit encoding=... to infer it safely." ) - model_kwargs = { - "contraction": contraction, - "chi": chi, - "cutoff": cutoff, - "contraction_opts": contraction_opts, - "dtype": dtype, - "device": device, - "site_order": metadata.site_order, - "graded_torch": graded_torch, - "amplitude_batching": amplitude_batching, - } - if _validate_contraction(contraction, chi) == "boundary": - model_kwargs.update( - proposal_batching=proposal_batching, - proposal_vmap_min_batch=proposal_vmap_min_batch, - ) - model = make_torch_peps_amplitude_model(peps, **model_kwargs) - model_device = _model_device(model, device=device) if generator is not None and seed is not None: raise ValueError("Pass either generator=... or seed=..., not both.") - if seed is not None: - try: - generator = torch.Generator(device=model_device) - except (RuntimeError, TypeError, ValueError): - generator = torch.Generator() - generator.manual_seed(int(seed)) from ..api import OperatorSum if terms is None: @@ -366,6 +410,159 @@ def __init__( hamiltonian = terms terms = _normalize_terms_site_labels(terms, metadata.site_order) + self.peps = peps + self.fermion = fermion + self.metadata = metadata + self.hamiltonian = hamiltonian + self.observables = self._compile_observables(observables) + self.physical_charges = metadata.physical_charges + if proposal is None: + if metadata.spinful: + proposal = { + "U1": "spinful_u1", + "U1U1": "spinful", + "Z2": "spinful_z2", + "Z2Z2": "spinful_z2z2", + }[metadata.symmetry] + else: + proposal = "spin" + self._driver_initialized = False + self._contraction_config = None + self._legacy_contraction_config = legacy_contraction + self._initial_configs = configs + self._initial_n_walkers = n_walkers + self._initial_generator = generator + self._initial_seed = seed + self._initial_amplitude_floor = amplitude_floor + self._initial_max_attempts = init_max_attempts + self._initial_max_states = init_max_states + self._hamiltonian_terms = terms + self._model_options = { + "dtype": dtype, + "device": device, + "graded_torch": graded_torch, + "amplitude_batching": amplitude_batching, + "proposal_batching": proposal_batching, + "proposal_vmap_min_batch": proposal_vmap_min_batch, + } + self._driver_options = { + "proposal": proposal, + "hopping_rate": hopping_rate, + "spin_flip_rate": spin_flip_rate, + "pair_toggle_rate": pair_toggle_rate, + "chunk_size": chunk_size, + "compile_kernels": compile_kernels, + "log_amplitude_fn": log_amplitude_fn, + } + if configs is not None or n_walkers is not None or legacy_contraction is not None: + self._ensure_initialized( + contraction=legacy_contraction, + n_walkers=n_walkers, + ) + + def _ensure_initialized( + self, + *, + sampling=None, + contraction=None, + contraction_opts=None, + n_walkers=None, + ): + """Initialize the native driver once, from the measurement recipe. + + The concise API deliberately leaves chain count and amplitude + contraction unset until a first measurement. This lets one + ``SamplingConfig`` own every sampling choice and one + ``contraction_opts`` mapping own every contraction choice. Once a + Markov state exists, changing either would silently mix incompatible + chains or amplitudes, so it is rejected explicitly. + """ + from ..api import SamplingConfig + + if sampling is not None and not isinstance(sampling, SamplingConfig): + raise TypeError("sampling must be a SamplingConfig or None.") + + requested_contraction = _contraction_config( + contraction, + contraction_opts=contraction_opts, + ) + if requested_contraction is None: + if self._driver_initialized: + requested_contraction = self._contraction_config + else: + requested_contraction = ( + self._legacy_contraction_config + or _default_fermion_contraction() + ) + + if sampling is not None: + requested_n_walkers = sampling.n_chains + elif n_walkers is not None: + requested_n_walkers = n_walkers + elif self._driver_initialized: + requested_n_walkers = self.n_walkers + elif self._initial_configs is not None: + requested_n_walkers = int(_as_long_matrix(self._initial_configs).shape[0]) + elif self._initial_n_walkers is not None: + requested_n_walkers = self._initial_n_walkers + else: + requested_n_walkers = 128 + + if self._driver_initialized: + if requested_contraction != self._contraction_config: + raise ValueError( + "contraction settings are fixed after the first native VMC " + "run; create a new TorchFermionVMC for a different " + "contraction." + ) + if requested_n_walkers != self.n_walkers: + raise ValueError( + "SamplingConfig.n_chains must match the existing native " + f"VMC chain count ({self.n_walkers}), got " + f"{requested_n_walkers}. Create a new TorchFermionVMC " + "for a different chain count." + ) + return + + self._initialize_driver( + requested_contraction, + n_walkers=requested_n_walkers, + ) + + def _initialize_driver(self, contraction, *, n_walkers): + """Build the amplitude model and initial walkers for a first run.""" + torch = _require_torch() + model_kwargs = { + "contraction": contraction, + "dtype": self._model_options["dtype"], + "device": self._model_options["device"], + "site_order": self.metadata.site_order, + "graded_torch": self._model_options["graded_torch"], + "amplitude_batching": self._model_options["amplitude_batching"], + } + if contraction.method == "boundary": + model_kwargs.update( + proposal_batching=self._model_options["proposal_batching"], + proposal_vmap_min_batch=self._model_options[ + "proposal_vmap_min_batch" + ], + ) + model = make_torch_peps_amplitude_model(self.peps, **model_kwargs) + model_device = _model_device( + model, + device=self._model_options["device"], + ) + + generator = self._initial_generator + if self._initial_seed is not None: + try: + generator = torch.Generator(device=model_device) + except (RuntimeError, TypeError, ValueError): + generator = torch.Generator() + generator.manual_seed(int(self._initial_seed)) + + metadata = self.metadata + configs = self._initial_configs if configs is None: if metadata.sector is None: raise ValueError( @@ -378,21 +575,23 @@ def __init__( n_walkers, device=model_device, generator=generator, - amplitude_floor=amplitude_floor, - max_attempts=init_max_attempts, - max_states=init_max_states, + amplitude_floor=self._initial_amplitude_floor, + max_attempts=self._initial_max_attempts, + max_states=self._initial_max_states, ) else: configs = _as_long_matrix(configs).to(device=model_device) if configs.shape[1] != metadata.n_sites: raise ValueError( - f"configs must have {metadata.n_sites} sites, got {configs.shape[1]}." + f"configs must have {metadata.n_sites} sites, got " + f"{configs.shape[1]}." ) metadata.encoding.validate(configs) actual_sector = _fermion_sector_from_configs(configs, metadata) if metadata.sector is not None and actual_sector != metadata.sector: raise ValueError( - f"configs are in sector {actual_sector}, expected {metadata.sector}." + f"configs are in sector {actual_sector}, expected " + f"{metadata.sector}." ) if metadata.sector is None: metadata = replace(metadata, sector=actual_sector) @@ -400,47 +599,35 @@ def __init__( amplitudes = _call_amplitude_fn(model, configs) valid = ( torch.isfinite(amplitudes.abs()) - & (amplitudes.abs() > float(amplitude_floor)) + & (amplitudes.abs() > float(self._initial_amplitude_floor)) ) if not bool(torch.all(valid)): raise ValueError( - "configs contain zero, non-finite, or below-floor PEPS amplitudes." + "configs contain zero, non-finite, or below-floor PEPS " + "amplitudes." ) - self.peps = peps - self.fermion = fermion self.metadata = metadata - self.hamiltonian = hamiltonian - self.observables = self._compile_observables(observables) self.physical_charges = metadata.physical_charges - if proposal is None: - if metadata.spinful: - proposal = { - "U1": "spinful_u1", - "U1U1": "spinful", - "Z2": "spinful_z2", - "Z2Z2": "spinful_z2z2", - }[metadata.symmetry] - else: - proposal = "spin" - super().__init__( model, metadata.graph, configs, - terms=terms, + terms=self._hamiltonian_terms, site_order=metadata.site_order, amplitudes=amplitudes, - proposal=proposal, - hopping_rate=hopping_rate, - spin_flip_rate=spin_flip_rate, - pair_toggle_rate=pair_toggle_rate, + proposal=self._driver_options["proposal"], + hopping_rate=self._driver_options["hopping_rate"], + spin_flip_rate=self._driver_options["spin_flip_rate"], + pair_toggle_rate=self._driver_options["pair_toggle_rate"], encoding=metadata.encoding, - chunk_size=chunk_size, - compile_kernels=compile_kernels, - log_amplitude_fn=log_amplitude_fn, + chunk_size=self._driver_options["chunk_size"], + compile_kernels=self._driver_options["compile_kernels"], + log_amplitude_fn=self._driver_options["log_amplitude_fn"], generator=generator, ) + self._contraction_config = contraction + self._driver_initialized = True @property def Lx(self): @@ -471,6 +658,7 @@ def measure_from_mps( encoding. A bare MPS additionally needs ``one_d_to_two_d`` and the constructor's native ``fermion`` object so its sampler can be built. """ + self._ensure_initialized() return self.measure_from_proposal( proposal, n_samples=n_samples, @@ -487,6 +675,346 @@ def measure_from_mps( deduplicate=deduplicate, ) + def _measurement_observables(self, observables, *, include_energy=False): + """Compile a user observable mapping for the native estimator.""" + if observables is None: + compiled = dict(self.observables) + else: + try: + entries = tuple(observables.items()) + except AttributeError as exc: + raise TypeError( + "observables must be a mapping of names to operators." + ) from exc + compiled = {} + for name, value in entries: + if value is None: + compiled[name] = None + else: + compiled[name] = self._compile_observables({name: value})[name] + if include_energy and "energy" not in compiled: + compiled = {"energy": None, **compiled} + if not compiled: + raise ValueError( + "No observables are configured. Pass observables=... or provide " + "observables=... when constructing TorchFermionVMC." + ) + return compiled + + @staticmethod + def _sampling_estimator_kwargs(sampling, kwargs): + """Lower a shared sampling config without silently overriding options.""" + kwargs = dict(kwargs) + if sampling is None: + return kwargs + from ..api import SamplingConfig + + if not isinstance(sampling, SamplingConfig): + raise TypeError("sampling must be a SamplingConfig or None.") + if kwargs.get("sampler") is not None: + raise ValueError("Pass either sampling=... or sampler=..., not both.") + configured = sampling.torch_kwargs() + expected = { + "n_samples": configured["n_samples"], + "n_chains": configured["n_chains"], + "n_discard_per_chain": configured["n_discard_per_chain"], + "n_discard": configured["n_discard_per_chain"], + "sweep_size": configured["n_thin"], + "n_thin": configured["n_thin"], + "seed": configured["seed"], + "sampler_seed": configured["sampler_seed"], + } + for name, value in expected.items(): + supplied = kwargs.get(name) + if supplied is not None and supplied != value: + raise ValueError(f"{name} conflicts with sampling.") + kwargs["sampling"] = sampling + return kwargs + + def estimate_observables( + self, + observables=None, + *, + sampling=None, + contraction=None, + contraction_opts=None, + **kwargs, + ): + """Estimate native PEPS observables from one shared Markov sample set. + + Values in ``observables`` may be native Fermion terms, + :class:`~pepsy.vmc.OperatorSum` objects, or ``None`` to reuse this + driver's Hamiltonian. Omit ``observables`` to measure the Hamiltonian + together with the supplemental observables supplied at construction. + ``sampling`` centralizes chains, burn-in, thinning, and seeds through + :class:`~pepsy.vmc.SamplingConfig`. + """ + self._ensure_initialized( + sampling=sampling, + contraction=contraction, + contraction_opts=contraction_opts, + ) + compiled = self._measurement_observables( + observables, + include_energy=observables is None, + ) + kwargs = self._sampling_estimator_kwargs(sampling, kwargs) + return super().estimate_observables(compiled, **kwargs) + + def sample( + self, + *, + sampling=None, + contraction=None, + contraction_opts=None, + proposal=None, + **kwargs, + ): + """Collect reusable Markov or external-proposal samples. + + On the first call, pass both ``sampling`` and ``contraction_opts``. + With no ``proposal``, the returned :class:`TorchMCMCSamples` retains + chain configurations and parent PEPS amplitudes. Pass an MPS/BP/tree + sampler or a sampled proposal batch as ``proposal=...`` to obtain + :class:`TorchImportanceSamples` instead. That path draws from ``q`` + once, stores ``log q(x)``, and lets :meth:`measure` form importance + estimates for any number of observables without another proposal draw. + + ``SamplingConfig`` describes target-Metropolis burn-in and thinning, + so it does not apply to independently drawn proposal samples; use + ``n_samples=...`` for that path. + """ + if proposal is not None: + if sampling is not None: + raise ValueError( + "sampling= describes target-Metropolis burn-in and " + "thinning; pass n_samples=... for proposal samples." + ) + n_samples = kwargs.pop("n_samples", 128) + seed = kwargs.pop("seed", None) + fermion = kwargs.pop("fermion", self.fermion) + one_d_to_two_d = kwargs.pop("one_d_to_two_d", None) + occupation_map = kwargs.pop("occupation_map", None) + sample_kwargs = kwargs.pop("sample_kwargs", None) + progress = kwargs.pop("progress", False) + amplitude_floor = kwargs.pop("amplitude_floor", 0.0) + if kwargs: + unexpected = ", ".join(sorted(kwargs)) + raise TypeError( + "Unsupported keyword arguments for proposal sampling: " + f"{unexpected}." + ) + self._ensure_initialized( + contraction=contraction, + contraction_opts=contraction_opts, + ) + return self.sample_from_proposal( + proposal, + n_samples=n_samples, + seed=seed, + fermion=fermion, + one_d_to_two_d=one_d_to_two_d, + occupation_map=occupation_map, + sample_kwargs=sample_kwargs, + progress=progress, + amplitude_floor=amplitude_floor, + ) + self._ensure_initialized( + sampling=sampling, + contraction=contraction, + contraction_opts=contraction_opts, + n_walkers=kwargs.get("n_chains"), + ) + return super().sample(sampling=sampling, **kwargs) + + def measure( + self, + samples, + observables=None, + *, + amplitudes=None, + weights=None, + proposal_log_probs=None, + profile=False, + deduplicate=True, + progress=False, + _include_energy=False, + ): + """Measure observables from retained samples without resampling. + + ``samples`` normally comes from :meth:`sample`; its stored parent + amplitudes are reused. Values in ``observables`` follow :meth:`run`'s + native mapping convention, including an explicit + ``{"energy": terms}`` entry. + """ + self._ensure_initialized() + compiled = self._measurement_observables( + observables, + include_energy=_include_energy or observables is None, + ) + return self.measure_samples( + samples, + observables=compiled, + amplitudes=amplitudes, + weights=weights, + proposal_log_probs=proposal_log_probs, + profile=profile, + deduplicate=deduplicate, + progress=progress, + ) + + def warmup( + self, + *, + sampling=None, + contraction=None, + contraction_opts=None, + n_sweeps=0, + progress=False, + ): + """Eagerly evaluate one PEPS amplitude and optionally equilibrate walkers. + + The direct amplitude evaluation initializes lazy contraction work with + a valid sector-preserving configuration. It is deliberately separate + from burn-in, which mutates the Markov chains only when + ``n_sweeps > 0``. + """ + if isinstance(n_sweeps, bool) or not isinstance(n_sweeps, int) or n_sweeps < 0: + raise ValueError("n_sweeps must be a non-negative integer.") + self._ensure_initialized( + sampling=sampling, + contraction=contraction, + contraction_opts=contraction_opts, + ) + torch = _require_torch() + start = time.perf_counter() + with torch.no_grad(): + config = self.configs[:1].detach().clone() + amplitude = _call_amplitude_fn( + self.model, + config, + chunk_size=self.chunk_size, + )[0].detach().clone() + burn_in = None + if n_sweeps: + burn_in = self.burn_in(n_sweeps, progress=progress) + return TorchVMCWarmupResult( + config=config[0], + amplitude=amplitude, + n_sweeps=n_sweeps, + elapsed_seconds=time.perf_counter() - start, + burn_in=burn_in, + ) + + def run_measurement( + self, + observables=None, + *, + sampling=None, + contraction=None, + contraction_opts=None, + warmup=True, + warmup_sweeps=0, + progress=False, + profile=False, + ): + """Warm up, sample, and estimate PEPS Fermion observables once. + + This is the concise measurement workflow. The returned record keeps + the warm-up amplitude, the exact chain-preserving samples, and the + observable estimates. ``progress=True`` reports optional burn-in, + MCMC sampling, then the connection/contraction/statistics phases. + """ + self._ensure_initialized( + sampling=sampling, + contraction=contraction, + contraction_opts=contraction_opts, + ) + start = time.perf_counter() + warmup_result = ( + self.warmup(n_sweeps=warmup_sweeps, progress=progress) + if warmup + else None + ) + samples = self.sample(sampling=sampling, progress=progress) + estimates = self.measure( + samples, + observables=observables, + profile=profile, + progress=progress, + _include_energy=True, + ) + return TorchVMCMeasurementRun( + warmup=warmup_result, + samples=samples, + estimates=estimates, + elapsed_seconds=time.perf_counter() - start, + ) + + def run( + self, + n_steps=None, + *, + observables=None, + sampling=None, + contraction=None, + contraction_opts=None, + warmup=None, + warmup_sweeps=0, + progress=False, + **kwargs, + ): + """Run either a PEPS measurement workflow or optimization updates. + + With no ``n_steps`` (or an observable mapping as the first argument), + this is an alias for :meth:`run_measurement` and defaults to one eager + amplitude warm-up. The first measurement receives the chain recipe in + ``sampling`` and the PEPS recipe in ``contraction_opts``. Pass an + integer ``n_steps`` to retain the + established optimization alias for :meth:`TorchVMCDriver.optimize`. + Keeping the two modes distinct avoids treating a measurement as an + optimization step while preserving existing ``run(n_steps=...)`` code. + """ + if n_steps is not None and hasattr(n_steps, "items"): + if observables is not None: + raise TypeError( + "Pass observables either positionally or as observables=..., " + "not both." + ) + observables = n_steps + n_steps = None + if n_steps is not None: + if ( + observables is not None + or sampling is not None + or contraction is not None + or contraction_opts is not None + or warmup_sweeps != 0 + ): + raise ValueError( + "observables, sampling, contraction settings, and " + "warmup_sweeps apply only to measurement runs; omit " + "n_steps to use them." + ) + if warmup is not None: + raise ValueError("warmup applies only to a measurement run.") + self._ensure_initialized() + return super().run(n_steps, progress=progress, **kwargs) + profile = kwargs.pop("profile", False) + if kwargs: + unexpected = ", ".join(sorted(kwargs)) + raise TypeError(f"Unexpected measurement run keyword(s): {unexpected}.") + return self.run_measurement( + observables, + sampling=sampling, + contraction=contraction, + contraction_opts=contraction_opts, + warmup=True if warmup is None else bool(warmup), + warmup_sweeps=warmup_sweeps, + progress=progress, + profile=profile, + ) + def make_bp_sampler( self, proposal_sampler=None, @@ -500,6 +1028,7 @@ def make_bp_sampler( sampler_seed=None, ): """Create a symmetry-aware BP independence sampler from this PEPS.""" + self._ensure_initialized(n_walkers=n_chains) if proposal_sampler is None: from ...sampling import PepsBpSampler # pylint: disable=import-outside-toplevel diff --git a/src/pepsy/vmc/torch/importance.py b/src/pepsy/vmc/torch/importance.py index f25f731..f4e939e 100644 --- a/src/pepsy/vmc/torch/importance.py +++ b/src/pepsy/vmc/torch/importance.py @@ -8,10 +8,15 @@ from __future__ import annotations import inspect +import time from ._common import _as_long_matrix, _model_device, _proposal_log_probabilities from .amplitude import _call_amplitude_fn -from .results import TorchVMCImportanceEstimate +from .results import ( + TorchImportanceSamples, + TorchVMCImportanceEstimate, + _torch_sample_provenance, +) from ..torch_types import ( FermionSiteEncoding, SpinlessSiteEncoding, @@ -19,7 +24,7 @@ _require_torch, ) -__all__ = ["measure_from_proposal"] +__all__ = ["measure_from_proposal", "sample_from_proposal"] def _target_site_order(driver, site_order): @@ -395,7 +400,7 @@ def _bridge_samples( return configs[valid], amplitudes[valid], log_q[valid], int(configs.shape[0]) -def measure_from_proposal( +def sample_from_proposal( driver, proposal, *, @@ -406,21 +411,20 @@ def measure_from_proposal( site_order=None, occupation_map=None, sample_kwargs=None, - observables=None, progress=False, amplitude_floor=0.0, - profile=False, - deduplicate=True, ): - """Measure VMC observables from an MPS, BP, tree, or proposal batch. + """Draw and bridge reusable samples from an MPS, BP, tree, or proposal. - The returned value is exactly the result of - :meth:`TorchVMCDriver.measure_samples`: one - :class:`TorchVMCEnergyEstimate` for the default Hamiltonian or a mapping - of estimates when ``observables`` is supplied. + The resulting :class:`TorchImportanceSamples` stores PEPS-code + configurations, the fixed proposal density ``log q(x)``, and the current + target amplitudes. Pass it to :meth:`TorchVMCDriver.measure_samples` (or + :meth:`TorchFermionVMC.measure`) to estimate any number of observables + without drawing the MPS proposal again. """ n_samples = _check_positive_int("n_samples", n_samples) - configs, amplitudes, log_q, _ = _bridge_samples( + start = time.perf_counter() + configs, amplitudes, log_q, n_drawn = _bridge_samples( driver, proposal, n_samples=n_samples, @@ -433,13 +437,57 @@ def measure_from_proposal( sample_kwargs=sample_kwargs, amplitude_floor=amplitude_floor, ) - return driver.measure_samples( - configs, - observables=observables, + elapsed = time.perf_counter() - start + n_valid = int(configs.shape[0]) + return TorchImportanceSamples( + configs=configs, amplitudes=amplitudes, proposal_log_probs=log_q, + n_samples=n_valid, + n_drawn=n_drawn, + elapsed_seconds=elapsed, + samples_per_second=n_valid / elapsed if elapsed > 0 else float("inf"), + target_provenance=_torch_sample_provenance(driver.model), + ) + + +def measure_from_proposal( + driver, + proposal, + *, + n_samples=128, + seed=None, + fermion=None, + one_d_to_two_d=None, + site_order=None, + occupation_map=None, + sample_kwargs=None, + observables=None, + progress=False, + amplitude_floor=0.0, + profile=False, + deduplicate=True, +): + """Compatibility one-shot wrapper around sample then measure.""" + samples = sample_from_proposal( + driver, + proposal, + n_samples=n_samples, + seed=seed, + fermion=fermion, + one_d_to_two_d=one_d_to_two_d, + site_order=site_order, + occupation_map=occupation_map, + sample_kwargs=sample_kwargs, + progress=progress, + amplitude_floor=amplitude_floor, + ) + return driver.measure_samples( + samples, + observables=observables, profile=profile, deduplicate=deduplicate, + progress=progress, ) diff --git a/src/pepsy/vmc/torch/proposals.py b/src/pepsy/vmc/torch/proposals.py index ef212a4..612fce8 100644 --- a/src/pepsy/vmc/torch/proposals.py +++ b/src/pepsy/vmc/torch/proposals.py @@ -712,7 +712,7 @@ def _warmup_proposal_mix( bar = _make_progress( progress, total=n_sweeps, - desc="Torch VMC proposal warm-up", + desc="Metropolis proposal warm-up", unit="sweep", ) try: diff --git a/src/pepsy/vmc/torch/results.py b/src/pepsy/vmc/torch/results.py index fbb12ad..81f4f9d 100644 --- a/src/pepsy/vmc/torch/results.py +++ b/src/pepsy/vmc/torch/results.py @@ -7,8 +7,10 @@ from __future__ import annotations +from collections.abc import Mapping from dataclasses import dataclass from numbers import Integral +from types import MappingProxyType from typing import Any import numpy as np @@ -33,6 +35,61 @@ def acceptance_rate(self): return self.n_accepted / self.n_proposed +def _freeze_provenance_value(value): + """Return a stable, equality-comparable description of an option value.""" + if isinstance(value, Mapping): + return tuple( + (str(key), _freeze_provenance_value(item)) + for key, item in sorted(value.items(), key=lambda pair: str(pair[0])) + ) + if isinstance(value, (tuple, list)): + return tuple(_freeze_provenance_value(item) for item in value) + if isinstance(value, (str, bytes, bool, int, float, type(None))): + return value + return repr(value) + + +@dataclass(frozen=True) +class TorchSampleProvenance: + """Identity of the amplitude model and contraction used to draw samples. + + The native measurement path compares this record with its current model + before reusing stored parent amplitudes. It prevents mixing configurations + drawn from one PEPS state with local estimators from a later state. + """ + + model_type: str + model_identity: int + parameter_versions: tuple[int, ...] + contraction_signature: tuple[Any, Any, Any, Any] + + +def _torch_sample_provenance(model): + """Capture the mutable model state relevant to stored MCMC amplitudes.""" + parameters = getattr(model, "parameters", None) + if callable(parameters): + try: + parameter_versions = tuple( + int(getattr(parameter, "_version", 0)) + for parameter in parameters() + ) + except (RuntimeError, TypeError, ValueError): + parameter_versions = () + else: + parameter_versions = () + return TorchSampleProvenance( + model_type=f"{type(model).__module__}.{type(model).__qualname__}", + model_identity=id(model), + parameter_versions=parameter_versions, + contraction_signature=( + _freeze_provenance_value(getattr(model, "contraction", None)), + _freeze_provenance_value(getattr(model, "chi", None)), + _freeze_provenance_value(getattr(model, "cutoff", None)), + _freeze_provenance_value(getattr(model, "contraction_opts", None)), + ), + ) + + @dataclass(frozen=True) class TorchMCMCSamples: """Chain-preserving samples and diagnostics from a torch sampler. @@ -40,7 +97,9 @@ class TorchMCMCSamples: ``configs`` and ``amplitudes`` have shape ``(n_samples_per_chain, n_chains, ...)``. ``n_samples`` is the actual number of returned samples, so it can be larger than the requested total - when that total is not divisible by ``n_chains``. + when that total is not divisible by ``n_chains``. Native samplers attach + ``provenance`` so a later measurement can reject a batch after its PEPS or + contraction settings have changed. """ configs: Any @@ -57,6 +116,7 @@ class TorchMCMCSamples: samples_per_second: float log_abs_amplitudes: Any = None proposal_stats: Any = None + provenance: TorchSampleProvenance | None = None def diagnostics(self, values=None, *, max_lag=None): """Compute chain diagnostics for a scalar observable. @@ -95,6 +155,46 @@ def to_common(self): ) +@dataclass(frozen=True) +class TorchImportanceSamples: + """Reusable PEPS configurations independently drawn from a proposal. + + ``proposal_log_probs`` stores ``log q(x)`` while ``amplitudes`` stores the + target PEPS values evaluated when the batch was bridged. Unlike Markov + samples, the configurations remain valid after a PEPS update because + their distribution is the fixed external proposal. The driver detects a + changed ``target_provenance`` and refreshes the parent PEPS amplitudes + before forming the importance weights. + """ + + configs: Any + amplitudes: Any + proposal_log_probs: Any + n_samples: int + n_drawn: int + elapsed_seconds: float + samples_per_second: float + target_provenance: TorchSampleProvenance | None = None + + def to_common(self): + """Convert to a backend-neutral externally weighted sample batch.""" + from ..api import VMCSamples + + return VMCSamples( + configs=self.configs, + amplitudes=self.amplitudes, + proposal_log_probs=self.proposal_log_probs, + diagnostics={ + "sample_source": "external-proposal", + "n_samples": self.n_samples, + "n_drawn": self.n_drawn, + "elapsed_seconds": self.elapsed_seconds, + "samples_per_second": self.samples_per_second, + }, + native=self, + ) + + @dataclass(frozen=True) class TorchChainDiagnostics: """MCMC convergence diagnostics for chain-shaped scalar values.""" @@ -165,6 +265,45 @@ class TorchVMCEnergyEstimate: proposal_log_probs: Any = None +@dataclass(frozen=True) +class TorchVMCWarmupResult: + """One eager PEPS-amplitude evaluation and optional walker burn-in. + + ``config`` and ``amplitude`` are a representative valid walker and its + freshly evaluated PEPS amplitude. ``burn_in`` is populated only when the + caller also asks for Metropolis equilibration sweeps. + """ + + config: Any + amplitude: Any + n_sweeps: int + elapsed_seconds: float + burn_in: TorchMetropolisResult | None = None + + +@dataclass(frozen=True) +class TorchVMCMeasurementRun: + """Result of the high-level fermionic PEPS measurement workflow. + + The record keeps the warm-up result, exact chain-preserving samples, and + all observable estimates separate so callers can reuse or inspect each + stage without rerunning the Markov chain. + """ + + warmup: TorchVMCWarmupResult | None + samples: TorchMCMCSamples + estimates: Mapping[str, TorchVMCEnergyEstimate] + elapsed_seconds: float + + def __post_init__(self): + object.__setattr__(self, "estimates", MappingProxyType(dict(self.estimates))) + + @property + def energy(self): + """Return the Hamiltonian estimate, when the run included energy.""" + return self.estimates.get("energy") + + @dataclass(frozen=True) class TorchVMCImportanceEstimate: """Energy estimate from an external proposal distribution.""" @@ -223,17 +362,112 @@ def _progress_scalar(value): return float(np.real(value)) -def _set_vmc_progress_postfix(bar, result, *, n_sites, include_energy=True): - """Update a VMC progress bar without affecting the numerical workflow.""" +def _model_progress_fields(model): + """Return short, display-only contraction and cache fields for a model.""" + if model is None: + return {} + fields = {} + contraction = getattr(model, "contraction", None) + if contraction is not None: + chi = getattr(model, "chi", None) + fields["amp"] = ( + str(contraction) + if chi is None + else f"{contraction} chi={chi}" + ) + return fields + + +def _proposal_environment_progress(model): + """Summarize the most recent boundary proposal-cache activity.""" + stats = getattr(model, "last_proposal_cache_stats", None) + if not stats: + return None + environment_hits = int(stats.get("num_environment_cache_hits", 0)) + environment_builds = int(stats.get("num_environment_builds", 0)) + transition_hits = int(stats.get("num_transition_cache_hits", 0)) + vmapped = int(stats.get("num_vmapped", 0)) + parts = [] + if environment_hits or environment_builds: + parts.append(f"{environment_hits} reuse/{environment_builds} build") + if transition_hits: + parts.append(f"{transition_hits} transition") + if vmapped: + parts.append(f"vmap={vmapped}") + return ",".join(parts) or None + + +def _connected_target_progress(model): + """Summarize the latest local-estimator target-amplitude route.""" + stats = getattr(model, "last_connected_reuse_stats", None) + if not stats: + return {} + fields = {} + diagonal = int(stats.get("num_diagonal", 0)) + reused = int(stats.get("num_reused", 0)) + batched = int(stats.get("num_batched", 0)) + fallback = int(stats.get("num_fallback", 0)) + if diagonal or reused or batched or fallback: + fields["targets"] = ( + f"diag={diagonal}, env={reused}, " + f"batch={batched}, direct={fallback}" + ) + environment_hits = int(stats.get("num_environment_cache_hits", 0)) + environment_builds = int(stats.get("num_environment_builds", 0)) + if environment_hits or environment_builds: + fields["env"] = f"{environment_hits} reuse/{environment_builds} build" + return fields + + +def _display_observables(observables): + """Keep a progress-bar observable list short enough for notebooks.""" + names = tuple(str(name) for name in observables) + if len(names) <= 3: + return ",".join(names) + return ",".join(names[:3]) + f",+{len(names) - 3}" + + +def _set_vmc_progress_postfix( + bar, + result=None, + *, + n_sites=None, + include_energy=True, + n_chains=None, + model=None, + proposal=None, + retained_per_walker=None, + burn_in=None, + thin=None, + phase=None, +): + """Update a Metropolis/VMC bar without affecting numerical work.""" if bar is None: return - postfix = {"accept": f"{result.acceptance_rate:.3f}"} + postfix = _model_progress_fields(model) + if n_chains is not None: + postfix["walkers"] = int(n_chains) + if proposal is not None: + postfix["move"] = str(proposal) + if retained_per_walker is not None and n_chains is not None: + postfix["retain"] = f"{int(retained_per_walker)}x{int(n_chains)}" + if burn_in is not None: + postfix["burn"] = int(burn_in) + if thin is not None: + postfix["thin"] = int(thin) + if phase is not None: + postfix["phase"] = str(phase) + proposal_environment = _proposal_environment_progress(model) + if proposal_environment is not None: + postfix["env"] = proposal_environment + if result is not None: + postfix["accept"] = f"{result.acceptance_rate:.3f}" no_op_rate = _proposal_no_op_rate( getattr(result, "proposal_stats", None) ) if no_op_rate is not None: postfix["no-op"] = f"{no_op_rate:.3f}" - if include_energy: + if include_energy and result is not None and n_sites is not None: postfix["E/site"] = ( f"{_progress_scalar(result.energy_mean) / n_sites:+.6f}" ) @@ -247,6 +481,33 @@ def _set_vmc_progress_postfix(bar, result, *, n_sites, include_energy=True): set_postfix(postfix) +def _set_evaluation_progress_postfix( + bar, + *, + model, + n_steps, + n_chains, + observables, + parent_amplitudes, + stage, + n_connections=None, +): + """Describe shared local-estimator work on the Evaluation progress bar.""" + if bar is None: + return + postfix = _model_progress_fields(model) + postfix["samples"] = f"{int(n_steps)}x{int(n_chains)}" + postfix["obs"] = _display_observables(observables) + postfix["parent psi"] = parent_amplitudes + if n_connections is not None: + postfix["connections"] = int(n_connections) + postfix.update(_connected_target_progress(model)) + postfix["stage"] = stage + set_postfix = getattr(bar, "set_postfix", None) + if callable(set_postfix): + set_postfix(postfix) + + def _cache_profile_snapshot(model): """Copy lightweight model-cache counters for an opt-in VMC profile.""" snapshot = {} @@ -278,8 +539,10 @@ def _accumulate_cache_profile(total, snapshot): __all__ = [ "TorchChainDiagnostics", + "TorchImportanceSamples", "TorchMCMCSamples", "TorchMetropolisResult", + "TorchSampleProvenance", "TorchVMCImportanceEstimate", "TorchVMCEnergyEstimate", "TorchVMCStepResult", @@ -287,5 +550,6 @@ def _accumulate_cache_profile(total, snapshot): "_cache_profile_snapshot", "_make_progress", "_progress_scalar", + "_set_evaluation_progress_postfix", "_set_vmc_progress_postfix", ] diff --git a/src/pepsy/vmc/torch/sampler.py b/src/pepsy/vmc/torch/sampler.py index 1b2bf5d..b7cf8ee 100644 --- a/src/pepsy/vmc/torch/sampler.py +++ b/src/pepsy/vmc/torch/sampler.py @@ -30,6 +30,7 @@ TorchMetropolisResult, _make_progress, _set_vmc_progress_postfix, + _torch_sample_provenance, ) __all__ = [ @@ -310,9 +311,19 @@ def burn_in( bar = _make_progress( True, total=n_sweeps, - desc="Torch VMC burn-in", + desc="Metropolis warm-up", unit="sweep", ) + _set_vmc_progress_postfix( + bar, + n_sites=self.n_sites, + include_energy=False, + n_chains=self.n_chains, + model=self.amplitude_fn, + proposal=self.proposal, + burn_in=0, + thin=1, + ) result = None n_proposed = 0 n_accepted = 0 @@ -333,6 +344,11 @@ def burn_in( result, n_sites=self.n_sites, include_energy=False, + n_chains=self.n_chains, + model=self.amplitude_fn, + proposal=self.proposal, + burn_in=0, + thin=1, ) finally: bar.close() @@ -388,7 +404,10 @@ def sample( ``n_samples`` is the requested total across all chains. As in NetKet, the chain length is rounded up so every chain contributes the same number of samples. ``n_discard`` and ``n_thin`` are aliases for - ``n_discard_per_chain`` and ``sweep_size`` respectively. + ``n_discard_per_chain`` and ``sweep_size`` respectively. Both the + discard and retained portions use that sweep interval, so the progress + bar totals ``(n_discard_per_chain + n_samples_per_chain) * + sweep_size`` batched Metropolis sweeps. """ torch = _require_torch() n_samples = _check_positive_int("n_samples", n_samples) @@ -416,25 +435,60 @@ def sample( bar = _make_progress( progress, total=total_sweeps, - desc="Torch Metropolis", + desc="Metropolis", + unit="sweep", + ) + _set_vmc_progress_postfix( + bar, + n_sites=self.n_sites, + include_energy=False, + n_chains=self.n_chains, + model=self.amplitude_fn, + proposal=self.proposal, + retained_per_walker=n_samples_per_chain, + burn_in=n_discard_per_chain, + thin=sweep_size, + phase=("equilibrate" if n_discard_per_chain else "retain 0/" + f"{n_samples_per_chain}"), ) start = time.perf_counter() n_proposed = 0 n_accepted = 0 + n_completed_sweeps = 0 proposal_stats = _empty_proposal_stats() if track_proposal_stats else None + def sampling_phase(): + if n_completed_sweeps <= n_discard_per_chain * sweep_size: + return "equilibrate" + retained = (n_completed_sweeps - n_discard_per_chain * sweep_size) // sweep_size + return f"retain {retained}/{n_samples_per_chain}" + def advance_one_sweep(): - nonlocal n_proposed, n_accepted + nonlocal n_proposed, n_accepted, n_completed_sweeps sweep_kwargs = {"n_sweeps": 1} if track_proposal_stats: sweep_kwargs["track_proposal_stats"] = True result = self.sample_sweep(**sweep_kwargs) n_proposed += result.n_proposed n_accepted += result.n_accepted + n_completed_sweeps += 1 if track_proposal_stats: _merge_proposal_stats(proposal_stats, result.proposal_stats) if bar is not None: bar.update(1) + _set_vmc_progress_postfix( + bar, + result, + n_sites=self.n_sites, + include_energy=False, + n_chains=self.n_chains, + model=self.amplitude_fn, + proposal=self.proposal, + retained_per_walker=n_samples_per_chain, + burn_in=n_discard_per_chain, + thin=sweep_size, + phase=sampling_phase(), + ) for _ in range(n_discard_per_chain * sweep_size): advance_one_sweep() @@ -485,6 +539,7 @@ def advance_one_sweep(): ), log_abs_amplitudes=log_abs_amplitudes, proposal_stats=proposal_stats, + provenance=_torch_sample_provenance(self.amplitude_fn), ) diff --git a/tests/test_vmc_api.py b/tests/test_vmc_api.py index c3868c5..e62b742 100644 --- a/tests/test_vmc_api.py +++ b/tests/test_vmc_api.py @@ -2,6 +2,7 @@ import numpy as np import pytest +from types import SimpleNamespace from pepsy.vmc import ( BackendCapabilityWarning, @@ -275,6 +276,37 @@ def forward(self, configs): ) assert samples.configs.shape == (2, 2, 2) assert samples.to_common().chain_shape == (2, 2) + assert samples.provenance is not None + assert samples.provenance.model_identity == id(driver.model) + + with torch.no_grad(): + driver.model.weights.add_(0.1) + with pytest.raises(RuntimeError, match="different PEPS/model state"): + driver.measure_samples(samples) + + samples = driver.sample( + sampling=SamplingConfig( + n_samples_per_chain=1, + n_chains=2, + burn_in=0, + thin=1, + seed=14, + ) + ) + driver.model.contraction = "exact" + with pytest.raises(RuntimeError, match="different PEPS/model state"): + driver.measure_samples(samples) + + measurement = driver.estimate_observable( + sampling=SamplingConfig( + n_samples_per_chain=1, + n_chains=2, + burn_in=0, + thin=1, + seed=13, + ) + ) + assert measurement.n_samples == 2 history = driver.optimize( optimization=OptimizationConfig( @@ -288,6 +320,247 @@ def forward(self, configs): assert len(history) == 1 +def test_torch_fermion_measurement_api_keeps_sampling_and_estimation_separate( + monkeypatch, +): + torch = pytest.importorskip("torch") + from pepsy.vmc import ( + TorchFermionVMC, + TorchVMCMeasurementRun, + TorchVMCWarmupResult, + ) + from pepsy.vmc.torch import TorchVMCDriver + + vmc = object.__new__(TorchFermionVMC) + vmc.observables = {"density": "compiled-density"} + vmc._compile_observables = lambda observables: { + name: f"compiled:{value}" for name, value in observables.items() + } + + seen = {} + vmc._ensure_initialized = lambda **kwargs: seen.setdefault( + "ensure", kwargs + ) + + def fake_estimate(self, observables, **kwargs): + seen["estimate_observables"] = observables + seen["estimate_kwargs"] = kwargs + return observables + + monkeypatch.setattr(TorchVMCDriver, "estimate_observables", fake_estimate) + sampling = SamplingConfig( + n_samples_per_chain=2, + n_chains=3, + burn_in=4, + thin=2, + seed=9, + ) + estimates = vmc.estimate_observables(sampling=sampling) + assert estimates == {"energy": None, "density": "compiled-density"} + assert seen["estimate_kwargs"] == {"sampling": sampling} + + class ConstantAmplitude(torch.nn.Module): + def forward(self, configs): + return torch.ones(configs.shape[0], dtype=torch.float64) + + warmup_driver = object.__new__(TorchFermionVMC) + warmup_driver.configs = torch.tensor([[1, 2], [2, 1]], dtype=torch.long) + warmup_driver.model = ConstantAmplitude() + warmup_driver.chunk_size = None + warmup_driver._ensure_initialized = lambda **kwargs: None + warmup = warmup_driver.warmup() + assert warmup.config.tolist() == [1, 2] + assert warmup.amplitude.item() == pytest.approx(1.0) + assert warmup.n_sweeps == 0 + + native_samples = object() + vmc.warmup = lambda *, n_sweeps, progress: warmup + vmc.sample = lambda *, sampling, progress: native_samples + + def fake_measure(samples, *, observables, profile, progress, **kwargs): + seen["run_observables"] = observables + seen["measure_kwargs"] = kwargs + assert samples is native_samples + assert not profile + assert not progress + return {name: name for name in observables} + + vmc.measure_samples = fake_measure + run = vmc.run(sampling=sampling, progress=False) + assert isinstance(run, TorchVMCMeasurementRun) + assert isinstance(run.warmup, TorchVMCWarmupResult) + assert run.samples is native_samples + assert run.energy == "energy" + assert run.estimates == {"energy": "energy", "density": "density"} + assert seen["run_observables"] == { + "energy": None, + "density": "compiled-density", + } + assert seen["measure_kwargs"] == { + "amplitudes": None, + "weights": None, + "proposal_log_probs": None, + "deduplicate": True, + } + with pytest.raises(TypeError): + run.estimates["other"] = "not allowed" + + positional_run = vmc.run({"eta": "raw-eta"}, sampling=sampling) + assert positional_run.estimates == {"energy": "energy", "eta": "eta"} + assert seen["run_observables"] == { + "energy": None, + "eta": "compiled:raw-eta", + } + + explicit_energy_run = vmc.run( + observables={"energy": "raw-energy", "eta": "raw-eta"}, + sampling=sampling, + contraction_opts={ + "method": "boundary", + "chi": 4, + "cutoff": 1.0e-10, + "mode": "mps", + }, + ) + assert explicit_energy_run.estimates == {"energy": "energy", "eta": "eta"} + assert seen["run_observables"] == { + "energy": "compiled:raw-energy", + "eta": "compiled:raw-eta", + } + + reused = vmc.measure( + native_samples, + {"energy": "raw-energy", "eta": "raw-eta"}, + ) + assert reused == {"energy": "energy", "eta": "eta"} + assert seen["run_observables"] == { + "energy": "compiled:raw-energy", + "eta": "compiled:raw-eta", + } + + def fake_optimization_run(self, n_steps, *, progress, **kwargs): + return n_steps, progress, kwargs + + monkeypatch.setattr(TorchVMCDriver, "run", fake_optimization_run) + assert vmc.run(3, progress=True, profile=True) == ( + 3, + True, + {"profile": True}, + ) + + +def test_torch_fermion_vmc_lazy_setup_uses_run_sampling_and_contraction_opts(): + from pepsy.vmc import TorchFermionVMC + + vmc = object.__new__(TorchFermionVMC) + vmc._driver_initialized = False + vmc._contraction_config = None + vmc._legacy_contraction_config = None + vmc._initial_configs = None + vmc._initial_n_walkers = None + seen = {} + + def fake_initialize(contraction, *, n_walkers): + seen["contraction"] = contraction + seen["n_walkers"] = n_walkers + + vmc._initialize_driver = fake_initialize + sampling = SamplingConfig(n_samples_per_chain=2, n_chains=3) + vmc._ensure_initialized( + sampling=sampling, + contraction_opts={ + "method": "boundary", + "chi": 8, + "cutoff": 1.0e-9, + "mode": "mps", + }, + ) + + assert seen["n_walkers"] == sampling.n_chains + assert seen["contraction"] == ContractionConfig( + method="boundary", + chi=8, + cutoff=1.0e-9, + options={"mode": "mps"}, + ) + + +def test_torch_vmc_progress_postfix_reports_contraction_and_reuse(): + from pepsy.vmc.torch.results import ( + _set_evaluation_progress_postfix, + _set_vmc_progress_postfix, + ) + + class Progress: + postfix = None + + def set_postfix(self, postfix): + self.postfix = dict(postfix) + + model = SimpleNamespace( + contraction="boundary", + chi=16, + last_proposal_cache_stats={ + "num_environment_cache_hits": 2, + "num_environment_builds": 1, + "num_transition_cache_hits": 3, + "num_vmapped": 0, + }, + last_connected_reuse_stats={ + "num_diagonal": 5, + "num_reused": 7, + "num_batched": 0, + "num_fallback": 2, + "num_environment_cache_hits": 4, + "num_environment_builds": 1, + }, + ) + bar = Progress() + _set_vmc_progress_postfix( + bar, + SimpleNamespace(acceptance_rate=0.75, proposal_stats=None, sr=None), + n_sites=4, + include_energy=False, + n_chains=3, + model=model, + proposal="spinful", + retained_per_walker=2, + burn_in=4, + thin=2, + ) + assert bar.postfix == { + "amp": "boundary chi=16", + "walkers": 3, + "move": "spinful", + "retain": "2x3", + "burn": 4, + "thin": 2, + "env": "2 reuse/1 build,3 transition", + "accept": "0.750", + } + + _set_evaluation_progress_postfix( + bar, + model=model, + n_steps=2, + n_chains=3, + observables=("energy", "eta"), + parent_amplitudes="stored", + n_connections=14, + stage="statistics", + ) + assert bar.postfix == { + "amp": "boundary chi=16", + "samples": "2x3", + "obs": "energy,eta", + "parent psi": "stored", + "connections": 14, + "targets": "diag=5, env=7, batch=0, direct=2", + "env": "4 reuse/1 build", + "stage": "statistics", + } + + def test_torch_vmc_modules_own_implementations_and_keep_core_aliases(): from pepsy.vmc.torch import _common, _core, _graded, amplitude, sr @@ -516,6 +789,166 @@ class Ansatz: ) assert samples.chain_shape == (3, 2) assert samples.configs.shape == (3, 2, 4) + assert samples.diagnostics["n_samples"] == 6 + assert samples.diagnostics["n_chains"] == 2 + assert samples.diagnostics["burn_in"] == 0 + + +def test_netket_optimize_result_reports_compile_and_step_timings(): + from pepsy.vmc.netket import VMCOptimizeResult + + result = VMCOptimizeResult( + steps=np.asarray([0, 1]), + energies=np.asarray([-1.0, -1.2]), + errors=np.asarray([0.1, 0.05]), + variances=np.asarray([0.2, 0.1]), + final_energy=-1.2, + final_error=0.05, + compile_seconds=3.0, + optimization_seconds=8.0, + total_seconds=11.5, + ) + + assert result.warmup_seconds == 3.0 + assert result.optimization_seconds_per_step == 4.0 + assert result.total_seconds == 11.5 + + +def test_netket_build_timing_reports_slowest_phase(): + from pepsy.vmc.netket import NetKetBuildTiming + + timing = NetKetBuildTiming( + settings_seconds=0.1, + geometry_seconds=0.2, + hamiltonian_seconds=2.0, + peps_seconds=0.3, + model_seconds=0.4, + sampler_seconds=1.0, + total_seconds=4.0, + ) + + assert timing.slowest_phase == ("hamiltonian_seconds", 2.0) + assert timing.as_dict()["total_seconds"] == 4.0 + + +def test_netket_boundary_zero_separation_has_flat_symmray_retry(): + import pepsy.vmc.netket as netket_vmc + + flat_array = type( + "Z2FermionicArrayFlat", + (), + {"__module__": "symmray.fake"}, + )() + + class Tensor: + data = flat_array + + class Network: + def __init__(self): + self.calls = [] + + def __iter__(self): + return iter((Tensor(),)) + + def contract_boundary(self, **kwargs): + self.calls.append(kwargs) + if len(self.calls) == 1: + raise ValueError("empty intermediate axis") + return (1.0, 0.0) + + netket_vmc._FLAT_SYMMRAY_BOUNDARY_FALLBACK_WARNED = False + network = Network() + with pytest.warns(RuntimeWarning, match="max_separation=0"): + result = netket_vmc._contract_boundary_for_vmc( + network, + max_bond=8, + cutoff=0.0, + method_opts={"max_separation": 0, "canonize": True}, + ) + assert result == (1.0, 0.0) + assert network.calls[0]["max_separation"] == 0 + assert network.calls[1]["max_separation"] == 1 + assert network.calls[1]["canonize"] is True + + +def test_netket_amplitude_timing_reports_batch_average(): + from pepsy.vmc.netket import NetKetAmplitudeTiming + + timing = NetKetAmplitudeTiming(n_samples=4, amplitude_seconds=2.0) + + assert timing.amplitude_seconds_per_sample == 0.5 + assert timing.as_dict()["n_samples"] == 4 + + +def test_netket_setup_benchmarks_amplitude_without_sampling_in_timer(): + from pepsy.vmc.netket import NetKetPEPSVMC + + class Ansatz: + n_sites = 2 + n_params = 1 + + class Model: + def apply(self, variables, configs): + assert "params" in variables + return configs.sum(axis=-1) + + class State: + samples = np.zeros((1, 2, 2), dtype=np.int8) + parameters = {"x": 1.0} + + State.model = Model() + setup = NetKetPEPSVMC( + hilbert=None, + graph=None, + hamiltonian=None, + sampler=None, + vstate=State(), + model=State.model, + ansatz=Ansatz(), + config_map=None, + preconditioner=None, + ) + + timing = setup.benchmark_amplitude(n_samples=1) + assert timing.n_samples == 1 + assert timing.amplitude_seconds >= 0.0 + + +def test_netket_vmc_config_validates_shared_settings(): + from pepsy.vmc import ContractionConfig, NetKetVMCConfig, SamplingConfig + + config = NetKetVMCConfig( + contraction=ContractionConfig(method="boundary", chi=4), + sampling=SamplingConfig(n_samples_per_chain=2, n_chains=1), + sampler_sweep_size=8, + conserving=False, + use_sr=False, + progress=True, + ) + + assert config.contraction.chi == 4 + assert config.sampling.n_samples == 2 + assert config.sampler_sweep_size == 8 + + +def test_netket_native_geometry_helpers_accept_coordinate_terms(): + from pepsy.vmc.netket import ( + _edges_from_fermi_terms, + _infer_lattice_shape_from_fermi_terms, + _infer_pbc_from_fermi_terms, + ) + + terms = { + (0, 0): object(), + (0, 1): object(), + (1, 0): object(), + (1, 1): object(), + ((0, 0), (1, 0)): object(), + ((0, 1), (1, 1)): object(), + } + assert _infer_lattice_shape_from_fermi_terms(terms) == (2, 2) + assert _infer_pbc_from_fermi_terms(terms, 2, 2) == (False, False) + assert _edges_from_fermi_terms(terms, 2, 2) == ((0, 2), (1, 3)) def test_netket_compiler_lowers_common_fermion_terms(): diff --git a/tests/test_vmc_importance.py b/tests/test_vmc_importance.py index 12a7f9a..b491c95 100644 --- a/tests/test_vmc_importance.py +++ b/tests/test_vmc_importance.py @@ -29,6 +29,10 @@ def _driver(torch, *, zero_first=False): from pepsy.vmc import FermionSiteEncoding, TorchVMCDriver class Amplitude(torch.nn.Module): + def __init__(self): + super().__init__() + self.scale = torch.nn.Parameter(torch.tensor(1.0, dtype=torch.float64)) + def forward(self, configs): values = torch.ones(configs.shape[0], dtype=torch.float64) if zero_first: @@ -37,7 +41,7 @@ def forward(self, configs): torch.zeros_like(values), values, ) - return values + return self.scale * values return TorchVMCDriver( Amplitude(), @@ -64,8 +68,12 @@ def test_mps_batch_bridge_reorders_occupations_and_supports_observables(): torch.tensor([0.5, 0.5], dtype=torch.float64), ) - result = driver.measure_from_proposal( - batch, + samples = driver.sample_from_proposal(batch) + assert samples.configs.tolist() == [[1, 2], [2, 1]] + assert samples.proposal_log_probs.shape == (2,) + + result = driver.measure_samples( + samples, observables={ "energy": None, "eta": {(1, 0): torch.eye(4, dtype=torch.float64)}, @@ -78,6 +86,30 @@ def test_mps_batch_bridge_reorders_occupations_and_supports_observables(): assert result["energy"].effective_sample_size == pytest.approx(2.0) +def test_importance_samples_refresh_target_amplitudes_after_a_peps_update(): + torch = pytest.importorskip("torch") + driver = _driver(torch) + batch = _FermionBatch( + torch.tensor( + [ + [[1, 0], [0, 1]], + [[0, 1], [1, 0]], + ], + dtype=torch.long, + ), + torch.tensor([0.5, 0.5], dtype=torch.float64), + ) + + samples = driver.sample_from_proposal(batch) + assert torch.all(samples.amplitudes == 1.0) + with torch.no_grad(): + driver.model.scale.add_(1.0) + + result = driver.measure_samples(samples) + assert torch.all(result.amplitudes == 2.0) + assert result.effective_sample_size == pytest.approx(2.0) + + def test_mps_bridge_drops_zero_amplitude_nodes_before_local_energy(): torch = pytest.importorskip("torch") driver = _driver(torch, zero_first=True) @@ -174,11 +206,14 @@ def test_real_u1u1_mps_sampler_feeds_fermionic_peps_vmc(): seed=5, init_max_states=4096, ) - result = vmc.measure_from_mps( - proposal, + samples = vmc.sample( + proposal=proposal, n_samples=4, seed=8, + fermion=fermion, + one_d_to_two_d=one_d_to_two_d, ) + result = vmc.measure(samples)["energy"] assert result.configs.shape == (1, 4, 6) assert result.configs[0].tolist() == [[2, 2, 2, 1, 1, 1]] * 4 From a36df43b3488a8eb60546850cd79fd2cac254164 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Tue, 28 Jul 2026 07:57:12 -0700 Subject: [PATCH 07/70] Add backend-aware optimizer workflows --- src/pepsy/__init__.py | 4 + src/pepsy/backends/__init__.py | 2 + src/pepsy/backends/config.py | 29 ++- src/pepsy/optimizers/mps/optimizer.py | 134 ++++++++++ src/pepsy/optimizers/tree/optimizer.py | 339 ++++++++++++++++++------- src/pepsy/tensors/__init__.py | 2 + src/pepsy/tensors/contractions.py | 15 +- src/pepsy/tensors/core.py | 17 +- tests/test_backends.py | 18 ++ tests/test_optimize_mps.py | 58 +++++ tests/test_optimize_tree.py | 80 ++++++ 11 files changed, 599 insertions(+), 99 deletions(-) diff --git a/src/pepsy/__init__.py b/src/pepsy/__init__.py index 5d469c2..22609e4 100644 --- a/src/pepsy/__init__.py +++ b/src/pepsy/__init__.py @@ -44,6 +44,7 @@ "backend_jax": ".backends", "backend_numpy": ".backends", "backend_torch": ".backends", + "build_backend": ".backends", "get_default_array_backend": ".backends", "get_default_grad_backend": ".backends", "register_torch_linalg": ".backends", @@ -194,6 +195,7 @@ "SymPEPS": ".tensors", "add_cycle": ".tensors", "build_compressed_optimizer": ".tensors", + "build_contraction": ".tensors", "build_optimizer": ".tensors", "contract_hypercompressed_tn": ".tensors", "contract_hypercompressed_tn_batch": ".tensors", @@ -276,6 +278,7 @@ def __getattr__(name): from .backends import ( # noqa: F401 get_default_array_backend, get_default_grad_backend, + build_backend, register_torch_linalg, reset_default_backends, set_default_array_backend, @@ -348,6 +351,7 @@ def __getattr__(name): from .solvers import FDSolver # noqa: F401 from .tensors import ( # noqa: F401 OneDMap, + build_contraction, Fermion, FermionLatticeSetup, SpinfulFermion, diff --git a/src/pepsy/backends/__init__.py b/src/pepsy/backends/__init__.py index ea14ce6..2f2d208 100644 --- a/src/pepsy/backends/__init__.py +++ b/src/pepsy/backends/__init__.py @@ -15,6 +15,7 @@ backend_jax, backend_numpy, backend_torch, + build_backend, get_default_array_backend, get_default_grad_backend, register_torch_linalg, @@ -24,6 +25,7 @@ ) __all__ = [ + "build_backend", "backend_cupy", "backend_jax", "backend_numpy", diff --git a/src/pepsy/backends/config.py b/src/pepsy/backends/config.py index 3ebfaf2..b56b771 100644 --- a/src/pepsy/backends/config.py +++ b/src/pepsy/backends/config.py @@ -11,7 +11,7 @@ torch = None __all__ = [ - "backend_torch", "backend_numpy", "backend_cupy", "backend_jax", + "build_backend", "backend_torch", "backend_numpy", "backend_cupy", "backend_jax", "register_torch_linalg", "reg_rel_svd_torch", "reg_real_svd_torch", "reg_complex_svd_torch", "reg_real_qr_torch", "reg_complex_qr_torch", "reg_rel_svd_jax", "reg_real_svd_jax", "reg_complex_svd_jax", @@ -157,6 +157,33 @@ def cast_array(x, device=device, dtype=dtype, requires_grad=requires_grad): return cast_array +def build_backend(device="cpu", dtype=None, requires_grad=False, *, set_default=True): + """Build the standard Torch array backend, defaulting to CPU. + + The existing :func:`backend_torch` name remains unchanged. This helper is + the concise public entry point for workflows that want one backend + converter and one package-wide default:: + + import pepsy as py + to_backend = py.build_backend() # Torch CPU + + Parameters are forwarded to :func:`backend_torch`. By default the + resulting converter is also installed as Pepsy's default array backend; + pass ``set_default=False`` when only the returned converter should be + used. Explicit ``to_backend=`` / ``array_backend=`` arguments continue to + take precedence in individual APIs. + """ + + converter = backend_torch( + device=device, + dtype=dtype, + requires_grad=requires_grad, + ) + if set_default: + set_default_array_backend(converter) + return converter + + def backend_numpy(dtype=np.float64): """Return a converter that materializes arrays as NumPy arrays.""" diff --git a/src/pepsy/optimizers/mps/optimizer.py b/src/pepsy/optimizers/mps/optimizer.py index f0b42c2..72de4b1 100644 --- a/src/pepsy/optimizers/mps/optimizer.py +++ b/src/pepsy/optimizers/mps/optimizer.py @@ -1090,6 +1090,7 @@ def __init__( # pylint: disable=too-many-arguments,too-many-positional-argument self._unitary_initial_norm = None self._unitary_previous_norm = None self._unitary_global_norm_tracking = False + self._backend_mismatch_warned = False self._init_canonicalization() def _info_for_state(self, p, info=None): @@ -1327,6 +1328,7 @@ def set_p(self, p): self._su_gauges_state = None self._su_force_regauge = self.mode == "su" self.p_ungauged = None + self._backend_mismatch_warned = False self._init_canonicalization() def normalize(self, eps=1e-15, insert=None): @@ -1395,6 +1397,7 @@ def copy(self) -> "MpsOptimizer": copied._unitary_initial_norm = self._unitary_initial_norm copied._unitary_previous_norm = self._unitary_previous_norm copied._unitary_global_norm_tracking = self._unitary_global_norm_tracking + copied._backend_mismatch_warned = self._backend_mismatch_warned copied._su_gauges_supplied = True copied._su_gauges_ready = self._su_gauges_ready copied._su_gauges_state = copied.p if self._su_gauges_ready else None @@ -2364,6 +2367,13 @@ def _execute_mode( # pylint: disable=too-many-arguments,too-many-positional-arg ``event_seq`` must contain only ``"gate"``/``"submpo"`` events. Control events (measure/cap/reset) are handled by :meth:`_run_segmented`. """ + # Prepare gate and sub-MPO payloads once per executable segment. The + # converter returns already-compatible arrays/networks unchanged, + # while foreign payloads are moved to the backend owned by the live + # MPS. This keeps exact, simple-update, and compressed modes on one + # backend contract. + G_seq = self._prepare_gate_stream_backend(G_seq, event_seq) + if self.mode == "dmrg": self._prepare_dmrg_state() self._run_dmrg( @@ -2721,6 +2731,18 @@ def _to_state_backend(self, array): like = self._state_backend_like() if like is None: return np.asarray(array, dtype=complex) + # Avoid any Autoray conversion for an already-compatible payload. This + # is important for Symmray, whose backend intentionally does not expose + # a generic ``array`` constructor, and keeps the common matching-gate + # path allocation-free for every backend. + try: + if ( + ar.infer_backend(array) == ar.infer_backend(like) + and getattr(array, "dtype", None) == getattr(like, "dtype", None) + ): + return array + except (AttributeError, TypeError, ValueError): + pass dtype = getattr(like, "dtype", complex) if "complex" not in str(dtype): dtype = getattr( @@ -2737,6 +2759,118 @@ def _to_state_backend(self, array): arr = ar.do("array", array, like=like) return ar.do("astype", arr, dtype) + def to_backend(self, array): + """Return ``array`` on the backend currently owned by ``self.p``. + + Already-compatible arrays are returned by identity. This public helper + is intentionally state-derived so replacing the MPS with :meth:`set_p` + automatically changes the target backend without stale converter state. + """ + return self._to_state_backend(array) + + def _prepare_gate_stream_backend(self, gates, event_types): + """Prepare gate and sub-MPO payloads for the live MPS backend lazily. + + Gate streams are commonly authored as NumPy arrays even when the live + MPS uses Torch, JAX, CuPy, or another Autoray backend. The fast path in + :meth:`_to_state_backend` returns an already-compatible payload + unchanged, so matching streams incur no array copy. One representative + gate is used for the stream-level backend decision. Explicit sub-MPO + payloads are copied and converted with ``apply_to_arrays`` when needed, + preserving their tensor labels and operator bonds. + """ + if not gates: + return gates + like = self._state_backend_like() + like_backend = None + like_dtype = None + if like is not None: + try: + like_backend = ar.infer_backend(like) + except (AttributeError, TypeError, ValueError): + pass + like_dtype = getattr(like, "dtype", None) + + # Gate streams are expected to be backend-homogeneous. Inspect one + # ordinary gate, then apply that decision to the whole executable + # segment so matching streams are left entirely untouched. + gate_needs_conversion = like_backend is None + gate_backend = None + if like_backend is not None: + for candidate, event_type in zip(gates, event_types): + if event_type != "gate": + continue + try: + gate_backend = ar.infer_backend(candidate) + gate_needs_conversion = ( + gate_backend != like_backend + or getattr(candidate, "dtype", None) != like_dtype + ) + except (AttributeError, TypeError, ValueError): + gate_needs_conversion = True + break + + if ( + gate_needs_conversion + and gate_backend is not None + and gate_backend != like_backend + and not self._backend_mismatch_warned + ): + warnings.warn( + "MpsOptimizer converted a gate payload from backend " + f"{gate_backend!r} to the live MPS backend " + f"{like_backend!r}; provide matching gate payloads to " + "avoid this conversion.", + UserWarning, + stacklevel=3, + ) + self._backend_mismatch_warned = True + + prepared = [] + for gate, event_type in zip(gates, event_types): + if event_type == "gate": + if gate_needs_conversion: + gate = self.to_backend(gate) + elif event_type == "submpo" and like_backend is not None: + # ``apply_to_arrays`` changes only the raw tensor payloads, + # unlike rebuilding an MPO, which can lose custom labels or + # operator bonds. Keep the caller's stream immutable by + # applying it to a shallow network copy. + # Tensor-network payloads are expected to use one backend and + # dtype throughout, so inspect one representative tensor only. + tensor = next(iter(getattr(gate, "tensors", ())), None) + if tensor is None: + needs_conversion = False + else: + array = tensor.data + try: + needs_conversion = ( + ar.infer_backend(array) != like_backend + or getattr(array, "dtype", None) != like_dtype + ) + except (AttributeError, TypeError, ValueError): + needs_conversion = True + if needs_conversion: + if not self._backend_mismatch_warned: + warnings.warn( + "MpsOptimizer converted a sub-MPO payload to the " + f"live MPS backend {like_backend!r}; provide matching " + "sub-MPO payloads to avoid this conversion.", + UserWarning, + stacklevel=3, + ) + self._backend_mismatch_warned = True + gate = gate.copy() + apply_to_arrays = getattr(gate, "apply_to_arrays", None) + if not callable(apply_to_arrays): + raise TypeError( + "sub-MPO payloads must provide apply_to_arrays() " + "for backend conversion." + ) + apply_to_arrays(self.to_backend) + prepared.append(gate) + return prepared + def _pauli_operator(self, pauli, where): """Return the dense Pauli operator (numpy) for ``pauli`` on ``where``.""" chars = [c for c in str(pauli).upper() if not c.isspace()] diff --git a/src/pepsy/optimizers/tree/optimizer.py b/src/pepsy/optimizers/tree/optimizer.py index 14563c3..e646c18 100644 --- a/src/pepsy/optimizers/tree/optimizer.py +++ b/src/pepsy/optimizers/tree/optimizer.py @@ -293,17 +293,17 @@ class TreeOptimizer: applies. cutoff : float Relative singular-value cutoff for truncations. - mode : {"auto", "direct", "mpo"} - Implementation used for two-site gates. ``"direct"`` uses the - specialised gate-SVD/QR path threading algorithm. ``"mpo"`` first - factorises every two-site gate with Quimb into a two-tensor sub-MPO, - routes that MPO bond exactly through the tree, then compresses the - affected path once. ``"auto"`` (the default) selects direct threading - for every two-site gate; choose ``"mpo"`` explicitly when inspecting - or benchmarking the operator-TN formulation. The modes represent the - same update and can differ only through floating-point roundoff at an - exact bond dimension, or through the usual final ``chi``/``cutoff`` - truncation. + mode : {"auto", "direct", "mpo", "submpo"} + Implementation used for two-site gates and explicit operator streams. + ``"direct"`` uses the specialised gate-SVD/QR path. ``"mpo"`` first + factorises ordinary two-site gates with Quimb into a two-tensor + sub-MPO. ``"submpo"`` declares that the stream is already made of + explicit :meth:`submpo_event` entries (with ordinary one-site gates + allowed for singleton supports); it rejects ordinary multi-site gate + entries and replays the MPO payloads natively. + Explicit sub-MPO entries are also accepted in the other modes for + backward compatibility. ``"auto"`` (the default) selects direct + threading for ordinary two-site gates. structure : {"quality", "balanced", "adaptive"} Tree-structure strategy used when ``tree`` is not supplied. max_arity : int, None, or iterable of ints @@ -348,6 +348,11 @@ class TreeOptimizer: Whether to probe the full local singular spectrum before each truncating split/compression and record discarded-weight diagnostics. The extra spectrum probes are disabled by default. + track_infidelity : bool + Whether to compute the norm-based progress infidelity and include it + in progress-bar updates. This is enabled by default for compatibility + with direct TreeOptimizer use, but can be disabled for non-unitary + transfer-operator streams where norm changes are physical. max_intermediate_bond : int, optional Conservative preflight limit for the untruncated crossing-bond bound. When set, eager replay raises :class:`MemoryError` before tensor work if @@ -387,10 +392,12 @@ class TreeOptimizer: @staticmethod def _normalize_mode(mode): - """Validate and normalize the two-site gate implementation mode.""" + """Validate and normalize the gate or sub-MPO replay mode.""" mode = str(mode).strip().lower() - if mode not in {"auto", "direct", "mpo"}: - raise ValueError("mode must be 'auto', 'direct', or 'mpo'.") + if mode not in {"auto", "direct", "mpo", "submpo"}: + raise ValueError( + "mode must be 'auto', 'direct', 'mpo', or 'submpo'." + ) return mode @staticmethod @@ -411,7 +418,7 @@ def __init__(self, gates=None, n=None, *, chi=64, cutoff=1e-12, star_frac=0.75, layout_objective="path", layout_weight_mode="count", layout=None, tree=None, dtype=complex, threads=1, seed=None, run=True, tn=None, - state=None, track_truncation=False, + state=None, track_truncation=False, track_infidelity=True, max_intermediate_bond=None, max_operator_qubits=_DEFAULT_MAX_OPERATOR_QUBITS, max_subtree_nodes=_DEFAULT_MAX_SUBTREE_NODES, @@ -549,6 +556,7 @@ def __init__(self, gates=None, n=None, *, chi=64, cutoff=1e-12, raise ValueError("threads must be positive or None.") self.rng = np.random.default_rng(seed) self.track_truncation = bool(track_truncation) + self.track_infidelity = bool(track_infidelity) self.max_intermediate_bond = self._positive_limit( max_intermediate_bond, "max_intermediate_bond" ) @@ -780,6 +788,34 @@ def _validate_event_stream_for_run(self): f"current active labels {active!r}: {out_of_range!r}." ) + def _validate_mode_for_stream(self): + """Validate the explicit ``mode='submpo'`` stream declaration.""" + if self.mode != "submpo" or not self.G: + return + # ``output_replay='submpo'`` still represents singleton supports as + # ordinary one-site gates. They do not introduce a competing + # multi-site lowering path and are therefore valid in this mode. + ordinary = [] + for step, (where, event_type) in enumerate( + zip(self.where, self.event_types), start=1 + ): + if event_type != "gate": + continue + width = len(_normalize_where(where)) + if width > 1: + ordinary.append((step, width)) + if ordinary: + raise ValueError( + "mode='submpo' requires explicit sub-MPO events for " + "multi-site operations; ordinary multi-site gate event(s) " + f"found at step/width {ordinary!r}. " + "Use mode='direct' or mode='mpo' for dense gate streams." + ) + if "submpo" not in self.event_types: + raise ValueError( + "mode='submpo' requires at least one explicit sub-MPO event." + ) + # -- construction --------------------------------------------------------- @staticmethod @@ -847,6 +883,31 @@ def backend_info(self): """Return the common backend, dtype, and device of the live TTN.""" return self._state_backend_info(self.tn) + def _warn_backend_conversion(self, source_signature, target_signature): + """Warn once for one explicit source/target backend conversion.""" + warning_key = (source_signature, target_signature) + if ( + source_signature[0] != "builtins" + and warning_key not in self._backend_conversion_warnings + ): + self._backend_conversion_warnings.add(warning_key) + warnings.warn( + "TreeOptimizer is converting a gate/operator payload from " + f"backend/dtype/device {source_signature!r} to the TTN state " + f"{target_signature!r}. Provide backend-compatible gate " + "arrays to avoid this transfer or cast.", + UserWarning, + stacklevel=3, + ) + + @staticmethod + def _backend_converter(like): + """Build one converter for a stream targeting ``like``.""" + converter = infer_backend_converter_from_sample(like) + if converter is not None: + return converter + return lambda array: ar.do("array", array, like=like) + def _as_state_backend(self, array, *, warn=True): """Return an operator payload compatible with the live TTN backend. @@ -862,31 +923,96 @@ def _as_state_backend(self, array, *, warn=True): source_signature = _array_backend_signature(array) if source_signature == target_signature: return array - warning_key = (source_signature, target_signature) # Python sequences/scalars are ordinary convenience inputs rather than # a selected numerical backend. Materialize those silently; explicit # array backends/dtypes still receive the transfer/cast warning. - is_untyped_input = source_signature[0] == "builtins" - if ( - warn - and not is_untyped_input - and warning_key not in self._backend_conversion_warnings - ): - self._backend_conversion_warnings.add(warning_key) - warnings.warn( - "TreeOptimizer is converting a gate/operator payload from " - f"backend/dtype/device {source_signature!r} to the TTN state " - f"{target_signature!r}. Provide backend-compatible gate " - "arrays to avoid this transfer or cast.", - UserWarning, - stacklevel=3, - ) - converter = infer_backend_converter_from_sample(like) - if converter is not None: - return converter(array) + if warn: + self._warn_backend_conversion(source_signature, target_signature) if state_info["backend"] == "numpy": return ar.to_numpy(array) - return ar.do("array", array, like=like) + return self._backend_converter(like)(array) + + def _prepare_gate_stream_backend(self, payloads, event_types): + """Prepare one executable gate/sub-MPO stream for the live backend. + + Ordinary gates are expected to be backend-homogeneous. One + representative gate decides whether the whole gate stream needs + conversion; matching payloads are returned by identity. Sub-MPOs use + one representative tensor and ``apply_to_arrays`` on a copied network, + preserving the caller's labels and bonds. + """ + if not payloads: + return payloads + + like = self._state_like() + state_info = self.backend_info() + target_signature = ( + state_info["backend"], state_info["dtype"], state_info["device"] + ) + converter = None + prepared = list(payloads) + + gate_index = None + gate_signature = None + for index, (payload, event_type) in enumerate( + zip(payloads, event_types) + ): + if event_type != "gate": + continue + gate_index = index + try: + gate_signature = _array_backend_signature(payload) + except (AttributeError, TypeError, ValueError): + gate_signature = None + break + + if gate_index is not None: + gate_needs_conversion = gate_signature != target_signature + if gate_needs_conversion: + if gate_signature is not None: + self._warn_backend_conversion( + gate_signature, target_signature + ) + converter = self._backend_converter(like) + for index, event_type in enumerate(event_types): + if event_type == "gate": + prepared[index] = converter(payloads[index]) + + for index, (payload, event_type) in enumerate( + zip(payloads, event_types) + ): + if event_type != "submpo": + continue + tensor = next(iter(getattr(payload, "tensors", ())), None) + if tensor is None: + continue + try: + source_signature = _array_backend_signature(tensor.data) + except (AttributeError, TypeError, ValueError): + source_signature = None + if source_signature == target_signature: + continue + if source_signature is not None: + self._warn_backend_conversion( + source_signature, target_signature + ) + if converter is None: + converter = self._backend_converter(like) + copied = payload.copy() + apply_to_arrays = getattr(copied, "apply_to_arrays", None) + if callable(apply_to_arrays): + apply_to_arrays(converter) + else: + tensor_map = getattr(copied, "tensor_map", None) + if tensor_map is None: + # Opaque payloads are lowered later and will be coerced + # when their dense operator is materialized. + continue + for op_tensor in tensor_map.values(): + op_tensor.modify(data=converter(op_tensor.data)) + prepared[index] = copied + + return prepared def _coerce_tensor_network_backend(self, tn, *, warn=True): """Convert every tensor of an operator TN to the live state backend.""" @@ -1460,7 +1586,7 @@ def _apply_gate_impl(self, gate, where, *, renormalize=False): def run(self, gates=None, *, progbar=False, mode=None, non_unitary=False, normalize_every=False, normalize_final=False, - normalize_eps=1e-15, seed=None): + normalize_eps=1e-15, seed=None, track_infidelity=None): """Replay ``gates`` (or the construction stream) on the tree. Parameters @@ -1468,15 +1594,16 @@ def run(self, gates=None, *, progbar=False, mode=None, non_unitary=False, gates : bundled gate stream, optional Replacement stream to replay. If omitted, replay the queued stream. progbar : bool, default=False - Show a tqdm progress bar with the two-qubit gate count and a - norm-based truncation proxy. Both dense and native trees report - ``1 - (norm / reference_norm)**2``; the reference is established - at run start and reset after control/non-unitary events. - mode : {"auto", "direct", "mpo"} | {"tree", "ttn"} | None, default=None - Optional persistent two-site implementation selection, matching - ``MpsOptimizer.run(mode=...)``: a supplied value updates - :attr:`mode` before replay and remains active for future runs and - copies. ``"tree"``/``"ttn"`` are deprecated no-op compatibility + Show a tqdm progress bar with the two-qubit gate count. When + ``track_infidelity`` is enabled, also report the norm-based + truncation proxy ``1 - (norm / reference_norm)**2``; the reference + is established at run start and reset after control/non-unitary + events. + mode : {"auto", "direct", "mpo", "submpo"} | {"tree", "ttn"} | None, default=None + Optional persistent gate/sub-MPO replay selection: a supplied + value updates :attr:`mode` before replay and remains active for + future runs and copies. ``"submpo"`` validates an explicit MPO + stream; ``"tree"``/``"ttn"`` are deprecated no-op compatibility selectors for shared coefficient frontends. non_unitary : bool, default=False Mark the stream as non-unitary when using automatic normalization. @@ -1488,20 +1615,27 @@ def run(self, gates=None, *, progbar=False, mode=None, non_unitary=False, Zero-state threshold used by automatic normalization. seed : int | None, default=None Reseed measurement/reset sampling before replay. + track_infidelity : bool | None, default=None + Override :attr:`track_infidelity` for this replay. When disabled, + the progress bar omits the norm-based infidelity field and avoids + the per-event norm readout. Truncation-spectrum diagnostics remain + controlled independently by :attr:`track_truncation`. """ if mode is not None: requested_mode = str(mode).strip().lower() if requested_mode in {"tree", "ttn", "tree_tensor_network"}: warnings.warn( "run(mode='tree'/'ttn') is a deprecated no-op; use " - "mode='auto', 'direct', or 'mpo' to select a two-site " - "implementation.", + "mode='auto', 'direct', 'mpo', or 'submpo' to select " + "a gate/sub-MPO implementation.", DeprecationWarning, stacklevel=2, ) else: self.mode = self._normalize_mode(requested_mode) non_unitary = bool(non_unitary) + if track_infidelity is not None: + self.track_infidelity = bool(track_infidelity) if not non_unitary and normalize_every not in (False, None): raise ValueError("normalize_every requires non_unitary=True.") if not non_unitary and normalize_final: @@ -1515,6 +1649,13 @@ def run(self, gates=None, *, progbar=False, mode=None, non_unitary=False, if gates is not None: self.G, self.where, self.event_types = self._normalize_gate_queue(gates) self._validate_event_stream_for_run() + self._validate_mode_for_stream() + # Prepare the executable payloads without mutating the public queue. + # Matching streams retain their original objects; mismatched ordinary + # gates and sub-MPOs are converted to the live TTN backend once here. + payloads = self._prepare_gate_stream_backend( + self.G, self.event_types + ) pbar = None if progbar: from tqdm import tqdm # pylint: disable=import-outside-toplevel @@ -1534,15 +1675,16 @@ def run(self, gates=None, *, progbar=False, mode=None, non_unitary=False, control_count = 0 progress_reference_norm = ( self.norm() - if pbar is not None + if pbar is not None and self.track_infidelity else None ) try: for step, (payload, where, event_type) in enumerate(zip( - self.G, self.where, self.event_types + payloads, self.where, self.event_types ), start=1): - support = _normalize_where(where) + logical_support = _normalize_where(where) + support = logical_support if event_type == "gate": if len(support) == 1: one_qubit_count += 1 @@ -1552,27 +1694,20 @@ def run(self, gates=None, *, progbar=False, mode=None, non_unitary=False, multi_qubit_count += 1 self.apply_gate(payload, support) elif event_type == "submpo": - started = self._begin_update(event_type, support) - try: - support = self._validate_support(support) - self._check_operator_limits(support, dense=False) - with self._thread_ctx(): - applied = self._try_apply_native_submpo( - payload, support, max_bond=self.chi, - cutoff=self.cutoff, - ) - if applied is None: - self._check_operator_limits(support) - operator = _submpo_to_dense(payload, support) - self._apply_subtree_operator_impl( - operator, support - ) - except Exception: - if started: - self._abort_update() - raise - if started: - self._finish_update() + # Reuse the public sub-MPO implementation so stream + # replay gets the two-site factor fast path as well as + # the native multi-site MPO router. Passing both forms + # of support is important after a stable-label cap: + # ``support`` addresses compact TTN leaves, while + # ``logical_support`` addresses the MPO site tags. + support = self._validate_support(logical_support) + self._apply_submpo_resolved( + payload, + support, + logical_where=logical_support, + max_bond=self.chi, + cutoff=self.cutoff, + ) multi_qubit_count += 1 else: control_count += 1 @@ -1600,29 +1735,31 @@ def run(self, gates=None, *, progbar=False, mode=None, non_unitary=False, }) if pbar is not None: - # Use the same squared survival proxy for every backend. - # Control and explicitly non-unitary events change the - # physical norm for reasons unrelated to truncation, so - # reset the reference after those events. - state_norm = self.norm() - reset_progress = non_unitary or event_type != "gate" - if reset_progress: - truncation_infidelity = 0.0 - progress_reference_norm = state_norm - elif progress_reference_norm in (None, 0.0): - truncation_infidelity = 0.0 - else: - survival = state_norm / progress_reference_norm - truncation_infidelity = max( - 0.0, - 1.0 - survival * survival, - ) postfix = { "2q": two_qubit_count, - "infidelity": self._format_progress_infidelity( - truncation_infidelity - ), } + if self.track_infidelity: + # Use the same squared survival proxy for every + # backend. Control and explicitly non-unitary events + # change the physical norm for reasons unrelated to + # truncation, so reset the reference after those + # events. + state_norm = self.norm() + reset_progress = non_unitary or event_type != "gate" + if reset_progress: + truncation_infidelity = 0.0 + progress_reference_norm = state_norm + elif progress_reference_norm in (None, 0.0): + truncation_infidelity = 0.0 + else: + survival = state_norm / progress_reference_norm + truncation_infidelity = max( + 0.0, + 1.0 - survival * survival, + ) + postfix["infidelity"] = self._format_progress_infidelity( + truncation_infidelity + ) if multi_qubit_count: postfix["kq"] = multi_qubit_count if control_count: @@ -1835,8 +1972,14 @@ def _apply_2q_impl(self, gate, qa, qb, *, max_bond=None, cutoff=None): make the equivalent two-tensor MPO. Both immediately enter the same two-factor attach/QR-thread/compress kernel. ``'auto'`` selects direct factorization for every backend; use ``'mpo'`` explicitly to select - Quimb's operator-TN factorization. + Quimb's operator-TN factorization. ``'submpo'`` is reserved for + explicit sub-MPO stream events and cannot be used with a dense gate. """ + if self.mode == "submpo": + raise ValueError( + "mode='submpo' accepts explicit sub-MPO stream events, not " + "ordinary dense gates; use mode='direct' or mode='mpo'." + ) gate = self._as_state_backend(gate) if self.mode in {"auto", "direct"}: return self._apply_2q_path_thread_impl( @@ -2487,9 +2630,11 @@ def apply_submpo(self, submpo, where, *, max_bond=None, cutoff=None): """Apply an explicit MPO on ``where`` using the native tree path. This is the backend-neutral coefficient-state entry point used by - stabilizer and ordinary operator-sum frontends. MPOs exposing Quimb's - site interface stay structured; opaque MPO-like payloads fall back to - dense :meth:`apply_subtree_operator` lowering. + stabilizer and ordinary operator-sum frontends. Two-site MPOs reuse + the factorized gate path; larger MPOs exposing Quimb's site interface + stay structured and are QR-routed through the Steiner subtree before + one compression sweep. Opaque MPO-like payloads fall back to dense + :meth:`apply_subtree_operator` lowering. """ self._invalidate_state_norm_cache() logical_where = _normalize_where(where) @@ -3907,6 +4052,7 @@ def copy(self): layout_objective=self.layout_objective, layout_weight_mode=self.layout_weight_mode, track_truncation=self.track_truncation, + track_infidelity=self.track_infidelity, max_intermediate_bond=self.max_intermediate_bond, max_operator_qubits=self.max_operator_qubits, max_subtree_nodes=self.max_subtree_nodes, @@ -3926,6 +4072,9 @@ def copy(self): other.infidelity_samples = deepcopy(self.infidelity_samples) other.normalizations = deepcopy(self.normalizations) other.projection_diagnostics = deepcopy(self.projection_diagnostics) + other._backend_conversion_warnings = set( + self._backend_conversion_warnings + ) other._logical_qubits = list(self._logical_qubits) other._logical_positions = dict(self._logical_positions) other._truncation_survival = self._truncation_survival diff --git a/src/pepsy/tensors/__init__.py b/src/pepsy/tensors/__init__.py index dd1e0c1..783b853 100644 --- a/src/pepsy/tensors/__init__.py +++ b/src/pepsy/tensors/__init__.py @@ -84,6 +84,7 @@ def _register(module, *names): ) _register( ".contractions", + "build_contraction", "build_compressed_optimizer", "build_optimizer", "contract_hypercompressed_tn", @@ -98,6 +99,7 @@ def _register(module, *names): "backend_jax", "backend_numpy", "backend_torch", + "build_backend", "get_default_array_backend", "get_default_grad_backend", "register_torch_linalg", diff --git a/src/pepsy/tensors/contractions.py b/src/pepsy/tensors/contractions.py index 2d5cee7..d9c1c64 100644 --- a/src/pepsy/tensors/contractions.py +++ b/src/pepsy/tensors/contractions.py @@ -9,7 +9,7 @@ import quimb.tensor as qtn __all__ = [ - "build_optimizer", "build_compressed_optimizer", + "build_contraction", "build_optimizer", "build_compressed_optimizer", "contract_hypercompressed_tn", "contract_hypercompressed_tn_batch", "tn_norm", ] @@ -109,6 +109,19 @@ def build_optimizer( return ctg.ReusableHyperOptimizer(**kwargs) +def build_contraction(*args, **kwargs): + """Build a reusable contraction optimizer. + + This is the short, backend-neutral alias for :func:`build_optimizer`. + Numerical contractions use the array backend of the tensors supplied to + Quimb; pair it with ``py.build_backend()`` and ``to_backend=`` to run the + calculation on Torch CPU while keeping the existing ``build_optimizer`` + name available. + """ + + return build_optimizer(*args, **kwargs) + + def build_compressed_optimizer( progbar=True, chi=4, diff --git a/src/pepsy/tensors/core.py b/src/pepsy/tensors/core.py index 61b8a59..6369953 100644 --- a/src/pepsy/tensors/core.py +++ b/src/pepsy/tensors/core.py @@ -11,6 +11,7 @@ from . import observables as _observables from .contractions import ( build_compressed_optimizer as _build_compressed_optimizer, + build_contraction as _build_contraction, build_optimizer as _build_optimizer, contract_hypercompressed_tn as _contract_hypercompressed_tn, contract_hypercompressed_tn_batch, @@ -43,6 +44,7 @@ backend_jax, backend_numpy, backend_torch, + build_backend, get_default_array_backend, get_default_grad_backend, reg_complex_qr_torch, @@ -77,6 +79,17 @@ def build_optimizer(*args, **kwargs): _contractions._ensure_cotengrust = original_ensure +def build_contraction(*args, **kwargs): + """Compatibility wrapper for :func:`pepsy.tensors.build_contraction`.""" + + original_ensure = _contractions._ensure_cotengrust + _contractions._ensure_cotengrust = _ensure_cotengrust + try: + return _build_contraction(*args, **kwargs) + finally: + _contractions._ensure_cotengrust = original_ensure + + def build_compressed_optimizer(*args, **kwargs): original_ensure = _contractions._ensure_cotengrust _contractions._ensure_cotengrust = _ensure_cotengrust @@ -104,13 +117,13 @@ def contract_hypercompressed_tn(*args, **kwargs): _contractions.build_compressed_optimizer = original_build_compressed_optimizer __all__ = [ - "OneDMap", "backend_torch", "backend_numpy", "backend_cupy", "backend_jax", + "OneDMap", "build_backend", "backend_torch", "backend_numpy", "backend_cupy", "backend_jax", "register_torch_linalg", "reg_rel_svd_torch", "reg_real_svd_torch", "reg_complex_svd_torch", "reg_real_qr_torch", "reg_complex_qr_torch", "reg_rel_svd_jax", "reg_real_svd_jax", "reg_complex_svd_jax", "reg_stop_gradient_torch", "stop_grad", "set_default_array_backend", "get_default_array_backend", "set_default_grad_backend", "get_default_grad_backend", - "reset_default_backends", "build_optimizer", "build_compressed_optimizer", + "reset_default_backends", "build_contraction", "build_optimizer", "build_compressed_optimizer", "contract_hypercompressed_tn", "contract_hypercompressed_tn_batch", "tn_fidelity", "tn_norm", "measure_obs", "tns_align", "expec_mpo", "id_to_mpo", "id_to_pepo", "ps_to_peps", "ps_to_3dpeps", "ps_to_mps", "ps_to_ttn", "ps_to_pepo", "ps_to_mpo", diff --git a/tests/test_backends.py b/tests/test_backends.py index a6c6b89..095588c 100644 --- a/tests/test_backends.py +++ b/tests/test_backends.py @@ -29,6 +29,24 @@ def __array__(self, *_args, **_kwargs): assert pepsy.to_float(BackendScalar()) == pytest.approx(1.25) +def test_register_torch_svd_for_autoray(): + """The opt-in Torch SVD registration enables Pepsy's robust autoray SVD.""" + torch = pytest.importorskip("torch") + import autoray as ar + + pepsy.reg_rel_svd_torch() + svd_fn = ar.get_lib_fn("torch", "linalg.svd") + assert getattr(svd_fn, "__self__", None).__module__ == ( + "pepsy.backends.linalg_torch" + ) + + matrix = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float64) + u, s, vh = ar.do("linalg.svd", matrix) + assert u.shape == (2, 2) + assert s.shape == (2,) + assert vh.shape == (2, 2) + + def test_to_float_rejects_non_scalar_backend_array_before_numpy_coercion(): class BackendVector: shape = (2,) diff --git a/tests/test_optimize_mps.py b/tests/test_optimize_mps.py index b0eba62..d198f68 100644 --- a/tests/test_optimize_mps.py +++ b/tests/test_optimize_mps.py @@ -185,6 +185,64 @@ def test_mps_optimizer_simple_update_routes_torch_u1u1_long_range_gate(): assert len(optimizer.gauges) == out.L - 1 +@pytest.mark.parametrize( + "mode", ["dmrg", "mpo", "svd", "swap", "perm", "mix", "su", "exact"] +) +def test_mps_optimizer_casts_numpy_gate_stream_to_torch_state_backend(mode): + """Backend conversion should happen only when a gate payload mismatches p.""" + torch = pytest.importorskip("torch") + + state = qtn.MPS_computational_state("00", dtype="complex128") + state.apply_to_arrays(py.backend_torch(dtype=torch.complex128, device="cpu")) + gate = qu.CNOT() # ordinary NumPy gate stream + optimizer = py.MpsOptimizer( + state, + gates=[(gate, (0, 1))], + chi=2, + mode=mode, + inplace=True, + ) + + with pytest.warns(UserWarning, match="converted a gate payload"): + optimizer.run(progbar=False, n_iter=2) + + assert all(isinstance(tensor.data, torch.Tensor) for tensor in optimizer.p.tensors) + dense = np.asarray(py.MpsOptimizer._real_float(optimizer.p.norm())) + assert dense == pytest.approx(1.0) + + matching_gate = torch.as_tensor(np.array(gate, copy=True), dtype=torch.complex128) + assert optimizer.to_backend(matching_gate) is matching_gate + + +def test_mps_optimizer_casts_submpo_stream_arrays_to_torch_state_backend(): + """Sub-MPO conversion preserves the network structure and input stream.""" + torch = pytest.importorskip("torch") + + state = qtn.MPS_computational_state("00", dtype="complex128") + state.apply_to_arrays(py.backend_torch(dtype=torch.complex128, device="cpu")) + submpo = _two_branch_flip_submpo(L=2, sites=(0, 1), targets=(0, 1)) + original_inds = tuple(submpo.outer_inds()) + optimizer = py.MpsOptimizer( + state, + gates=[py.MpsOptimizer.submpo_event(submpo, (0, 1))], + chi=2, + mode="mpo", + inplace=True, + ) + + with pytest.warns(UserWarning, match="converted a sub-MPO payload"): + optimizer.run(progbar=False, n_iter=2) + + # ``run`` prepares a local executable segment, leaving the public queue + # unchanged. Re-run the same helper to inspect the converted payload. + converted = optimizer._prepare_gate_stream_backend([submpo], ["submpo"])[0] + assert converted is not submpo + assert tuple(converted.outer_inds()) == original_inds + assert all(isinstance(tensor.data, torch.Tensor) for tensor in converted.tensors) + assert all(isinstance(tensor.data, np.ndarray) for tensor in submpo.tensors) + assert all(isinstance(tensor.data, torch.Tensor) for tensor in optimizer.p.tensors) + + def test_mps_optimizer_accepts_perm_mode(): """Perm mode should expose an identity logical-to-physical ordering initially.""" p0 = qtn.MPS_computational_state("0000", dtype="complex128") diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index 0ef4dca..da09f6a 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -1801,6 +1801,19 @@ def fail_to_dense(): assert report["events"][0]["crossing_edges"] +def test_tree_submpo_mode_declares_and_validates_mpo_streams(): + """The explicit sub-MPO mode accepts MPO events and rejects dense gates.""" + mpo = _two_branch_flip_submpo(L=4, sites=(0, 3), targets=(0, 3)) + opt = TreeOptimizer(None, n=4, chi=8, mode="submpo", run=False) + opt.run([TreeOptimizer.submpo_event(mpo, (0, 3))]) + assert opt.mode == "submpo" + + dense = np.asarray(mpo.to_dense()) + ordinary = TreeOptimizer(None, n=4, chi=8, mode="submpo", run=False) + with pytest.raises(ValueError, match="requires explicit sub-MPO"): + ordinary.run([(dense, (0, 3))]) + + def test_tree_estimate_bonds_includes_submpo_operator_schmidt_rank(): """Sub-MPO markers participate in the same conservative bond estimate.""" mpo = _two_branch_flip_submpo(L=4, sites=(0, 3), targets=(0, 3)) @@ -1963,6 +1976,51 @@ def test_tree_warns_once_when_a_gate_does_not_match_the_state_backend(): assert opt.backend_info()["backend"] == "torch" +def test_tree_gate_stream_backend_preparation_is_stream_level(): + """One representative gate decides conversion for the whole stream.""" + torch = pytest.importorskip("torch") + to_backend = pepsy.backend_torch(device="cpu", dtype=torch.complex128) + plan = TreePlan.from_order(range(2), structure="balanced") + state = TreeTensorNetwork.from_plan(plan) + for tensor in state.tensor_map.values(): + tensor.modify(data=to_backend(tensor.data)) + gates = [ + np.eye(2, dtype=complex), + np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex), + ] + opt = TreeOptimizer(None, state=state, tree=plan, run=False) + + with pytest.warns(UserWarning, match="converting a gate/operator payload"): + prepared = opt._prepare_gate_stream_backend(gates, ["gate", "gate"]) + + assert all(torch.is_tensor(gate) for gate in prepared) + assert all(isinstance(gate, np.ndarray) for gate in gates) + + matching = [to_backend(gate) for gate in gates] + untouched = opt._prepare_gate_stream_backend(matching, ["gate", "gate"]) + assert untouched[0] is matching[0] + assert untouched[1] is matching[1] + + +def test_tree_submpo_stream_backend_preparation_preserves_input(): + """A mismatched stream sub-MPO is copied and converted by its arrays.""" + torch = pytest.importorskip("torch") + to_backend = pepsy.backend_torch(device="cpu", dtype=torch.complex128) + plan = TreePlan.from_order(range(2), structure="balanced") + state = TreeTensorNetwork.from_plan(plan) + for tensor in state.tensor_map.values(): + tensor.modify(data=to_backend(tensor.data)) + submpo = _two_branch_flip_submpo(L=2, sites=(0, 1), targets=(0, 1)) + opt = TreeOptimizer(None, state=state, tree=plan, run=False) + + with pytest.warns(UserWarning, match="converting a gate/operator payload"): + prepared = opt._prepare_gate_stream_backend([submpo], ["submpo"])[0] + + assert prepared is not submpo + assert all(torch.is_tensor(tensor.data) for tensor in prepared.tensors) + assert all(isinstance(tensor.data, np.ndarray) for tensor in submpo.tensors) + + def test_tree_rejects_a_mixed_backend_initial_state(): """A TTN must use one backend, dtype, and device across all tensors.""" torch = pytest.importorskip("torch") @@ -2777,6 +2835,28 @@ def test_tree_stable_labels_route_submpo_by_payload_sites(monkeypatch): assert opt.norm() == pytest.approx(1.0) +def test_tree_stream_stable_labels_route_submpo_natively(monkeypatch): + """Stream replay preserves native MPO routing after a stable-label cap.""" + x = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex) + mpo = qtn.MatrixProductOperator.from_dense( + np.kron(x, x), dims=(2, 2), sites=(2, 3), L=4 + ) + events = [ + TreeOptimizer.cap_event(1, [1.0, 0.0], compact_labels=False), + TreeOptimizer.submpo_event(mpo, (2, 3)), + ] + opt = TreeOptimizer(None, n=4, chi=8, run=False) + + monkeypatch.setattr( + mpo, + "to_dense", + lambda: (_ for _ in ()).throw(AssertionError("dense MPO fallback")), + ) + opt.run(events) + assert opt.qubits == [0, 2, 3] + assert opt.norm() == pytest.approx(1.0) + + def test_tree_estimate_bonds_tracks_compact_plan_after_cap(): """Bond preflight follows the live logical mapping across a cap event.""" cnot = np.array( From 080144ee33d9c20afc51b630a0f941eabf911224 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Tue, 28 Jul 2026 12:45:27 -0700 Subject: [PATCH 08/70] Add root physical sites and tree layout search controls --- .github/skills/tree-optimizer/SKILL.md | 35 +- docs/api/optimizers/tree.md | 98 ++++-- src/pepsy/optimizers/tree/layout.py | 467 +++++++++++++++++++------ src/pepsy/optimizers/tree/optimizer.py | 198 ++++++----- src/pepsy/optimizers/tree/ttn.py | 245 ++++++++----- tests/test_optimize_tree.py | 135 +++++++ 6 files changed, 877 insertions(+), 301 deletions(-) diff --git a/.github/skills/tree-optimizer/SKILL.md b/.github/skills/tree-optimizer/SKILL.md index a0c000c..5e9f71c 100644 --- a/.github/skills/tree-optimizer/SKILL.md +++ b/.github/skills/tree-optimizer/SKILL.md @@ -24,7 +24,8 @@ The rooted tree-tensor-network circuit simulator of *Simulating quantum circuits using tree tensor networks* (Seitz, Medina, Cruz, Huang, Mendl; Quantum 7, 964, 2023; arXiv:2206.01000). The state is a rooted TTN (internal nodes of **any arity**; binary is the default, see *Non-binary trees* below) -whose leaves carry the physical qubit indices; a bundled gate stream +whose leaves carry physical qubit indices. An optional ``root_qubit`` is instead +carried by the top tensor; all other physical sites remain leaves. A bundled gate stream `[(gate, where), ...]` is replayed. `where` is an `int` (1q) or a pair of `int` (2q); supports with `len(where) >= 3` route through `apply_subtree_operator` (see *Multi-qubit / sub-MPO application*). @@ -91,10 +92,12 @@ naming, so all inherited quimb methods (`canonize_around`, `canonize_between`, `if isinstance(ts, TensorNetwork): super().__init__(ts, **o); return` lets the base copy the extra props without the fresh-construction defaults clobbering `_plan`. -- Each leaf tensor carries **both** the structural node tag `N{nid}` and the - quimb site tag `I{q}` plus physical index `k{q}`; internal nodes carry only - `N{nid}`. So quimb sees the leaves as the `nsites` sites and internal nodes as - ancillary bond carriers -- +- Each physical-site tensor carries **both** the structural node tag `N{nid}` + and the quimb site tag `I{q}` plus physical index `k{q}`. These tensors are + structural leaves by default. When + `plan.root_qubit` is set, the root carries that site tag/index too; other + internal nodes carry only `N{nid}`. So quimb sees the leaf sites and optional + root site as `nsites`, with remaining internal nodes as ancillary bond carriers -- `ttn.local_expectation(G, where=[q], max_bond=None, optimize="auto")` uses the tree's canonical contraction for dense states and an exact complete doubled-tree contraction for native fermionic states. @@ -103,8 +106,8 @@ naming, so all inherited quimb methods (`canonize_around`, `canonize_between`, - Builders: `from_plan(plan)` (product `|0...0>`), `from_order(order, structure=...)` (plan + product in one call), `rand(plan, D=, seed=, canonicalize=True)` (random state, canonicalised around the root). -- `show()` prints a top-down ASCII tree (root on top, qubit leaves `◆ q{q}` at - the bottom, internal `●`, each branch annotated with its bond dim); +- `show()` prints a top-down ASCII tree (root on top, structural leaves at the + bottom, physical nodes labelled `q{q}`, each branch annotated with its bond dim); `ascii_tree()` returns that string. `TreeOptimizer.show()` delegates to it. - `TreeOptimizer.tn` **is** a `TreeTensorNetwork`; the optimizer delegates `_phys->tn.site_ind`, `_tag->tn.node_tag`, `_tid->tn.node_tid`, @@ -117,13 +120,15 @@ naming, so all inherited quimb methods (`canonize_around`, `canonize_between`, - Node ids are ints from `TreePlan`. Tensor tag = `N{nid}` (`TreeTensorNetwork. node_tag`, via optimizer `_tag`). - Physical index of qubit `q` = `k{q}` (`TreeTensorNetwork.site_ind`, via - optimizer `_phys`) -- ket-leg convention. Leaves also carry site tag `I{q}`. + optimizer `_phys`) -- ket-leg convention. Physical nodes also carry site tag + `I{q}`. Resolve their geometry with `plan.node_of_qubit[q]`; use + `leaf_of_qubit` only when a true structural leaf is required. - Newly created virtual bonds between adjacent nodes `u,v` use `_tb{lo}_{hi}` with `lo= {"k4"} +``` + +`root_qubit` is first-class rather than an unregistered outer leg: +`to_dense()` retains it in normal qubit order, `cap(root_qubit, vec)` contracts +only that physical leg, and direct gates, dense subtree operators, and +structured sub-MPOs may include it in their support. `TreeLayoutFinder` keeps +the site fixed at the root while its path, Steiner, congestion, greedy, and +Nevergrad objectives permute only the remaining leaf sites. + +Gates are absorbed into the tree: + +- **single-qubit gates** are contracted into their site tensor with no bond growth; a unitary one-qubit gate preserves the tree canonical form regardless of where the orthogonality centre sits; -- **two-qubit gates** on leaves `a` and `b` are split by SVD into two factors - joined by a virtual bond; the factors are absorbed into the two leaves and the +- **two-qubit gates** on sites `a` and `b` are split by SVD into two factors + joined by a virtual bond; the factors are absorbed into the two site nodes and the bond is *threaded exactly* (lossless economical QR) along the tree path from `a` to `b`. Only once **both** factors are in place is a single canonical compression sweep run back along the path, truncating every touched bond to @@ -37,7 +67,7 @@ into the tree: compression when the payload exposes Quimb's MPO site interface, so `to_dense()` is not required; opaque MPO-like payloads fall back to the dense recursive subtree-operator path. - `cap` contracts and removes one tree leaf, compacts the remaining qubit + `cap` contracts and removes one physical site, compacts the remaining qubit labels above it, and keeps the live tree canonical. The orthogonality centre is a single node id tracked on the @@ -93,7 +123,7 @@ exactly as the single centre tensor does for a one-node region. Disconnected `nodes` raise unless `span=True` auto-expands to the minimal connected subtree that spans them (`subtree_span`). `canonize_around_qubits_(qubits)` is the qubit-level entry point: it canonicalises around the minimal subtree spanning -those qubits' leaves, so the reduced state on a set of qubits is captured by one +those qubits' physical nodes, so the reduced state on a set of qubits is captured by one subtree. `is_subtree_canonical_form(nodes)` verifies the outside-is-isometric property directly; `is_canonical_form` is its one-node case. `TreeOptimizer` mirrors this too: `canonical_region`, `canonize_subtree(nodes, span=...)`, @@ -109,7 +139,8 @@ generalisation of the two-qubit gate: a `k`-qubit gate, a multi-site analogue of a sub-MPO applied over the covering range and then compressed (cf. Quimb's `MatrixProductState.gate_with_submpo`, which exists for the 1D chain only). The dense operator is first factorized into an exact tree-MPO on the -**minimal connected subtree** (Steiner subtree) spanning the target leaves. +**minimal connected subtree** (Steiner subtree) spanning the target physical +nodes. Application then proceeds recursively from subtree leaves to a hub: each local state/operator message is losslessly QR-split on one edge and absorbed by its parent, carrying every still-open operator virtual leg. No dense state tensor @@ -183,8 +214,9 @@ ordinary one-/two-/multi-qubit gates, structured sub-MPOs, Pauli expectation and projection, measurement, reset, measure-reset, cap, normalization, copying, canonicalization, layout construction, dense readout, and truncation diagnostics for dense two-level qubit TTNs. A cap's `absorb` argument is -accepted for stream compatibility, but a tree always absorbs into the leaf's -unique parent. `cap(q, vec)` compacts labels by default; use +accepted for stream compatibility. A leaf site absorbs into its unique parent; +a root site is contracted directly without changing the tree edges. +`cap(q, vec)` compacts labels by default; use `stable_labels=True` (or `compact_labels=False`) to preserve caller-facing logical IDs across the cap while the internal TTN stays compact. `TreeOptimizer.qubits`, `logical_order`, `position`, and `logical_site` expose @@ -230,7 +262,8 @@ class adds the naming and geometry glue on top of a node tag `node_tag_id.format(nid)` (default `"N{}"`); - leaf tensors additionally carry the Quimb site tag `site_tag_id.format(q)` (default `"I{}"`) and physical index `site_ind_id.format(q)` (default `"k{}"`) - for qubit `q`, so Quimb treats the leaves as the sites; + for qubit `q`; when `plan.root_qubit` is set, the root tensor carries that + qubit's site tag and physical index as well; - adjacent nodes share one live virtual bond. Newly constructed edges use the deterministic `_tb{lo}_{hi}` name, but Quimb may replace it with a UUID during threading or canonicalisation; `TreeTensorNetwork.bond(a, b)` resolves the @@ -247,7 +280,7 @@ geometry queries to it. `TreeTensorNetwork.local_expectation(op, where, max_bond=None)` has two backend-specific exact paths. Dense/nonfermionic TTNs move the centre to the -target leaf/subtree, cancel the ordinary isometric exterior, and contract only +target physical node/subtree, cancel the ordinary isometric exterior, and contract only the minimal Steiner subtree. Native fermionic TTNs insert the Symmray operator without densifying it and contract the complete doubled tree, preserving every graded boundary phase. For native fermionic states, `max_bond` is accepted for @@ -291,9 +324,9 @@ tree with the requested charge-sector bond dimension. These constructors keep the Symmray arrays native; they do not materialize dense tensor data. `TreeTensorNetwork.show()` prints a top-down ASCII drawing of the tree -- the -tree analogue of a quimb MPS `show()` -- with the root at the top and the qubit -leaves at the bottom, internal nodes marked `●`, leaves `◆` labelled by their -qubit, and every branch annotated with its current virtual bond dimension +tree analogue of a quimb MPS `show()` -- with the root at the top, structural +leaves at the bottom, physical sites labelled by qubit, and every branch +annotated with its current virtual bond dimension (`ascii_tree()` returns the same drawing as a string). `TreeOptimizer.show()` delegates to it. @@ -305,10 +338,12 @@ recursive spectral (Fiedler) partition, keeping the recursion as the rooted tree (`structure="quality"`). This reuses the interaction-graph and spectral machinery of `pepsy.optimizers.mps.layout`; where the MPS finder flattens the recursion into a 1D order, the tree finder keeps the tree. Strongly coupled -qubits become nearby leaves, minimising the tree-path length that two-qubit -gates thread across. `structure="balanced"` splits the qubit index order in -half at each level. `TreeLayoutFinder.score(plan)` returns the total -interaction-weighted tree-path length that the structure minimises. +qubits become nearby physical nodes, minimising the tree-path length that +two-qubit gates thread across. With `root_qubit=q`, that physical node stays +fixed at the top while the finder searches over the remaining leaf sites. +`structure="balanced"` splits the leaf-qubit order in half at each level. +`TreeLayoutFinder.score(plan)` returns the total interaction-weighted tree-path +length that the structure minimises. For circuits with gates of different operator-Schmidt ranks, use `TreeLayoutFinder(..., objective="congestion")` or @@ -459,6 +494,27 @@ choice = finder.recommend_layered( ) ``` +The same fixed-plan quality controls can be supplied directly to `run()`, +which gives finder-based frontends the same define-then-search shape as the MPS +layout API: + +```python +tree_plan = finder.run( + refine="greedy", + refine_budget=64, + search="nevergrad", + search_budget=128, + seed=0, + nevergrad_optimizer="OnePlusOne", + progbar=True, +) +``` + +Omitted `run()` options inherit the finder configuration, preserving the +zero-argument API. Structure, arity candidates, objective, and event weighting +remain finder-construction options because they define the Tree search space +and scoring model rather than one refinement pass. + Nevergrad evaluates every candidate plan, so reserve it for offline circuit studies rather than routine short simulations. Candidate records expose their initial/final leaf order and the greedy/Nevergrad diagnostics under @@ -553,7 +609,7 @@ backend. The dominant lever for accuracy at fixed `chi` is the tree structure, so the finder and optimizer expose diagnostics to choose it: -- `TreeLayoutFinder.report(plan=None)` summarises the leaf-to-leaf geodesic +- `TreeLayoutFinder.report(plan=None)` summarises the physical-node geodesic lengths over the interaction graph (`score`, `max_path`, `mean_path`, `weighted_mean_path`) and compares against a balanced index tree (`balanced_score`, `score_ratio_vs_balanced`). It also reports diff --git a/src/pepsy/optimizers/tree/layout.py b/src/pepsy/optimizers/tree/layout.py index 8183b3f..a2e4981 100644 --- a/src/pepsy/optimizers/tree/layout.py +++ b/src/pepsy/optimizers/tree/layout.py @@ -312,24 +312,38 @@ def _chi_cut_fields(plan, chi): class TreePlan: - """A rooted tree over ``n`` qubit leaves (any internal-node arity). - - Nodes are integer ids. Leaves map one-to-one to qubits; internal nodes have - one or more children. A strictly-binary tree (every internal node with two - children) is the common default, but the structure supports arbitrary arity - so a level can branch into as many subtrees as the gate stream suggests. - The plan is a pure structure description: it carries no tensor data and is - consumed by :class:`~pepsy.optimizers.tree.TreeOptimizer` to build the tree - tensor network. + """A rooted tree over ``n`` qubits (any internal-node arity). + + Nodes are integer ids. Leaves map one-to-one to qubits. Optionally, one + additional qubit can be carried by the structural root via ``root_qubit``; + this gives a binary top tensor two child bonds plus one open physical leg. + Other internal nodes carry no physical qubit. A strictly-binary tree (every + internal node with two children) is the common default, but the structure + supports arbitrary arity so a level can branch into as many subtrees as the + gate stream suggests. The plan is a pure structure description: it carries + no tensor data and is consumed by + :class:`~pepsy.optimizers.tree.TreeOptimizer` to build the tree tensor + network. """ - def __init__(self, root, children, parent, qubit_of_leaf): + def __init__( + self, root, children, parent, qubit_of_leaf, *, root_qubit=None + ): self.root = root self.children = dict(children) self.parent = dict(parent) self.qubit_of_leaf = dict(qubit_of_leaf) self.leaf_of_qubit = {q: nid for nid, q in self.qubit_of_leaf.items()} - self.n = len(self.qubit_of_leaf) + self.root_qubit = ( + None if root_qubit is None else int(root_qubit) + ) + self.qubit_of_node = dict(self.qubit_of_leaf) + if self.root_qubit is not None: + self.qubit_of_node[self.root] = self.root_qubit + self.node_of_qubit = { + q: nid for nid, q in self.qubit_of_node.items() + } + self.n = len(self.node_of_qubit) self._path_cache = {} # -- construction --------------------------------------------------------- @@ -337,13 +351,15 @@ def __init__(self, root, children, parent, qubit_of_leaf): @classmethod def from_order(cls, order, *, weights=None, structure="quality", max_arity=2, community_frac=0.35, star_frac=0.75, - dense_max=512): + dense_max=512, root_qubit=None): """Build a rooted tree by recursive partition of ``order``. Parameters ---------- order : sequence of int - The qubit labels to place as leaves. + The qubit labels to place as leaves. When ``root_qubit`` is given, + ``order`` contains every other qubit and the combined labels must + still be ``0..n-1``. weights : mapping, optional Unordered ``(qi, qj) -> weight`` interaction weights. Used to spectrally reorder each recursion level (``structure="quality"``) @@ -372,17 +388,26 @@ def from_order(cls, order, *, weights=None, structure="quality", (all pairwise geodesics length two) instead of being bisected. dense_max : int Maximum subsystem size for dense spectral reordering. + root_qubit : int, optional + Qubit label carried by the top tensor rather than a leaf. """ order = list(order) - if not order: + if not order and root_qubit is None: raise ValueError("order must contain at least one qubit.") try: order = [int(q) for q in order] except (TypeError, ValueError) as exc: raise ValueError("order must contain integer qubit labels.") from exc - if sorted(order) != list(range(len(order))): + if root_qubit is not None: + try: + root_qubit = int(root_qubit) + except (TypeError, ValueError) as exc: + raise ValueError("root_qubit must be an integer or None.") from exc + all_qubits = order + ([] if root_qubit is None else [root_qubit]) + if sorted(all_qubits) != list(range(len(all_qubits))): raise ValueError( - "order must be a permutation of qubit labels 0..n-1." + "leaf order plus root_qubit must be a permutation of " + "qubit labels 0..n-1." ) if structure not in {"quality", "balanced", "adaptive"}: raise ValueError( @@ -533,11 +558,23 @@ def build(qs): child_ids = [build(g) for g in groups] return make_internal(child_ids) - root = build(order) - return cls(root, children, parent, qubit_of_leaf) + if order: + root = build(order) + else: + root = new_node() + children[root] = () + return cls( + root, + children, + parent, + qubit_of_leaf, + root_qubit=root_qubit, + ) @classmethod - def from_children(cls, children, qubit_of_leaf, *, root=None): + def from_children( + cls, children, qubit_of_leaf, *, root=None, root_qubit=None + ): """Build and validate a :class:`TreePlan` from an explicit tree. This is the general entry point for arbitrary (non-binary) trees: a @@ -554,6 +591,8 @@ def from_children(cls, children, qubit_of_leaf, *, root=None): root : int, optional The root node id. Inferred as the unique parent-less node when omitted. + root_qubit : int, optional + Qubit label carried by ``root`` rather than by a leaf. """ children = {int(k): tuple(int(c) for c in v) for k, v in children.items()} @@ -581,6 +620,11 @@ def from_children(cls, children, qubit_of_leaf, *, root=None): root = int(root) if root not in children or root in parent: raise ValueError(f"invalid root {root}") + if root_qubit is not None: + try: + root_qubit = int(root_qubit) + except (TypeError, ValueError) as exc: + raise ValueError("root_qubit must be an integer or None.") from exc leaves = set() for nid, ch in children.items(): @@ -591,15 +635,32 @@ def from_children(cls, children, qubit_of_leaf, *, root=None): ) else: leaves.add(nid) - if nid not in qubit_of_leaf: + if ( + nid not in qubit_of_leaf + and not (nid == root and root_qubit is not None) + ): raise ValueError(f"leaf node {nid} is missing a qubit") - if set(qubit_of_leaf) != leaves: + expected_leaf_nodes = ( + leaves - {root} + if root_qubit is not None and not children[root] + else leaves + ) + if set(qubit_of_leaf) != expected_leaf_nodes: raise ValueError( "qubit_of_leaf must map exactly the leaf nodes" ) - qs = sorted(qubit_of_leaf.values()) + if root_qubit is not None and root in qubit_of_leaf: + raise ValueError("the root cannot carry both a leaf and root qubit") + qs = sorted( + [ + *qubit_of_leaf.values(), + *([] if root_qubit is None else [root_qubit]), + ] + ) if qs != list(range(len(qs))): - raise ValueError("leaf qubits must be 0..n-1 without repeats") + raise ValueError( + "leaf qubits plus root_qubit must be 0..n-1 without repeats" + ) seen = set() stack = [root] @@ -614,13 +675,19 @@ def from_children(cls, children, qubit_of_leaf, *, root=None): raise ValueError( f"nodes not reachable from root {root}: {sorted(unreached)}" ) - return cls(root, children, parent, qubit_of_leaf) + return cls( + root, + children, + parent, + qubit_of_leaf, + root_qubit=root_qubit, + ) #: Fixed number of legs on the top tensor of a :meth:`build_layered` tree. LAYERED_ROOT_ARITY = 3 @classmethod - def build_layered(cls, order, *, block_size=4): + def build_layered(cls, order, *, block_size=4, root_qubit=None): """Build a fixed-structure layered tree with a ternary top tensor. The structure is fixed; only ``block_size`` is tunable: @@ -636,22 +703,30 @@ def build_layered(cls, order, *, block_size=4): Parameters ---------- order : sequence of int - Qubit labels ``0..n-1`` in the desired spatial order. Strongly - coupled qubits should be consecutive so they land in the same - block; use :meth:`TreeLayoutFinder.qubit_order` to obtain an + Leaf-qubit labels in the desired spatial order. Strongly coupled + qubits should be consecutive so they land in the same block; use + :meth:`TreeLayoutFinder.qubit_order` to obtain an entanglement-adapted ordering, or :meth:`TreeLayoutFinder.recommend_layered` to also search - ``block_size``. + ``block_size``. Together with an optional ``root_qubit``, the + labels must cover ``0..n-1``. block_size : int Number of physical qubits per leaf-parent node. Default 4. + root_qubit : int, optional + Qubit label carried by the top tensor rather than a leaf. """ order = list(order) - if not order: + if not order and root_qubit is None: raise ValueError("order must be non-empty.") order = [int(q) for q in order] + if root_qubit is not None: + root_qubit = int(root_qubit) + all_qubits = order + ([] if root_qubit is None else [root_qubit]) n = len(order) - if sorted(order) != list(range(n)): - raise ValueError("order must be a permutation of 0..n-1.") + if sorted(all_qubits) != list(range(len(all_qubits))): + raise ValueError( + "leaf order plus root_qubit must be a permutation of 0..n-1." + ) if not isinstance(block_size, Integral): raise ValueError("block_size must be an integer >= 1.") block_size = int(block_size) @@ -675,6 +750,16 @@ def new_node(): qubit_of_leaf[nid] = q leaf_ids.append(nid) + if not leaf_ids: + root_nid = new_node() + children_map[root_nid] = () + return cls.from_children( + children_map, + qubit_of_leaf, + root=root_nid, + root_qubit=root_qubit, + ) + # First layer: group block_size leaves into one blocking node. # A single-leaf chunk skips the parent and uses the leaf directly. block_nodes = [] @@ -705,7 +790,10 @@ def binary_subtree(nodes): # not add a unary wrapper for n=1 or n <= block_size: it adds a # useless bond and makes the fixed layered family less efficient. return cls.from_children( - children_map, qubit_of_leaf, root=block_nodes[0] + children_map, + qubit_of_leaf, + root=block_nodes[0], + root_qubit=root_qubit, ) root_arity = min(cls.LAYERED_ROOT_ARITY, num_blocks) if num_blocks <= root_arity: @@ -724,7 +812,12 @@ def binary_subtree(nodes): root_nid = new_node() children_map[root_nid] = tuple(root_children) - return cls.from_children(children_map, qubit_of_leaf, root=root_nid) + return cls.from_children( + children_map, + qubit_of_leaf, + root=root_nid, + root_qubit=root_qubit, + ) # -- queries -------------------------------------------------------------- @@ -769,7 +862,8 @@ def max_bond_cut(self): size = {} for x in reversed(visit): ch = self.children[x] - size[x] = sum(size[c] for c in ch) if ch else 1 + local = 1 if x in self.qubit_of_node else 0 + size[x] = local + sum(size[c] for c in ch) best = 0 for x, s in size.items(): if x == self.root: @@ -821,20 +915,38 @@ def subtree_qubit_masks(self): stack.extend(self.children[node]) masks = {} for node in reversed(visit): - if self.is_leaf(node): - masks[node] = 1 << self.qubit_of_leaf[node] - else: - mask = 0 - for child in self.children[node]: - mask |= masks[child] - masks[node] = mask + mask = 0 + q = self.qubit_of_node.get(node) + if q is not None: + mask |= 1 << q + for child in self.children[node]: + mask |= masks[child] + masks[node] = mask return masks def tree_distance(self, qa, qb): - """Return the leaf-to-leaf path length between qubits ``qa`` and ``qb``.""" - la = self.leaf_of_qubit[qa] - lb = self.leaf_of_qubit[qb] - return len(self.node_path(la, lb)) - 1 + """Return the node-path length between physical qubits ``qa`` and ``qb``.""" + na = self.node_of_qubit[qa] + nb = self.node_of_qubit[qb] + return len(self.node_path(na, nb)) - 1 + + def remove_qubit(self, q): + """Return a plan with physical qubit ``q`` removed and labels compacted.""" + q = int(q) + if q == self.root_qubit: + if self.n <= 1: + raise ValueError("cannot remove the only qubit from a tree.") + qubit_of_leaf = { + node: old_q - 1 if old_q > q else old_q + for node, old_q in self.qubit_of_leaf.items() + } + return type(self).from_children( + self.children, + qubit_of_leaf, + root=self.root, + root_qubit=None, + ) + return self.remove_leaf(q) def remove_leaf(self, q): """Return a plan with qubit ``q`` capped and its unary parent removed. @@ -859,9 +971,13 @@ def remove_leaf(self, q): del children[leaf] del qubit_of_leaf[leaf] - # A tree node may not become unary. Keep the old parent id and absorb - # its only surviving child into it; this also handles a two-leaf root. - if len(children[parent]) == 1: + # A virtual-only tree node may not become unary. A physical root is + # different: its one child plus root physical leg is still a meaningful + # rank-two top tensor, so retain that unary structural root. + physical_root = ( + parent == self.root and self.root_qubit is not None + ) + if len(children[parent]) == 1 and not physical_root: child = children[parent][0] children[parent] = children[child] del children[child] @@ -871,15 +987,27 @@ def remove_leaf(self, q): for node, old_q in tuple(qubit_of_leaf.items()): if old_q > q: qubit_of_leaf[node] = old_q - 1 + root_qubit = self.root_qubit + if root_qubit is not None and root_qubit > q: + root_qubit -= 1 return type(self).from_children( - children, qubit_of_leaf, root=self.root + children, + qubit_of_leaf, + root=self.root, + root_qubit=root_qubit, ) def __repr__(self): n_internal = sum(1 for nid in self.nodes() if not self.is_leaf(nid)) + root_site = ( + "" + if self.root_qubit is None + else f", root_qubit={self.root_qubit}" + ) return ( f"TreePlan(n={self.n}, root={self.root}, " - f"internal_nodes={n_internal}, max_arity={self.max_arity()})" + f"internal_nodes={n_internal}, " + f"max_arity={self.max_arity()}{root_site})" ) @@ -895,6 +1023,9 @@ class TreeLayoutFinder: :class:`TreeOptimizer`. Ignored when ``supports`` is given. n : int, optional Number of qubits. Inferred from the stream when omitted. + root_qubit : int, optional + Designated qubit carried by the top tensor instead of a leaf. It remains + part of every path, Steiner-subtree, and congestion calculation. supports : sequence of sequences, optional Explicit interaction supports, used instead of extracting them from ``gates``. @@ -957,7 +1088,7 @@ def __init__(self, gates=None, n=None, *, supports=None, structure="quality", dense_max=512, objective="path", weight_mode="count", chi=None, max_operator_qubits=8, hybrid_weights=None, refine=None, refine_budget=None, search=None, search_budget=128, seed=0, - nevergrad_optimizer="OnePlusOne"): + nevergrad_optimizer="OnePlusOne", root_qubit=None): if ( _looks_like_tree_tensor_network(gates) or _looks_like_tree_tensor_network(supports) @@ -982,6 +1113,14 @@ def __init__(self, gates=None, n=None, *, supports=None, structure="quality", for site in support: if isinstance(site, Integral): inferred = max(inferred, site) + if root_qubit is not None: + try: + root_qubit = int(root_qubit) + except (TypeError, ValueError) as exc: + raise ValueError( + "root_qubit must be an integer or None." + ) from exc + inferred = max(inferred, root_qubit) if n is None: n = inferred + 1 try: @@ -992,6 +1131,11 @@ def __init__(self, gates=None, n=None, *, supports=None, structure="quality", raise ValueError( "Could not infer qubit count; pass n explicitly." ) + if root_qubit is not None: + if not 0 <= root_qubit < n: + raise ValueError( + f"root_qubit {root_qubit!r} is outside 0..{n - 1}." + ) normalized_supports = [] for support in supports: if len(set(support)) != len(support): @@ -1010,6 +1154,10 @@ def __init__(self, gates=None, n=None, *, supports=None, structure="quality", ) normalized_supports.append(tuple(int(site) for site in support)) self.n = n + self.root_qubit = root_qubit + self.leaf_qubits = tuple( + q for q in range(self.n) if q != self.root_qubit + ) self.supports = tuple(normalized_supports) self.structure = structure self.max_arity, self.arity_candidates = _normalize_arity_candidates( @@ -1157,7 +1305,7 @@ def _resolve_search_settings( refine_budget, "refine_budget" ) if refine is not None and refine_budget is None: - refine_budget = max(1, min(self.n - 1, 64)) + refine_budget = max(1, min(len(self.leaf_qubits) - 1, 64)) if search is _DEFAULT_SEARCH_OPTION: search = self.search @@ -1201,13 +1349,20 @@ def _leaf_order(self, plan): def _plan_with_leaf_order(self, plan, order): """Return ``plan``'s immutable topology with a new leaf assignment.""" order = tuple(int(q) for q in order) - if sorted(order) != list(range(self.n)): - raise ValueError("leaf order must be a permutation of 0..n-1.") + if set(order) != set(self.leaf_qubits) or len(order) != len( + self.leaf_qubits + ): + raise ValueError( + "leaf order must contain every non-root qubit exactly once." + ) qubit_of_leaf = dict(plan.qubit_of_leaf) for leaf, qubit in zip(self._leaf_nodes(plan), order): qubit_of_leaf[leaf] = qubit return TreePlan.from_children( - plan.children, qubit_of_leaf, root=plan.root + plan.children, + qubit_of_leaf, + root=plan.root, + root_qubit=plan.root_qubit, ) def _plan_with_leaf_swap(self, plan, left_leaf, right_leaf): @@ -1218,7 +1373,10 @@ def _plan_with_leaf_swap(self, plan, left_leaf, right_leaf): qubit_of_leaf[left_leaf], ) return TreePlan.from_children( - plan.children, qubit_of_leaf, root=plan.root + plan.children, + qubit_of_leaf, + root=plan.root, + root_qubit=plan.root_qubit, ) def _path_score_and_max(self, plan): @@ -1320,7 +1478,7 @@ def _discard_plan_cache(self, plan): if cached is not None and cached[0] is plan: del self._edge_load_cache[id(plan)] - def _refine_plan_greedy(self, plan, *, chi, budget): + def _refine_plan_greedy(self, plan, *, chi, budget, progbar=False): """Greedily improve a fixed topology through adjacent leaf swaps.""" initial_key = self._selection_key(plan, chi) leaf_nodes = self._leaf_nodes(plan) @@ -1339,10 +1497,21 @@ def _refine_plan_greedy(self, plan, *, chi, budget): evaluations = 0 accepted_moves = 0 position = 0 + progress = None + if progbar: + from tqdm import tqdm # pylint: disable=import-outside-toplevel + + progress = tqdm( + total=budget, + desc="tree layout greedy", + leave=False, + ) while position < len(leaf_nodes) - 1 and evaluations < budget: left_leaf = leaf_nodes[position] right_leaf = leaf_nodes[position + 1] evaluations += 1 + if progress is not None: + progress.update() if self.objective == "path": candidate_path_score = self._path_score_after_leaf_swap( current, left_leaf, right_leaf, current_path_score @@ -1374,6 +1543,8 @@ def _refine_plan_greedy(self, plan, *, chi, budget): else: self._discard_plan_cache(candidate) position += 1 + if progress is not None: + progress.close() return current, { "method": "greedy", "evaluations": evaluations, @@ -1383,7 +1554,7 @@ def _refine_plan_greedy(self, plan, *, chi, budget): } def _refine_plan_nevergrad( - self, plan, *, chi, budget, seed, optimizer_name + self, plan, *, chi, budget, seed, optimizer_name, progbar=False ): """Use Nevergrad to refine a leaf assignment before simulation starts.""" try: @@ -1404,23 +1575,48 @@ def _refine_plan_nevergrad( initial_plan = plan initial_key = self._selection_key(initial_plan, chi) initial_order = self._leaf_order(initial_plan) - priorities = np.empty(self.n, dtype=float) - for position, qubit in enumerate(initial_order): - priorities[qubit] = position + if len(initial_order) < 2 or budget < 1: + return initial_plan, { + "method": "nevergrad", + "optimizer": optimizer_name, + "budget": budget, + "evaluations": 0, + "seed": seed, + "initial_key": initial_key, + "final_key": initial_key, + "improved": False, + } + leaf_qubits = tuple(initial_order) + priorities = np.arange(len(leaf_qubits), dtype=float) parametrization = ng.p.Array(init=priorities) if hasattr(parametrization, "set_bounds"): - parametrization.set_bounds(-float(self.n), float(2 * self.n)) + parametrization.set_bounds( + -float(len(leaf_qubits)), + float(2 * len(leaf_qubits)), + ) optimizer = optimizer_class(parametrization=parametrization, budget=budget) random_state = getattr(optimizer.parametrization, "random_state", None) if random_state is not None: random_state.seed(seed) losses = {} + progress = None + if progbar: + from tqdm import tqdm # pylint: disable=import-outside-toplevel + + progress = tqdm( + total=budget, + desc="tree layout nevergrad", + leave=False, + ) def loss(values): + if progress is not None: + progress.update() values = np.asarray(values) order = tuple( - int(q) for q in np.argsort(values, kind="stable") + leaf_qubits[int(position)] + for position in np.argsort(values, kind="stable") ) cached = losses.get(order) if cached is not None: @@ -1431,10 +1627,16 @@ def loss(values): self._discard_plan_cache(candidate) return value - recommendation = optimizer.minimize(loss) + try: + recommendation = optimizer.minimize(loss) + finally: + if progress is not None: + progress.close() final_order = tuple( - int(q) - for q in np.argsort(np.asarray(recommendation.value), kind="stable") + leaf_qubits[int(position)] + for position in np.argsort( + np.asarray(recommendation.value), kind="stable" + ) ) candidate = self._plan_with_leaf_order(initial_plan, final_order) candidate_key = self._selection_key(candidate, chi) @@ -1456,7 +1658,7 @@ def loss(values): "improved": improved, } - def _improve_plan(self, plan, *, chi, settings): + def _improve_plan(self, plan, *, chi, settings, progbar=False): """Run the requested pre-simulation plan refinements in sequence.""" initial_order = self._leaf_order(plan) initial_key = self._selection_key(plan, chi) @@ -1468,7 +1670,10 @@ def _improve_plan(self, plan, *, chi, settings): } if settings["refine"] == "greedy": plan, info["refinement"] = self._refine_plan_greedy( - plan, chi=chi, budget=settings["refine_budget"] + plan, + chi=chi, + budget=settings["refine_budget"], + progbar=progbar, ) if settings["search"] == "nevergrad": plan, info["search"] = self._refine_plan_nevergrad( @@ -1477,6 +1682,7 @@ def _improve_plan(self, plan, *, chi, settings): budget=settings["search_budget"], seed=settings["seed"], optimizer_name=settings["nevergrad_optimizer"], + progbar=progbar, ) info["final_order"] = self._leaf_order(plan) info["final_key"] = self._selection_key(plan, chi) @@ -1496,13 +1702,14 @@ def _build_plan(self, weights, *, structure=None, if cached is not None and cached[0] is weights: return cached[1] plan = TreePlan.from_order( - range(self.n), + self.leaf_qubits, weights=weights, structure=structure, max_arity=max_arity, community_frac=self.community_frac, star_frac=self.star_frac, dense_max=self.dense_max, + root_qubit=self.root_qubit, ) self._plan_cache[key] = (weights, plan) return plan @@ -1580,22 +1787,24 @@ def _select_plan(self, max_arity): def qubit_order(self): """Return a spectral qubit ordering adapted to the gate-stream interactions. - The order is the global Fiedler spectral reordering of all qubits + The order is the global Fiedler spectral reordering of the leaf qubits under the similarity weights used internally by the layout finder. - Strongly coupled qubits end up consecutive, which is the ideal input - for :meth:`TreePlan.build_layered` so that blocks group entangled - qubits together. + A configured ``root_qubit`` is fixed at the root and omitted from this + returned leaf order. Strongly coupled leaf qubits end up consecutive, + which is the ideal input for :meth:`TreePlan.build_layered` so that + blocks group entangled qubits together. Returns ------- list of int - A permutation of ``0..n-1``. + Every non-root qubit exactly once (all ``0..n-1`` qubits when + ``root_qubit`` is ``None``). """ weights = self._similarity_weights() order = _gate_stream_spectral_order( - list(range(self.n)), weights, dense_max=self.dense_max + list(self.leaf_qubits), weights, dense_max=self.dense_max ) - return order if order else list(range(self.n)) + return order if order else list(self.leaf_qubits) def layered(self, block_size=4, *, order=None): """Build a fixed layered tree for a chosen ``block_size`` (no search). @@ -1632,7 +1841,11 @@ def layered(self, block_size=4, *, order=None): order = self.qubit_order() else: order = [int(q) for q in order] - return TreePlan.build_layered(order, block_size=block_size) + return TreePlan.build_layered( + order, + block_size=block_size, + root_qubit=self.root_qubit, + ) def recommend_layered( self, @@ -1646,6 +1859,7 @@ def recommend_layered( search_budget=_DEFAULT_SEARCH_OPTION, seed=_DEFAULT_SEARCH_OPTION, nevergrad_optimizer=_DEFAULT_SEARCH_OPTION, + progbar=False, ): """Optimize the fixed layered structure over ``block_size``. @@ -1687,6 +1901,8 @@ def recommend_layered( only the returned fixed plan; it never mutates a live TTN. search_budget, seed, nevergrad_optimizer Optional Nevergrad configuration for each candidate plan. + progbar : bool, optional + Display greedy and Nevergrad search progress for each candidate. Returns ------- @@ -1725,9 +1941,16 @@ def recommend_layered( candidates = [] for bs in options: - plan = TreePlan.build_layered(order, block_size=bs) + plan = TreePlan.build_layered( + order, + block_size=bs, + root_qubit=self.root_qubit, + ) plan, planning = self._improve_plan( - plan, chi=chi, settings=settings + plan, + chi=chi, + settings=settings, + progbar=progbar, ) report = self.report( plan, include_edge_loads=self.objective != "path" @@ -1772,7 +1995,18 @@ def candidate_key(candidate): "candidates": candidates, } - def run(self): + def run( + self, + *, + chi=_DEFAULT_CHI, + refine=_DEFAULT_SEARCH_OPTION, + refine_budget=_DEFAULT_SEARCH_OPTION, + search=_DEFAULT_SEARCH_OPTION, + search_budget=_DEFAULT_SEARCH_OPTION, + seed=_DEFAULT_SEARCH_OPTION, + nevergrad_optimizer=_DEFAULT_SEARCH_OPTION, + progbar=False, + ): """Return a TreePlan for the selected layout objective. When the finder was built with a set of candidate arities (the default @@ -1780,9 +2014,31 @@ def run(self): :meth:`recommend_arities` -- ``chi``-aware when the finder carries a ``chi`` -- and returns the objective-best plan. A scalar ``max_arity`` builds one fixed plan. + + ``chi`` and the fixed-plan ``refine`` / ``search`` controls can be + overridden for this call. Pass ``progbar=True`` to display greedy and + Nevergrad search progress. Omitted values inherit the corresponding + finder settings, so the original zero-argument behavior is unchanged. """ + if chi is _DEFAULT_CHI: + chi = self.chi + else: + chi = _validate_chi(chi) + settings = self._resolve_search_settings( + refine=refine, + refine_budget=refine_budget, + search=search, + search_budget=search_budget, + seed=seed, + nevergrad_optimizer=nevergrad_optimizer, + ) if self.arity_candidates is not None: - rec = self.recommend_arities(self.arity_candidates, chi=self.chi) + rec = self.recommend_arities( + self.arity_candidates, + chi=chi, + progbar=progbar, + **settings, + ) self._last_arity_recommendation = rec self._selected_candidate = f"arity={rec['recommended_max_arity']}" self._last_candidate_scores = { @@ -1793,19 +2049,23 @@ def run(self): } return rec["plan"] candidates = self._candidate_plans(self.max_arity) - settings = self._resolve_search_settings() if settings["refine"] is not None or settings["search"] is not None: candidates = { - name: self._improve_plan(plan, chi=self.chi, settings=settings)[0] + name: self._improve_plan( + plan, + chi=chi, + settings=settings, + progbar=progbar, + )[0] for name, plan in candidates.items() } selected = min( candidates, - key=lambda name: self._selection_key(candidates[name], self.chi), + key=lambda name: self._selection_key(candidates[name], chi), ) self._last_candidates = candidates self._last_candidate_scores = { - name: self._selection_key(plan, self.chi) + name: self._selection_key(plan, chi) for name, plan in candidates.items() } self._selected_candidate = selected @@ -1822,6 +2082,7 @@ def recommend_arities( search_budget=_DEFAULT_SEARCH_OPTION, seed=_DEFAULT_SEARCH_OPTION, nevergrad_optimizer=_DEFAULT_SEARCH_OPTION, + progbar=False, ): """Compare binary and wider trees and return the best candidate. @@ -1848,6 +2109,8 @@ def recommend_arities( Optional fixed-plan search controls with the same meaning as in :meth:`recommend_layered`. They are applied to each arity candidate before selecting one final immutable plan. + progbar : bool, optional + Display greedy and Nevergrad search progress for each candidate. """ if chi is _DEFAULT_CHI: chi = self.chi @@ -1878,7 +2141,10 @@ def recommend_arities( for arity in options: plan = self._select_plan(arity) plan, planning = self._improve_plan( - plan, chi=chi, settings=settings + plan, + chi=chi, + settings=settings, + progbar=progbar, ) report = self.report( plan, include_edge_loads=self.objective != "path" @@ -2008,24 +2274,24 @@ def edge_loads(self, plan=None): support_mask |= 1 << site # An edge crosses the support iff it belongs to the minimal - # subtree spanning the support leaves. Scanning every tree edge + # subtree spanning the support nodes. Scanning every tree edge # is needlessly O(n) for each event; for the dominant two-qubit - # case this reduces the work to the leaf-to-leaf geodesic. - leaves = [plan.leaf_of_qubit[site] for site in support] - if len(leaves) == 2: + # case this reduces the work to the site-to-site geodesic. + site_nodes = [plan.node_of_qubit[site] for site in support] + if len(site_nodes) == 2: # This branch dominates ordinary circuit layout. Avoid sets # and an all-node parent scan: every path hop is one crossed # rooted tree edge. - path = plan.node_path(leaves[0], leaves[1]) + path = plan.node_path(site_nodes[0], site_nodes[1]) crossed_edges = [ (u, v) if plan.parent.get(v) == u else (v, u) for u, v in zip(path, path[1:]) ] else: span_nodes = set() - anchor = leaves[0] - for leaf in leaves: - span_nodes.update(plan.node_path(anchor, leaf)) + anchor = site_nodes[0] + for site_node in site_nodes: + span_nodes.update(plan.node_path(anchor, site_node)) crossed_edges = [ (parent, node) for node in span_nodes @@ -2102,7 +2368,7 @@ def score(self, plan): """Return the total interaction-weighted tree-path length of ``plan``. Lower is better: this is the quantity the tree structure minimises - (short leaf-to-leaf paths for strongly coupled qubits). + (short physical-node paths for strongly coupled qubits). """ return self._path_score_and_max(plan)[0] @@ -2110,7 +2376,9 @@ def _balanced_plan(self): """Return the cached index-order balanced comparison plan.""" if self._balanced_plan_cache is None: self._balanced_plan_cache = TreePlan.from_order( - range(self.n), structure="balanced" + self.leaf_qubits, + structure="balanced", + root_qubit=self.root_qubit, ) return self._balanced_plan_cache @@ -2118,8 +2386,8 @@ def report(self, plan=None, *, include_edge_loads=True): """Return layout-quality diagnostics for ``plan`` (or a fresh run). The dominant lever for tree-tensor-network accuracy at fixed ``chi`` is - how well the tree keeps strongly coupled qubits as nearby leaves: a - two-qubit gate threads its virtual bond along the whole leaf-to-leaf + how well the tree keeps strongly coupled qubits as nearby nodes: a + two-qubit gate threads its virtual bond along the whole site-to-site geodesic, and every crossed bond can grow. This report summarises those geodesic lengths over the interaction graph and compares the chosen structure against a naive balanced index-order tree (lower ``score`` is @@ -2173,6 +2441,7 @@ def report(self, plan=None, *, include_edge_loads=True): ), "hybrid_cost": hybrid_cost, "root": plan.root, + "root_qubit": plan.root_qubit, "is_binary": plan.is_binary(), "max_arity": plan.max_arity(), "arity_histogram": arity_histogram, diff --git a/src/pepsy/optimizers/tree/optimizer.py b/src/pepsy/optimizers/tree/optimizer.py index e646c18..efbad09 100644 --- a/src/pepsy/optimizers/tree/optimizer.py +++ b/src/pepsy/optimizers/tree/optimizer.py @@ -4,18 +4,19 @@ simulator of *Simulating quantum circuits using tree tensor networks* (Seitz, Medina, Cruz, Huang, Mendl; Quantum 7, 964, 2023; arXiv:2206.01000). -A quantum state is stored as a rooted tree tensor network whose leaves -carry the physical qubit indices. Internal nodes may have any arity: the -default structure is a strictly-binary tree, but flatter ``k``-ary trees +A quantum state is stored as a rooted tree tensor network whose leaves carry +physical qubit indices. One optional physical qubit may instead live on the +root tensor. Internal nodes may have any arity: the default structure is a +strictly-binary tree, but flatter ``k``-ary trees (``max_arity``) or gate-connectivity-driven communities (``structure="adaptive"``) are supported unchanged. A bundled gate stream ``[(gate, where), ...]`` is replayed: -* single-qubit gates are absorbed into the leaf tensor (no bond growth); a +* single-qubit gates are absorbed into their physical-site tensor (no bond growth); a unitary one-qubit gate preserves the tree canonical form regardless of where the orthogonality centre sits; -* two-qubit gates on leaves ``a`` and ``b`` are SVD-split into two factors - joined by a virtual bond; the factors are absorbed into the two leaves and +* two-qubit gates on sites ``a`` and ``b`` are SVD-split into two factors + joined by a virtual bond; the factors are absorbed into the two site nodes and the virtual bond is *threaded exactly* (no truncation) along the tree path from ``a`` to ``b``. Only once both factors are in place is a single canonical compression sweep run back along the path, truncating every @@ -115,6 +116,7 @@ def _same_tree_plan(left, right): and left.root == right.root and left.children == right.children and left.qubit_of_leaf == right.qubit_of_leaf + and left.root_qubit == right.root_qubit ) @@ -331,6 +333,12 @@ class TreeOptimizer: tree : TreePlan, optional Explicit tree structure (any arity). When omitted a :class:`TreeLayoutFinder` builds one from the gate stream. + root_qubit : int, optional + Designated qubit carried by the top tensor instead of a leaf when the + layout is built automatically. Gates, sub-MPOs, readout, capping, and + layout scoring treat it as an ordinary physical site at the root. When + ``tree`` or ``layout`` is supplied, this must match its + ``TreePlan.root_qubit``. dtype : numpy dtype Data type of the initial product state (default ``complex128``). threads : int or None @@ -417,6 +425,7 @@ def __init__(self, gates=None, n=None, *, chi=64, cutoff=1e-12, structure="quality", max_arity=(2, 3, 4), community_frac=0.35, star_frac=0.75, layout_objective="path", layout_weight_mode="count", layout=None, tree=None, + root_qubit=None, dtype=complex, threads=1, seed=None, run=True, tn=None, state=None, track_truncation=False, track_infidelity=True, max_intermediate_bond=None, @@ -453,6 +462,13 @@ def __init__(self, gates=None, n=None, *, chi=64, cutoff=1e-12, "tree= expects a TreePlan, not a TreeTensorNetwork; " "pass the entangled state as state= or tn=." ) + if root_qubit is not None: + try: + root_qubit = int(root_qubit) + except (TypeError, ValueError) as exc: + raise ValueError( + "root_qubit must be an integer or None." + ) from exc self.G, self.where, self.event_types = self._normalize_gate_queue(gates) self.layout_finder = layout if isinstance(layout, TreeLayoutFinder) else None @@ -514,9 +530,15 @@ def __init__(self, gates=None, n=None, *, chi=64, cutoff=1e-12, (max(w) for w in self.where if len(w) > 0), default=-1, ) + if root_qubit is not None: + n = max(n, root_qubit + 1) self.n = int(n) if self.n <= 0: raise ValueError("Could not infer qubit count; pass n explicitly.") + if root_qubit is not None and not 0 <= root_qubit < self.n: + raise ValueError( + f"root_qubit {root_qubit!r} is outside 0..{self.n - 1}." + ) # The TTN itself always uses compact physical positions. This facade # optionally preserves caller-facing logical labels across a cap while # keeping Quimb's internal site/index space contiguous. @@ -590,10 +612,15 @@ def __init__(self, gates=None, n=None, *, chi=64, cutoff=1e-12, weight_mode=self.layout_weight_mode, chi=self.chi, max_operator_qubits=self.max_operator_qubits, + root_qubit=root_qubit, ) tree = self.layout_finder.run() if not isinstance(tree, TreePlan): raise TypeError("tree must be a TreePlan or None.") + if root_qubit is not None and tree.root_qubit != root_qubit: + raise ValueError( + "root_qubit does not match the supplied tree/layout plan." + ) self.plan = tree if product_state_source is not None: @@ -1060,9 +1087,9 @@ def _neighbors(self, nid): """Return the adjacent node ids of ``nid`` (children plus parent).""" return self.tn.neighbors(nid) - def _steiner_nodes(self, leaves): - """Return the node set of the minimal subtree spanning ``leaves``.""" - return self.tn.steiner_nodes(leaves) + def _steiner_nodes(self, nodes): + """Return the node set of the minimal subtree spanning ``nodes``.""" + return self.tn.steiner_nodes(nodes) def _build_product_state(self): return TreeTensorNetwork.from_plan(self.plan, dtype=self.dtype) @@ -1071,7 +1098,7 @@ def _build_product_state(self): def _product_site_vector(state, q): """Extract one qubit's vector from a bond-dimension-one TN site.""" if isinstance(state, TreeTensorNetwork): - tensor = state.node_tensor(state.leaf_of_qubit(q)) + tensor = state.node_tensor(state.node_of_qubit(q)) physical_index = state.site_ind(q) else: site_tag = state.site_tag(q) @@ -1105,8 +1132,8 @@ def _remount_product_state(self, state): sample = next(iter(state.tensor_map.values())).data for node in self.plan.nodes(): tensor = target.node_tensor(node) - if self.plan.is_leaf(node): - q = self.plan.qubit_of_leaf[node] + q = self.plan.qubit_of_node.get(node) + if q is not None: vector = self._product_site_vector(state, q) tensor.modify(data=ar.do("reshape", vector, ar.shape(tensor.data))) else: @@ -1120,7 +1147,7 @@ def _remount_product_state(self, state): if isinstance(state, TreeTensorNetwork): factor = None for node in state.plan.nodes(): - if state.plan.is_leaf(node): + if node in state.plan.qubit_of_node: continue scalar = ar.do("reshape", state.node_tensor(node).data, ()) factor = scalar if factor is None else factor * scalar @@ -1191,7 +1218,7 @@ def _validate_support(self, where, *, min_size=1, resolve=True): if resolve: return tuple(self._validate_qubit(q) for q in where) for q in where: - if q not in self.plan.leaf_of_qubit: + if q not in self.plan.node_of_qubit: raise ValueError(f"tree position {q} is outside the state.") return where @@ -1226,8 +1253,8 @@ def _check_operator_limits(self, where, *, dense=True): f"max_operator_qubits={self.max_operator_qubits}." ) if self.max_subtree_nodes is not None and len(where) > 1: - leaves = [self.plan.leaf_of_qubit[q] for q in where] - span = self._steiner_nodes(leaves) + site_nodes = [self.plan.node_of_qubit[q] for q in where] + span = self._steiner_nodes(site_nodes) if len(span) > self.max_subtree_nodes: raise MemoryError( f"operator Steiner subtree has {len(span)} nodes, exceeding " @@ -1236,8 +1263,8 @@ def _check_operator_limits(self, where, *, dense=True): def _projection_snapshot(self, where): """Return compact support/span/bond diagnostics for a projection.""" - leaves = [self.plan.leaf_of_qubit[q] for q in where] - span = frozenset(self._steiner_nodes(leaves)) + site_nodes = [self.plan.node_of_qubit[q] for q in where] + span = frozenset(self._steiner_nodes(site_nodes)) bonds = {} for node in span: for neighbour in self._neighbors(node): @@ -1395,7 +1422,7 @@ def canonize_around_qubits(self, qubits): """Canonicalise around the minimal subtree spanning ``qubits``. The qubit-level "range canonicalisation" entry point: gauge every tensor - outside the minimal connected subtree spanning the given qubits' leaves + outside the minimal connected subtree spanning the given qubits' nodes to point inward, so the reduced state on those qubits is captured by that subtree. Delegates to :meth:`TreeTensorNetwork.canonize_around_qubits_`. Returns ``self``. @@ -1435,10 +1462,10 @@ def canonize_mps(self, p, where, *, info=None): if len(sites) not in {1, 2}: raise ValueError("where must be an int, (int,), or (int, int).") for q in sites: - if q not in p.plan.leaf_of_qubit: + if q not in p.plan.node_of_qubit: raise ValueError(f"qubit {q} is outside the tree state.") if len(sites) == 1: - p.shift_orthogonality_center(p.plan.leaf_of_qubit[sites[0]]) + p.shift_orthogonality_center(p.plan.node_of_qubit[sites[0]]) target = (sites[0], sites[0]) else: p.canonize_around_qubits_(sites) @@ -1698,7 +1725,7 @@ def run(self, gates=None, *, progbar=False, mode=None, non_unitary=False, # replay gets the two-site factor fast path as well as # the native multi-site MPO router. Passing both forms # of support is important after a stable-label cap: - # ``support`` addresses compact TTN leaves, while + # ``support`` addresses compact TTN sites, while # ``logical_support`` addresses the MPO site tags. support = self._validate_support(logical_support) self._apply_submpo_resolved( @@ -1889,7 +1916,7 @@ def is_submpo_event(entry): return submpo_event_parts(entry) is not None def apply_1q(self, gate, q, *, renormalize=False): - """Absorb a one-qubit gate into the leaf tensor of qubit ``q``.""" + """Absorb a one-qubit gate into the site tensor of qubit ``q``.""" self._invalidate_state_norm_cache() with self._thread_ctx(): return self._apply_1q_impl(gate, q, renormalize=renormalize) @@ -1903,7 +1930,7 @@ def _apply_1q_impl(self, gate, q, *, renormalize=False): # A Symmray gate is already a (d, d) symmetric operator; reshaping # into base-2 sub-legs would destroy its block/charge structure. Its # unitarity cannot be cheaply certified here, so take the always-safe - # non-unitary branch (move the centre onto the leaf first). + # non-unitary branch (move the centre onto the site node first). unitary = False else: if tuple(ar.shape(gate)) != (d, d): @@ -1914,7 +1941,7 @@ def _apply_1q_impl(self, gate, q, *, renormalize=False): rtol=1e-10, atol=1e-12, ) if not unitary: - self._move_center(self.plan.leaf_of_qubit[q]) + self._move_center(self.plan.node_of_qubit[q]) region = self.tn.canonical_region self.tn.gate_inds_(gate, [self._phys(q)], contract=True) if unitary: @@ -1924,13 +1951,13 @@ def _apply_1q_impl(self, gate, q, *, renormalize=False): # canonical-preserving operation. self.tn.canonical_region = region if not unitary: - self.center = self.plan.leaf_of_qubit[q] + self.center = self.plan.node_of_qubit[q] if renormalize: self.normalize() return self def apply_2q(self, gate, qa, qb): - """Apply a two-qubit gate to leaves ``qa`` and ``qb``. + """Apply a two-qubit gate to physical sites ``qa`` and ``qb``. Following Seitz et al. (Figs. 3-6): SVD-split the gate into two factors joined by a virtual bond, absorb the left factor into leaf ``a`` and the @@ -2147,10 +2174,15 @@ def _apply_2q_factors_impl( one canonical compression sweep. """ plan = self.plan - la = plan.leaf_of_qubit[qa] - lb = plan.leaf_of_qubit[qb] + la = plan.node_of_qubit[qa] + lb = plan.node_of_qubit[qb] parent = plan.parent.get(la) - if parent is not None and plan.parent.get(lb) == parent: + if ( + plan.is_leaf(la) + and plan.is_leaf(lb) + and parent is not None + and plan.parent.get(lb) == parent + ): return self._apply_2q_sibling_factors( factors, outputs, qa, qb, la, lb, parent, max_bond=max_bond, cutoff=cutoff, @@ -2161,12 +2193,12 @@ def _apply_2q_factors_impl( if self._nearest_anchor((la, lb)) == la else (qb, qa) ) - source_leaf = plan.leaf_of_qubit[source] - destination_leaf = plan.leaf_of_qubit[destination] - self._move_center(source_leaf) + source_node = plan.node_of_qubit[source] + destination_node = plan.node_of_qubit[destination] + self._move_center(source_node) self._thread_ind = thread_ind try: - source_tensor = self.tn.tensor_map[self._tid(source_leaf)] + source_tensor = self.tn.tensor_map[self._tid(source_node)] merged_source = qtn.tensor_contract( source_tensor, factors[source] ).reindex_({outputs[source]: self._phys(source)}) @@ -2174,18 +2206,18 @@ def _apply_2q_factors_impl( data=merged_source.data, inds=merged_source.inds, ) - path = plan.node_path(source_leaf, destination_leaf) + path = plan.node_path(source_node, destination_node) for u, v in zip(path, path[1:]): self._thread_hop(u, v) - destination_tensor = self.tn.tensor_map[self._tid(destination_leaf)] + destination_tensor = self.tn.tensor_map[self._tid(destination_node)] merged_destination = qtn.tensor_contract( factors[destination], destination_tensor, ).reindex_({outputs[destination]: self._phys(destination)}) destination_tensor.modify( data=merged_destination.data, inds=merged_destination.inds, ) - self.center = destination_leaf + self.center = destination_node self._compress_path(path, max_bond=max_bond, cutoff=cutoff) finally: self._thread_ind = None @@ -2824,7 +2856,7 @@ def _apply_subtree_operator_impl(self, op, where, *, max_bond=None, The operator is first factorized into an exact tree-MPO on the *minimal connected subtree* (Steiner subtree) spanning the target - leaves. It is then applied recursively from the subtree leaves toward + physical nodes. It is then applied recursively from the subtree leaves toward a hub: each local state/operator message is QR-split losslessly on one edge and immediately absorbed by its parent. Thus no dense state tensor for the whole Steiner subtree is formed. This is the tree analogue of a @@ -2912,23 +2944,23 @@ def _apply_subtree_operator_impl(self, op, where, *, max_bond=None, with self._thread_ctx(): if k == 1: - # Single-site operator (possibly non-unitary): centre on the leaf + # Single-site operator (possibly non-unitary): centre on its node # so it holds the (rescaled) norm, then absorb the operator. - leaf = self.plan.leaf_of_qubit[where[0]] - self._move_center(leaf) + site_node = self.plan.node_of_qubit[where[0]] + self._move_center(site_node) self.tn.gate_inds_(op_arr, [phys[0]], contract=True) - self.center = leaf + self.center = site_node if renormalize: self.normalize() return self - leaves = [self.plan.leaf_of_qubit[q] for q in where] - snodes = self._steiner_nodes(leaves) - # Centre on a target leaf so the whole exterior is isometric toward + site_nodes = [self.plan.node_of_qubit[q] for q in where] + snodes = self._steiner_nodes(site_nodes) + # Centre on a target physical node so the whole exterior is isometric toward # the subtree. Operator bonds are routed losslessly first; the # final subtree sweep then measures true state error against that # complete operator update. - anchor = self._nearest_anchor(leaves) + anchor = self._nearest_anchor(site_nodes) if self.center != anchor: self._move_center(anchor) @@ -2993,9 +3025,9 @@ def _try_apply_native_submpo( if set(present) != set(payload_where): return None - leaves = [self.plan.leaf_of_qubit[q] for q in where] - snodes = self._steiner_nodes(leaves) - self._move_center(self._nearest_anchor(leaves)) + site_nodes = [self.plan.node_of_qubit[q] for q in where] + snodes = self._steiner_nodes(site_nodes) + self._move_center(self._nearest_anchor(site_nodes)) order, hub = self._peel_order(snodes) local = {} state_inds = {} @@ -3003,7 +3035,7 @@ def _try_apply_native_submpo( for nid in snodes: state_t = self.tn.tensor_map[self._tid(nid)].copy() state_inds[nid] = set(state_t.inds) - q = self.plan.qubit_of_leaf.get(nid) + q = self.plan.qubit_of_node.get(nid) if q is None: local[nid] = state_t operator_inds[nid] = set() @@ -3064,10 +3096,10 @@ def _apply_factorized_subtree_operator_impl( state_t = self.tn.tensor_map[self._tid(nid)].copy() state_inds[nid] = set(state_t.inds) op_t = op_factors[nid] - q = self.plan.qubit_of_leaf.get(nid) + q = self.plan.qubit_of_node.get(nid) if q is not None and q in where: # Operator sites are packed into one dimension-four leg. Split - # that leg only at physical leaves, then contract its input leg + # that leg only at physical sites, then contract its input leg # with the live state physical index. op_t = self._expand_tree_operator_leaf( op_t, @@ -3125,12 +3157,12 @@ def _decompose_tree_operator(self, op_arr, where, snodes, order, hub): interleaved = ar.do("reshape", interleaved, (4,) * len(where)) blob = qtn.Tensor(interleaved, inds=op_axes) - leaf_for_q = { - self.plan.leaf_of_qubit[q]: q for q in where + node_for_q = { + self.plan.node_of_qubit[q]: q for q in where } owned = {nid: set() for nid in snodes} - for leaf, q in leaf_for_q.items(): - owned[leaf].add(op_axes[where.index(q)]) + for node, q in node_for_q.items(): + owned[node].add(op_axes[where.index(q)]) factors = {} op_bonds = {"physical": dict(zip(where, op_axes))} @@ -3166,7 +3198,7 @@ def _apply_product_pauli_projector_impl( branch_index = {} for nid in snodes: state_t = self.tn.tensor_map[self._tid(nid)].copy() - q = self.plan.qubit_of_leaf.get(nid) + q = self.plan.qubit_of_node.get(nid) if q in target_axes: p = self._phys(q) branch = f"_ttn_pauli_branch_{qtn.rand_uuid()}" @@ -3319,10 +3351,10 @@ def _apply_product_pauli_projector( def _product_pauli_expectation(self, axes, where): """Evaluate a product-Pauli expectation using one-site insertions.""" - leaves = [self.plan.leaf_of_qubit[q] for q in where] - snodes = self._steiner_nodes(leaves) + site_nodes = [self.plan.node_of_qubit[q] for q in where] + snodes = self._steiner_nodes(site_nodes) if self.center not in snodes: - self._move_center(leaves[0]) + self._move_center(site_nodes[0]) internal = set() for nid in snodes: @@ -3407,10 +3439,10 @@ def show(self, *, bond_dims=True, node_ids=False, color=True): """Print a top-down ASCII drawing of the tree with current bond dims. Delegates to :meth:`TreeTensorNetwork.show`: the root sits at the top, - the qubit leaves at the bottom, internal nodes are ``●``, leaves ``◆`` - labelled with their qubit, and each edge carries its virtual-bond - dimension -- the tree analogue of a ``quimb`` MPS ``show``. Markers are - coloured by tree layer by default; pass ``color=False`` for plain text. + structural leaves at the bottom, physical nodes are labelled with their + qubits, and each edge carries its virtual-bond dimension -- the tree + analogue of a ``quimb`` MPS ``show``. Markers are coloured by tree layer + by default; pass ``color=False`` for plain text. """ self.tn.show(bond_dims=bond_dims, node_ids=node_ids, color=color) @@ -3486,13 +3518,13 @@ def edge_key(a, b): edge_bonds = {} def plan_steiner(plan, support): - leaves = [plan.leaf_of_qubit[q] for q in support] - if len(leaves) == 1: - return {leaves[0]} + site_nodes = [plan.node_of_qubit[q] for q in support] + if len(site_nodes) == 1: + return {site_nodes[0]} nodes = set() - anchor = leaves[0] - for leaf in leaves[1:]: - nodes.update(plan.node_path(anchor, leaf)) + anchor = site_nodes[0] + for site_node in site_nodes[1:]: + nodes.update(plan.node_path(anchor, site_node)) return nodes def plan_edges(plan): @@ -3590,7 +3622,7 @@ def plan_edges(plan): raise ValueError( f"cap event at step {index + 1} cannot remove the only site." ) - sim_plan = sim_plan.remove_leaf(support[0]) + sim_plan = sim_plan.remove_qubit(support[0]) capped = logical_support[0] active.remove(capped) if payload.get("compact_labels", True): @@ -3803,9 +3835,9 @@ def _leaf_canonical_norm(self): """ if self.tn.fermionic: return self.norm() - leaf = self.plan.leaf_of_qubit[min(self.plan.leaf_of_qubit)] - self._move_center(leaf) - t = self.tn.tensor_map[self._tid(leaf)] + site_node = self.plan.node_of_qubit[min(self.plan.node_of_qubit)] + self._move_center(site_node) + t = self.tn.tensor_map[self._tid(site_node)] val = qtn.tensor_contract(t.H, t, output_inds=[]) return float(np.sqrt(abs(to_float(val, real=True)))) @@ -3937,18 +3969,18 @@ def _apply_control_event_impl(self, name, payload, where): def measure(self, q, outcome=None): """Projectively measure qubit ``q`` in the computational basis. - Moves the orthogonality centre onto the leaf, reads the single-site Born + Moves the orthogonality centre onto the site node, reads the Born probabilities from that one canonical tensor, samples (or forces via - ``outcome``) a result, projects the leaf onto it, and renormalises. - Returns the outcome bit. Because the centre sits on the leaf the + ``outcome``) a result, projects the site, and renormalises. + Returns the outcome bit. Because the centre sits on the site node the probabilities are exact regardless of the global state norm. """ self._require_dense_qubit_state("measure") with self._thread_ctx(): q = self._validate_qubit(q) - leaf = self.plan.leaf_of_qubit[q] - self._move_center(leaf) - t = self.tn.tensor_map[self._tid(leaf)] + site_node = self.plan.node_of_qubit[q] + self._move_center(site_node) + t = self.tn.tensor_map[self._tid(site_node)] p = self._phys(q) ax = t.inds.index(p) arr = ar.do("reshape", ar.do("moveaxis", t.data, ax, 0), (2, -1)) @@ -4112,6 +4144,7 @@ def find_tree_layout(cls, gates, n=None, *, structure="quality", max_arity=(2, 3, 4), community_frac=0.35, star_frac=0.75, layout_objective="path", layout_weight_mode="count", + root_qubit=None, max_operator_qubits=_DEFAULT_MAX_OPERATOR_QUBITS): """Return the :class:`TreePlan` a :class:`TreeLayoutFinder` would use.""" return TreeLayoutFinder( @@ -4119,6 +4152,7 @@ def find_tree_layout(cls, gates, n=None, *, structure="quality", max_arity=max_arity, community_frac=community_frac, star_frac=star_frac, objective=layout_objective, weight_mode=layout_weight_mode, + root_qubit=root_qubit, max_operator_qubits=max_operator_qubits, ).run() @@ -4126,7 +4160,7 @@ def find_tree_layout(cls, gates, n=None, *, structure="quality", def convergence_sweep(cls, gates, n=None, chi_values=(2, 4, 8, 16, 32), *, ops=None, structure="quality", max_arity=(2, 3, 4), community_frac=0.35, star_frac=0.75, tree=None, - dense_cap=1 << 14): + root_qubit=None, dense_cap=1 << 14): """Replay ``gates`` at several ``chi`` and report convergence. The tree structure is built once and reused for every ``chi`` so the @@ -4163,7 +4197,7 @@ def convergence_sweep(cls, gates, n=None, chi_values=(2, 4, 8, 16, 32), *, if tree is None: probe = cls(gates, n=n, structure=structure, max_arity=max_arity, community_frac=community_frac, star_frac=star_frac, - run=False) + root_qubit=root_qubit, run=False) tree = probe.plan n = probe.n elif n is None: diff --git a/src/pepsy/optimizers/tree/ttn.py b/src/pepsy/optimizers/tree/ttn.py index c915585..93a3683 100644 --- a/src/pepsy/optimizers/tree/ttn.py +++ b/src/pepsy/optimizers/tree/ttn.py @@ -15,11 +15,14 @@ ---------------------- * every node of the plan (leaf **and** internal) is one tensor, tagged with the structural node tag ``node_tag_id.format(nid)`` (default ``"N{}"``); -* leaf tensors additionally carry the ``quimb`` site tag +* physical-site tensors carry the ``quimb`` site tag ``site_tag_id.format(q)`` (default ``"I{}"``) and the physical index - ``site_ind_id.format(q)`` (default ``"k{}"``) for qubit ``q`` -- so Quimb - treats the leaves as the sites and the internal nodes as ancillary bond - carriers. This class supplies the tree-specific ``local_expectation`` path; + ``site_ind_id.format(q)`` (default ``"k{}"``) for qubit ``q``; these are + structural leaves by default; +* a plan may designate one additional ``root_qubit`` carried by the top tensor. + A binary root then has exactly two child bonds plus this physical leg. Other + internal nodes remain ancillary bond carriers. This class supplies the + tree-specific ``local_expectation`` path for both leaf and root sites; * adjacent nodes ``a`` and ``b`` share the deterministic virtual bond index ``_tb{lo}_{hi}`` with ``lo, hi = sorted((a, b))``. @@ -149,7 +152,7 @@ def _color(s, code, enable): class TreeTensorNetwork(TensorNetworkGenVector): - """A rooted tree-tensor-network state over qubit leaves. + """A rooted tree-tensor-network state over physical qubit nodes. Subclasses :class:`quimb.tensor.TensorNetworkGenVector`, so it *is* a ``quimb`` tensor network: all of ``quimb``'s arbitrary-geometry methods @@ -201,6 +204,7 @@ def __init__(self, ts=(), *, plan=None, sites=None, site_tag_id="I{}", plan.root != ts.plan.root or plan.children != ts.plan.children or plan.qubit_of_leaf != ts.plan.qubit_of_leaf + or plan.root_qubit != ts.plan.root_qubit ): raise ValueError( "plan does not match the TreeTensorNetwork being copied." @@ -421,7 +425,7 @@ def local_expectation( ): """Evaluate a local observable with a backend-specific exact path. - Dense/nonfermionic trees use the canonical target leaf or minimal + Dense/nonfermionic trees use the canonical target physical node or minimal Steiner subtree and cancel its ordinary isometric exterior. Native fermionic trees keep the Symmray operator structured and contract the complete doubled tree so graded boundary phases are never discarded. @@ -437,7 +441,7 @@ def local_expectation( where = tuple(int(site) for site in where) if not where or len(set(where)) != len(where): raise ValueError("where must contain distinct tree sites.") - if any(site not in self.plan.leaf_of_qubit for site in where): + if any(site not in self.plan.node_of_qubit for site in where): raise ValueError(f"site(s) {where!r} are outside this tree state.") if not self.fermionic and preserve_gauge: @@ -459,7 +463,7 @@ def local_expectation( else: original_region = None - leaves = [self.plan.leaf_of_qubit[site] for site in where] + site_nodes = [self.plan.node_of_qubit[site] for site in where] phys = [self.site_ind(site) for site in where] op = operator if self.symmetry is not None and not _is_symmray_array(op): @@ -481,8 +485,8 @@ def local_expectation( normalized=normalized, ) if len(where) == 1: - self.shift_orthogonality_center(leaves[0]) - tensor = self.node_tensor(leaves[0]) + self.shift_orthogonality_center(site_nodes[0]) + tensor = self.node_tensor(site_nodes[0]) physical = phys[0] if not _is_symmray_array(op): dim = int(tensor.shape[tensor.inds.index(physical)]) @@ -498,9 +502,9 @@ def local_expectation( self._restore_readout_region(original_region) return result - span = self.steiner_nodes(leaves) + span = self.steiner_nodes(site_nodes) if self.orthogonality_center not in span: - self.shift_orthogonality_center(leaves[0]) + self.shift_orthogonality_center(site_nodes[0]) internal = { self.bond(node, neighbor) @@ -513,10 +517,10 @@ def local_expectation( ]) if not _is_symmray_array(op): dims = [ - int(self.node_tensor(leaf).shape[ - self.node_tensor(leaf).inds.index(physical) + int(self.node_tensor(node).shape[ + self.node_tensor(node).inds.index(physical) ]) - for leaf, physical in zip(leaves, phys) + for node, physical in zip(site_nodes, phys) ] op = ar.do("reshape", op, tuple(dims + dims)) elif len(ar.shape(op)) != 2 * len(where): @@ -637,7 +641,7 @@ def _with_center(self, nid): @property def nqubits(self): - """Number of qubit leaves (an alias of :attr:`nsites`).""" + """Number of physical qubits (an alias of :attr:`nsites`).""" return self._plan.n def node_tag(self, nid): @@ -684,9 +688,10 @@ def validate(self, *, check_canonical=False, tol=1e-9): The check is intentionally structural by default and therefore cheap enough for construction and resource-preflight paths. It verifies that - every planned node has exactly one tensor, every leaf owns exactly one - physical index, every plan edge has exactly one live bond, and there are - no extra tensors or malformed shared indices. Pass + every planned node has exactly one tensor, every planned physical site + (a leaf or the optional root site) owns exactly one physical index, + every plan edge has exactly one live bond, and there are no extra + tensors, outer legs, or malformed shared indices. Pass ``check_canonical=True`` to additionally verify the tracked canonical region; that part performs tensor contractions and is more expensive. @@ -723,22 +728,23 @@ def validate(self, *, check_canonical=False, tol=1e-9): for nid, tid in node_tids.items(): tensor = self.tensor_map[tid] node_phys = physical_inds.intersection(tensor.inds) - if self._plan.is_leaf(nid): - q = self._plan.qubit_of_leaf[nid] + q = self._plan.qubit_of_node.get(nid) + if q is not None: expected = self.site_ind(q) if expected not in tensor.inds: raise ValueError( - f"leaf node {nid} (qubit {q}) is missing physical " + f"tree node {nid} (qubit {q}) is missing physical " f"index {expected!r}." ) if len(node_phys) != 1 or expected not in node_phys: raise ValueError( - f"leaf node {nid} must own only physical index " + f"tree node {nid} must own only physical index " f"{expected!r}; found {sorted(node_phys)!r}." ) if self.site_tag(q) not in tensor.tags: raise ValueError( - f"leaf node {nid} is missing site tag {self.site_tag(q)!r}." + f"tree node {nid} is missing site tag " + f"{self.site_tag(q)!r}." ) elif node_phys: raise ValueError( @@ -751,9 +757,16 @@ def validate(self, *, check_canonical=False, tol=1e-9): for ind, ind_owners in owners.items(): if len(ind_owners) != 1: raise ValueError( - f"physical index {ind!r} must belong to one leaf; " + f"physical index {ind!r} must belong to one planned node; " f"found nodes {ind_owners!r}." ) + unexpected_outer = set(self.outer_inds()) - physical_inds + if unexpected_outer: + raise ValueError( + "live tree has unregistered outer indices " + f"{sorted(unexpected_outer)!r}; represent a top physical leg " + "with TreePlan(root_qubit=...)." + ) expected_edges = { frozenset((parent, child)) @@ -819,7 +832,7 @@ def validate(self, *, check_canonical=False, tol=1e-9): # -- plan delegators ------------------------------------------------------ def is_leaf(self, nid): - """Whether node ``nid`` is a leaf (carries a physical qubit).""" + """Whether ``nid`` is a structural leaf (has no children).""" return self._plan.is_leaf(nid) def parent(self, nid): @@ -846,39 +859,47 @@ def leaf_of_qubit(self, q): """Return the leaf node id carrying qubit ``q``.""" return self._plan.leaf_of_qubit[q] + def node_of_qubit(self, q): + """Return the leaf or physical root node carrying qubit ``q``.""" + return self._plan.node_of_qubit[q] + def qubit_of_leaf(self, nid): """Return the qubit label carried by leaf node ``nid``.""" return self._plan.qubit_of_leaf[nid] + def qubit_of_node(self, nid): + """Return the qubit carried by ``nid``, or ``None`` for a virtual node.""" + return self._plan.qubit_of_node.get(nid) + def tree_distance(self, qa, qb): - """Return the leaf-to-leaf path length between qubits ``qa`` and ``qb``.""" + """Return the site-node path length between qubits ``qa`` and ``qb``.""" return self._plan.tree_distance(qa, qb) - def steiner_nodes(self, leaves): - """Return the node set of the minimal subtree spanning ``leaves``. + def steiner_nodes(self, nodes): + """Return the node set of the minimal subtree spanning ``nodes``. The tree has a unique path between any two nodes, so the union of the - paths from ``leaves[0]`` to every other leaf is exactly the minimal + paths from ``nodes[0]`` to every other node is exactly the minimal connected subtree (Steiner tree) that contains all of them. """ - leaves = list(leaves) - if not leaves: - raise ValueError("need at least one leaf to span a subtree.") - for leaf in leaves: - if leaf not in self._plan.children: - raise ValueError(f"{leaf!r} is not a node of the tree.") - root_leaf = leaves[0] - nodes = set() - for lf in leaves: - nodes.update(self._plan.node_path(root_leaf, lf)) - return nodes + nodes = list(nodes) + if not nodes: + raise ValueError("need at least one node to span a subtree.") + for node in nodes: + if node not in self._plan.children: + raise ValueError(f"{node!r} is not a node of the tree.") + root_node = nodes[0] + span = set() + for node in nodes: + span.update(self._plan.node_path(root_node, node)) + return span def subtree_span(self, nodes): """Return the node set of the minimal connected subtree spanning ``nodes``. - Generalises :meth:`steiner_nodes` to *arbitrary* nodes (leaves and - internal): the union of the unique tree paths from ``nodes[0]`` to every - other node is the minimal connected subtree containing them all. + Alias-like generalisation retained for callers that work directly with + structural nodes: the union of the unique tree paths from ``nodes[0]`` + to every other node is the minimal connected subtree containing them all. """ nodes = list(nodes) if not nodes: @@ -1114,14 +1135,14 @@ def canonize_around_qubits_(self, qubits, *, absorb="right"): The qubit-level "range canonicalisation" entry point: given a set of qubit labels, gauge every tensor outside the minimal connected subtree - that spans those qubits' leaves to point inward, so the reduced state on + that spans those qubits' physical nodes to point inward, so the reduced state on those qubits is captured by that subtree. Equivalent to - ``canonize_subtree_(leaves_of(qubits), span=True)``. Returns ``self``. + ``canonize_subtree_(nodes_of(qubits), span=True)``. Returns ``self``. """ if isinstance(qubits, Integral): qubits = (qubits,) - leaves = [self.leaf_of_qubit(q) for q in qubits] - return self.canonize_subtree_(leaves, span=True, absorb=absorb) + site_nodes = [self.node_of_qubit(q) for q in qubits] + return self.canonize_subtree_(site_nodes, span=True, absorb=absorb) def _recover_center_from_region(self, region, target, *, absorb="right"): """Recover one centre by peeling a tracked canonical region. @@ -1294,25 +1315,23 @@ def is_subtree_canonical_form(self, nodes=None, *, span=False, tol=1e-9): return True def cap_qubit_(self, q, vec): - """Contract qubit ``q`` with ``vec`` and remove that leaf in place. + """Contract qubit ``q`` with ``vec`` and remove that site in place. This is the tree counterpart of an MPS physical-index cap. The capped - leaf is absorbed into its parent; if that creates a unary parent, the - parent and its remaining child are fused. Remaining qubit labels are - compacted above ``q`` so subsequent stream entries retain MPS-style - positional semantics. + leaf is absorbed into its parent; if that creates a virtual-only unary + parent, the parent and its remaining child are fused. A physical root + qubit is contracted directly on the root without changing the tree + edges. Remaining qubit labels are compacted above ``q`` so subsequent + stream entries retain MPS-style positional semantics. """ q = int(q) self._invalidate_norm_cache() - if q not in self._plan.leaf_of_qubit: + if q not in self._plan.node_of_qubit: raise ValueError(f"cap qubit {q} is outside the tree state.") if self._plan.n <= 1: raise ValueError("cannot cap the only qubit in a tree state.") - leaf = self._plan.leaf_of_qubit[q] - parent = self._plan.parent.get(leaf) - if parent is None: - raise ValueError("cannot cap the root leaf of a multi-qubit tree.") - like = self.node_tensor(leaf).data + site_node = self._plan.node_of_qubit[q] + like = self.node_tensor(site_node).data try: compatible = ( ar.infer_backend(vec) == ar.infer_backend(like) @@ -1332,6 +1351,55 @@ def cap_qubit_(self, q, vec): f"physical dimension {self.ind_size(phys)} of qubit {q}." ) + if q == self._plan.root_qubit: + self.shift_orthogonality_center(self._plan.root) + root_t = self.node_tensor(self._plan.root) + cap_t = qtn.Tensor(vec, inds=(phys,)) + merged = qtn.tensor_contract(root_t.copy(), cap_t) + tags = set(root_t.tags) + tags.discard(self.site_tag(q)) + root_t.modify(data=merged.data, inds=merged.inds, tags=tags) + + old_n = self._plan.n + temp_inds = { + old: f"_ttn_cap_ind_{old}" for old in range(q + 1, old_n) + } + temp_tags = { + old: f"_ttn_cap_tag_{old}" for old in range(q + 1, old_n) + } + if temp_inds: + self.reindex_({ + self.site_ind(old): temp + for old, temp in temp_inds.items() + }) + if temp_tags: + self.retag_({ + self.site_tag(old): temp + for old, temp in temp_tags.items() + }) + if temp_inds: + self.reindex_({ + temp: self.site_ind(old - 1) + for old, temp in temp_inds.items() + }) + if temp_tags: + self.retag_({ + temp: self.site_tag(old - 1) + for old, temp in temp_tags.items() + }) + + self._plan = self._plan.remove_qubit(q) + self._sites = tuple(range(self._plan.n)) + self.__dict__.pop("_node_tid_cache", None) + self._canonical_region = frozenset({self._plan.root}) + self.validate() + return self + + leaf = self._plan.leaf_of_qubit[q] + parent = self._plan.parent.get(leaf) + if parent is None: + raise ValueError("cannot cap the root leaf of a multi-qubit tree.") + # Put the absorbing parent at the centre before contraction, so the # remaining state stays canonical without a full-tree rescan. self.shift_orthogonality_center(parent) @@ -1342,7 +1410,13 @@ def cap_qubit_(self, q, vec): leaf_message = qtn.tensor_contract(leaf_t, cap_t) merged = qtn.tensor_contract(parent_t, leaf_message) - collapse = len(self._plan.children[parent]) == 2 + collapse = ( + len(self._plan.children[parent]) == 2 + and not ( + parent == self._plan.root + and self._plan.root_qubit is not None + ) + ) child = None tags = set(parent_t.tags) if collapse: @@ -1377,7 +1451,7 @@ def cap_qubit_(self, q, vec): if temp_tags: self.retag_({temp: self.site_tag(old - 1) for old, temp in temp_tags.items()}) - self._plan = self._plan.remove_leaf(q) + self._plan = self._plan.remove_qubit(q) self._sites = tuple(range(self._plan.n)) self.__dict__.pop("_node_tid_cache", None) self._canonical_region = frozenset({parent}) @@ -1413,11 +1487,11 @@ def ascii_tree(self, *, bond_dims=True, node_ids=False, color=False): """Return a top-down ASCII drawing of the tree, drawn root-first. The tree analogue of a ``quimb`` MPS ``show``: the root sits at the top - and the qubit leaves at the bottom, each internal node marked ``●`` and - each leaf ``◆`` with its qubit label ``q{q}`` beneath it. When + and structural leaves at the bottom. Physical nodes are labelled + ``q{q}``; the optional root site appears beside the top marker. When ``bond_dims`` is true every edge is annotated with the dimension of the - virtual bond joining a node to its parent (so growing entanglement shows - up as growing numbers on the branches):: + virtual bond joining a node to its parent (so growing entanglement + shows up as growing numbers on the branches):: ● ┌────┴────┐ @@ -1447,7 +1521,7 @@ def ascii_tree(self, *, bond_dims=True, node_ids=False, color=False): def render(nid, depth): """Return ``(lines, root_col, width)`` for the subtree rooted at ``nid``.""" if plan.is_leaf(nid): - label = f"q{plan.qubit_of_leaf[nid]}" + label = f"q{plan.qubit_of_node[nid]}" w = max(1, len(label)) col = (w - 1) // 2 marker = _color("◆", _LEAF_COLOR, color) @@ -1456,6 +1530,8 @@ def render(nid, depth): _ascii_place(lbl, w, col)], col, w dot = f"●{nid}" if node_ids else "●" + if nid in plan.qubit_of_node: + dot += f" q{plan.qubit_of_node[nid]}" marker = _color(dot, _LAYER_COLORS[depth % len(_LAYER_COLORS)], color) blocks = [] for child in plan.children[nid]: @@ -1522,10 +1598,10 @@ def show(self, *, bond_dims=True, node_ids=False, color=True): """Print the top-down ASCII drawing of the tree (see :meth:`ascii_tree`). The tree analogue of a ``quimb`` MPS ``show``: the root is at the top, - the qubit leaves at the bottom, internal nodes are ``●`` and leaves - ``◆`` labelled with their qubit, and every edge is annotated with its - current virtual-bond dimension. By default the markers are coloured by - tree layer; pass ``color=False`` for plain text. + structural leaves are at the bottom, physical nodes are labelled with + their qubits, and every edge is annotated with its current virtual-bond + dimension. By default the markers are coloured by tree layer; pass + ``color=False`` for plain text. """ print(self.ascii_tree(bond_dims=bond_dims, node_ids=node_ids, color=color)) @@ -1546,9 +1622,8 @@ def from_plan(cls, plan, *, dtype=complex, phys_dim=2, site_tag_id="I{}", inds = [] shape = [] tags = [node_tag_id.format(nid)] - leaf = plan.is_leaf(nid) - if leaf: - q = plan.qubit_of_leaf[nid] + q = plan.qubit_of_node.get(nid) + if q is not None: inds.append(site_ind_id.format(q)) shape.append(int(phys_dim)) tags.append(site_tag_id.format(q)) @@ -1559,7 +1634,7 @@ def from_plan(cls, plan, *, dtype=complex, phys_dim=2, site_tag_id="I{}", if up is not None: inds.append(_bond_index(nid, up)) shape.append(1) - if leaf: + if q is not None: data = np.zeros(shape, dtype=dtype) data[tuple([0] * len(shape))] = 1.0 # |0> else: @@ -1594,8 +1669,9 @@ def from_symmray_plan( Symmray supplies the block-sparse / fermionic tensor construction and all subsequent QR/SVD/contraction primitives. This method supplies - only the tree geometry: leaf nodes receive physical legs, while - internal tree nodes remain virtual tensors with neutral charge. + only the tree geometry: leaf nodes and the optional physical root + receive physical legs, while other internal tree nodes remain virtual + tensors with neutral charge. """ try: import symmray as sr @@ -1642,8 +1718,8 @@ def zero_charge(value): dtype=dtype, site_tag_id=node_tag_id, site_charge=lambda nid: ( - leaf_charges[plan.qubit_of_leaf[nid]] - if plan.is_leaf(nid) + leaf_charges[plan.qubit_of_node[nid]] + if nid in plan.qubit_of_node else neutral ), subsizes=subsizes, @@ -1663,8 +1739,8 @@ def zero_charge(value): bond_duals = [] inds = [] - if plan.is_leaf(nid): - q = plan.qubit_of_leaf[nid] + q = plan.qubit_of_node.get(nid) + if q is not None: physical_index = sr.utils.rand_index( symmetry, physical_sectors, @@ -1710,7 +1786,7 @@ def zero_charge(value): def from_order(cls, order, *, weights=None, structure="quality", max_arity=2, community_frac=0.35, star_frac=0.75, dtype=complex, site_tag_id="I{}", site_ind_id="k{}", - node_tag_id="N{}"): + node_tag_id="N{}", root_qubit=None): """Build a product state on a tree partitioned from ``order``. Convenience wrapper that first builds a :class:`TreePlan` with @@ -1721,7 +1797,7 @@ def from_order(cls, order, *, weights=None, structure="quality", plan = TreePlan.from_order( order, weights=weights, structure=structure, max_arity=max_arity, community_frac=community_frac, - star_frac=star_frac, + star_frac=star_frac, root_qubit=root_qubit, ) return cls.from_plan( plan, @@ -1755,9 +1831,8 @@ def _rand(shape): inds = [] shape = [] tags = [node_tag_id.format(nid)] - leaf = plan.is_leaf(nid) - if leaf: - q = plan.qubit_of_leaf[nid] + q = plan.qubit_of_node.get(nid) + if q is not None: inds.append(site_ind_id.format(q)) shape.append(phys_dim) tags.append(site_tag_id.format(q)) diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index da09f6a..42678c8 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -345,6 +345,104 @@ def test_user_supplied_plan_runs(): assert _fidelity(psi, opt.to_dense()) > 1 - 1e-8 +def test_root_physical_qubit_is_first_class_tree_site(): + """A binary top tensor can own one physical qubit alongside two bonds.""" + plan = TreePlan.from_order( + range(4), structure="balanced", root_qubit=4, + ) + state = TreeTensorNetwork.from_plan(plan) + root = state.node_tensor(plan.root) + + assert plan.n == 5 + assert plan.root_qubit == 4 + assert plan.node_of_qubit[4] == plan.root + assert 4 not in plan.leaf_of_qubit + assert set(root.inds) == { + state.site_ind(4), + *(state.bond(plan.root, child) for child in plan.children[plan.root]), + } + assert root.ndim == 3 + assert set(state.outer_inds()) == { + state.site_ind(q) for q in range(plan.n) + } + expected = np.zeros(2**plan.n) + expected[0] = 1.0 + assert np.array_equal(state.to_statevector(), expected) + assert state.validate(check_canonical=True) is state + + +def test_root_physical_qubit_gate_and_submpo_replay_are_exact(): + """Direct gates and a structured sub-MPO can target the top physical leg.""" + plan = TreePlan.from_order( + range(4), structure="balanced", root_qubit=4, + ) + direct_stream = [(pepsy.h(), 0), (pepsy.cnot(), (0, 4))] + direct = TreeOptimizer( + direct_stream, tree=plan, chi=32, cutoff=0.0, + ) + assert _fidelity( + _exact_state(direct_stream, plan.n), direct.to_dense() + ) > 1 - 1e-10 + assert direct.tn.validate(check_canonical=True) is direct.tn + + where = (1, 3, 4) + mpo = _two_branch_flip_submpo( + L=plan.n, + sites=where, + targets=where, + w0=0.0, + w1=1.0, + ) + submpo = TreeOptimizer( + None, tree=plan, chi=32, cutoff=0.0, run=False, + ) + submpo.apply_submpo(mpo, where) + expected = np.zeros(2**plan.n, dtype=complex) + expected[int("01011", 2)] = 1.0 + assert np.allclose(submpo.to_dense(), expected) + assert submpo.tn.validate(check_canonical=True) is submpo.tn + + +def test_root_physical_qubit_layout_and_cap_are_root_aware(): + """Layout scoring reaches the root site and capping removes only its leg.""" + finder = TreeLayoutFinder( + supports=[(0, 4), (0, 4), (1, 2)], + n=5, + root_qubit=4, + structure="balanced", + max_arity=2, + ) + plan = finder.run(refine="greedy", refine_budget=16) + root_path = plan.node_path(plan.node_of_qubit[0], plan.root) + loads = finder.edge_loads(plan) + path_edges = { + (u, v) if plan.parent.get(v) == u else (v, u) + for u, v in zip(root_path, root_path[1:]) + } + + assert plan.root_qubit == 4 + assert plan.node_of_qubit[4] == plan.root + assert all(loads[edge] > 0.0 for edge in path_edges) + assert finder.report(plan)["root_qubit"] == 4 + + automatic = TreeOptimizer( + None, root_qubit=4, max_arity=2, chi=16, run=False, + ) + assert automatic.n == 5 + assert automatic.plan.root_qubit == 4 + + opt = TreeOptimizer(None, tree=plan, chi=16, run=False) + opt.apply_1q(pepsy.h(), 4) + x = np.array([[0.0, 1.0], [1.0, 0.0]]) + assert opt.tn.local_expectation(x, 4) == pytest.approx(1.0) + opt.cap(4, [1.0, 0.0]) + assert opt.plan.root_qubit is None + assert opt.n == 4 + assert opt.to_dense().shape == (2**4,) + assert opt.norm() == pytest.approx(1 / np.sqrt(2)) + assert opt.tn.validate(check_canonical=True) is opt.tn + + def test_layout_finder_builds_valid_tree(): """With max_arity=2 the finder returns a rooted binary tree over all qubits.""" rng = np.random.default_rng(8) @@ -733,6 +831,43 @@ def test_layout_finder_searches_arities_by_default(): assert fixed.run().is_binary() +def test_layout_finder_run_accepts_search_overrides(monkeypatch): + """Tree ``run`` mirrors MPS by accepting per-run quality-search controls.""" + finder = TreeLayoutFinder([], n=4, max_arity=(2, 3), chi=8) + captured = {} + recommend_arities = finder.recommend_arities + + def capture(max_arities, **kwargs): + captured.update(kwargs) + return recommend_arities(max_arities, **kwargs) + + monkeypatch.setattr(finder, "recommend_arities", capture) + plan = finder.run( + chi=None, + refine="greedy", + refine_budget=2, + search=None, + search_budget=7, + seed=11, + nevergrad_optimizer="OnePlusOne", + progbar=True, + ) + + assert isinstance(plan, TreePlan) + assert captured == { + "chi": None, + "refine": "greedy", + "refine_budget": 2, + "search": None, + "search_budget": 7, + "seed": 11, + "nevergrad_optimizer": "OnePlusOne", + "progbar": True, + } + assert finder._last_arity_recommendation["refine"] == "greedy" + assert finder._last_arity_recommendation["chi"] is None + + def test_layout_finder_default_search_is_chi_aware_with_chi(): """A finder built with ``chi`` makes its default arity search chi-aware.""" rng = np.random.default_rng(214) From 6db7efc3c4b01ce0c2c661bdf016653512a48c4b Mon Sep 17 00:00:00 2001 From: rezaquant Date: Tue, 28 Jul 2026 13:52:19 -0600 Subject: [PATCH 09/70] add Torch MC convergence diagnostics --- docs/api/vmc.md | 27 +++ src/pepsy/vmc/__init__.py | 2 + src/pepsy/vmc/api.py | 4 + src/pepsy/vmc/torch/__init__.py | 4 + src/pepsy/vmc/torch/_core.py | 4 + src/pepsy/vmc/torch/driver.py | 343 ++++++++++++++++++++++++++++ src/pepsy/vmc/torch/fermion.py | 45 ++++ src/pepsy/vmc/torch/local_energy.py | 72 ++++-- src/pepsy/vmc/torch/results.py | 107 ++++++++- tests/test_vmc_convergence.py | 61 +++++ 10 files changed, 649 insertions(+), 20 deletions(-) create mode 100644 tests/test_vmc_convergence.py diff --git a/docs/api/vmc.md b/docs/api/vmc.md index b2fe985..2183582 100644 --- a/docs/api/vmc.md +++ b/docs/api/vmc.md @@ -1194,5 +1194,32 @@ callback = pvmc.make_netket_autochunk_callback( driver.run(n_iter=100, out="vmc_run", callback=callback) ``` +### Torch MCMC convergence check + +After a native Torch run, `TorchFermionVMC.check_mc_convergence(...)` runs a +separate, non-mutating diagnostic sampler from the current walker positions. +It retains one local-observable value after every raw Metropolis sweep, then +reports ordinary and split R-hat, average and maximum integrated +autocorrelation time, effective sample size, acceptance, and a suggested +production `sweep_size`: + +```python +report = vmc.check_mc_convergence( + observables={"energy": hamiltonian}, + min_chain_length=100, + max_chain_length=500, + target_effective_samples_per_chain=50, + rhat_threshold=1.05, + progress=True, +) + +print(report.reliable, report.recommended_sweep_size) +print(report.energy.split_r_hat, report.energy.tau_max) +``` + +The live progress display is sampling acceptance only. The check uses a +cloned random stream, leaves the active walker configurations and RNG state +unchanged, and is intentionally separate from fixed-size production sampling. + > API details are maintained as handwritten Markdown in this page. diff --git a/src/pepsy/vmc/__init__.py b/src/pepsy/vmc/__init__.py index babbef6..c509403 100644 --- a/src/pepsy/vmc/__init__.py +++ b/src/pepsy/vmc/__init__.py @@ -41,6 +41,8 @@ "TorchFermionVMC": ".torch", "TorchFermionVMCMetadata": ".torch", "TorchChainDiagnostics": ".torch", + "TorchVMCConvergenceEstimate": ".torch", + "TorchVMCConvergenceReport": ".torch", "TorchImportanceSamples": ".torch", "TorchMCMCSamples": ".torch", "TorchSampleProvenance": ".torch", diff --git a/src/pepsy/vmc/api.py b/src/pepsy/vmc/api.py index e1f5938..497fb20 100644 --- a/src/pepsy/vmc/api.py +++ b/src/pepsy/vmc/api.py @@ -833,6 +833,10 @@ def expect( proposal_log_probs=proposal_log_probs, ) + def check_mc_convergence(self, **kwargs): + """Run the selected backend's explicit post-run mixing diagnostic.""" + return self._setup.check_mc_convergence(**kwargs) + def optimize(self, optimization=None, *, n_steps=None, **kwargs): """Optimize the variational state through the selected backend.""" return self._setup.optimize( diff --git a/src/pepsy/vmc/torch/__init__.py b/src/pepsy/vmc/torch/__init__.py index 5ddda1a..c685261 100644 --- a/src/pepsy/vmc/torch/__init__.py +++ b/src/pepsy/vmc/torch/__init__.py @@ -46,6 +46,8 @@ TorchMCMCSamples, TorchSampleProvenance, TorchChainDiagnostics, + TorchVMCConvergenceEstimate, + TorchVMCConvergenceReport, TorchVMCEnergyEstimate, TorchVMCImportanceEstimate, TorchVMCMeasurementRun, @@ -67,6 +69,8 @@ "TorchMCMCSamples", "TorchSampleProvenance", "TorchChainDiagnostics", + "TorchVMCConvergenceEstimate", + "TorchVMCConvergenceReport", "TorchMetropolisSampler", "TorchBPMetropolisSampler", "TorchVMCDriver", diff --git a/src/pepsy/vmc/torch/_core.py b/src/pepsy/vmc/torch/_core.py index 74f9049..ba96a86 100644 --- a/src/pepsy/vmc/torch/_core.py +++ b/src/pepsy/vmc/torch/_core.py @@ -30,6 +30,8 @@ ) from .results import ( TorchChainDiagnostics, + TorchVMCConvergenceEstimate, + TorchVMCConvergenceReport, TorchImportanceSamples, TorchMCMCSamples, TorchMetropolisResult, @@ -124,6 +126,8 @@ "TorchMCMCSamples", "TorchSampleProvenance", "TorchChainDiagnostics", + "TorchVMCConvergenceEstimate", + "TorchVMCConvergenceReport", "TorchMetropolisSampler", "TorchBPMetropolisSampler", "TorchVMCDriver", diff --git a/src/pepsy/vmc/torch/driver.py b/src/pepsy/vmc/torch/driver.py index ce3d542..3e2bf50 100644 --- a/src/pepsy/vmc/torch/driver.py +++ b/src/pepsy/vmc/torch/driver.py @@ -3,6 +3,7 @@ from __future__ import annotations from dataclasses import replace +import math import time import warnings @@ -26,6 +27,7 @@ _local_energies_from_connection_map, _normalized_sample_weights, _observable_statistics, + torch_chain_diagnostics, _weighted_energy_statistics, local_energy_from_connections, ) @@ -36,6 +38,8 @@ metropolis_exchange_sweep, ) from .results import ( + TorchVMCConvergenceEstimate, + TorchVMCConvergenceReport, TorchVMCEnergyEstimate, TorchVMCImportanceEstimate, TorchVMCStepResult, @@ -737,6 +741,345 @@ def local_observables( compile_kernels=self.compile_kernels, ) + def check_mc_convergence( + self, + observables=None, + *, + burn_in=0, + min_chain_length=100, + max_chain_length=500, + target_effective_samples_per_chain=50.0, + rhat_threshold=1.05, + check_interval=None, + seed=None, + progress=False, + ): + """Check MCMC convergence without advancing the active VMC chains. + + A temporary driver starts from the current walker positions and keeps + one retained value per raw Metropolis sweep. Each observable is + checked with ordinary and split R-hat, an FFT initial-positive- + sequence autocorrelation estimate, and a conservative maximum IAT + across chains. Sampling stops once every observable has at least + ``target_effective_samples_per_chain`` effective samples and satisfies + ``rhat_threshold``. The hard ``max_chain_length`` cap is always + respected. + + The returned :class:`TorchVMCConvergenceReport` contains a suggested + ``sweep_size`` for a later production :class:`SamplingConfig`. The + diagnostic uses raw spacing intentionally, so the recommendation is + not confused with the spacing of an earlier production batch. + """ + torch = _require_torch() + if not hasattr(self, "configs") or not hasattr(self, "amplitudes"): + raise RuntimeError( + "Initialize the Torch VMC driver before checking convergence." + ) + if isinstance(burn_in, bool) or not isinstance(burn_in, int) or burn_in < 0: + raise ValueError("burn_in must be a non-negative integer.") + min_chain_length = _check_positive_int( + "min_chain_length", + min_chain_length, + ) + max_chain_length = _check_positive_int( + "max_chain_length", + max_chain_length, + ) + if max_chain_length < min_chain_length: + raise ValueError( + "max_chain_length must be at least min_chain_length." + ) + try: + target_effective_samples_per_chain = float( + target_effective_samples_per_chain + ) + except (TypeError, ValueError) as exc: + raise ValueError( + "target_effective_samples_per_chain must be positive and finite." + ) from exc + if ( + not math.isfinite(target_effective_samples_per_chain) + or target_effective_samples_per_chain <= 0 + ): + raise ValueError( + "target_effective_samples_per_chain must be positive and finite." + ) + if rhat_threshold is not None: + try: + rhat_threshold = float(rhat_threshold) + except (TypeError, ValueError) as exc: + raise ValueError( + "rhat_threshold must be at least 1 or None." + ) from exc + if not math.isfinite(rhat_threshold) or rhat_threshold < 1.0: + raise ValueError( + "rhat_threshold must be at least 1 or None." + ) + if check_interval is None: + check_interval = max(1, min_chain_length // 10) + check_interval = _check_positive_int("check_interval", check_interval) + + if observables is None: + observable_items = (("energy", None),) + else: + try: + observable_items = tuple(observables.items()) + except AttributeError as exc: + raise TypeError( + "observables must be a mapping of names to native terms." + ) from exc + if not observable_items: + raise ValueError("observables must contain at least one entry.") + observable_items = tuple( + (str(name), terms) for name, terms in observable_items + ) + observable_map = dict(observable_items) + n_chains = self.n_walkers + if n_chains < 2: + raise ValueError( + "check_mc_convergence requires at least two active chains." + ) + model_device = self.configs.device + + def make_generator(): + if seed is None and self.generator is None: + return None + try: + generator = torch.Generator(device=model_device) + except (RuntimeError, TypeError, ValueError): + generator = torch.Generator() + if seed is not None: + generator.manual_seed(int(seed)) + else: + try: + generator.set_state(self.generator.get_state()) + except (AttributeError, RuntimeError, TypeError, ValueError) as exc: + raise RuntimeError( + "Could not clone the active Torch RNG state for the " + "temporary convergence sampler. Pass seed=... instead." + ) from exc + return generator + + if self.terms is not None: + temporary_terms = self.terms + temporary_connection_fn = None + temporary_connection_kwargs = None + else: + temporary_terms = None + temporary_connection_fn = self.connection_fn + temporary_connection_kwargs = self.connection_kwargs + + fork_devices = [] + if getattr(model_device, "type", None) == "cuda": + fork_devices = [ + 0 if model_device.index is None else int(model_device.index) + ] + + histories = {name: [] for name in observable_map} + n_proposed = 0 + n_accepted = 0 + start = time.perf_counter() + bar = _make_progress( + progress, + total=burn_in + max_chain_length, + desc="MC convergence", + unit="sweep", + ) + + def update_progress(result): + nonlocal n_proposed, n_accepted + n_proposed += result.n_proposed + n_accepted += result.n_accepted + if bar is not None: + bar.update(1) + _set_vmc_progress_postfix( + bar, + replace( + result, + n_proposed=n_proposed, + n_accepted=n_accepted, + ), + progress_postfix="acceptance", + ) + + def make_estimates(): + estimates = {} + for name, values in histories.items(): + chain_values = torch.stack(values, dim=0) + diagnostics = torch_chain_diagnostics( + chain_values, + split_rhat=True, + ) + flat_values = chain_values.reshape(-1) + mean, variance = _energy_mean_and_variance(flat_values) + effective_sample_size = diagnostics.effective_sample_size + stderr = torch.sqrt( + variance + / torch.clamp(effective_sample_size, min=1.0) + ) + split_r_hat = ( + diagnostics.split_r_hat + if diagnostics.split_r_hat is not None + else diagnostics.r_hat + ) + max_tau = diagnostics.max_integrated_autocorrelation_time + effective_per_chain = effective_sample_size / n_chains + max_tau_float = float(max_tau.detach().cpu()) + recommended_sweep_size = ( + max(1, int(math.ceil(max_tau_float))) + if math.isfinite(max_tau_float) + else 1 + ) + finite = all( + bool(torch.isfinite(torch.as_tensor(value)).all()) + for value in ( + mean, + variance, + split_r_hat, + diagnostics.integrated_autocorrelation_time, + max_tau, + effective_sample_size, + ) + ) + reasons = [] + if not finite: + reasons.append("non-finite statistics") + if rhat_threshold is not None and ( + not bool(torch.isfinite(split_r_hat)) + or float(split_r_hat.detach().cpu()) > rhat_threshold + ): + reasons.append( + f"split R-hat={float(split_r_hat.detach().cpu()):.4f} " + f"> {rhat_threshold:.4f}" + ) + if ( + not bool(torch.isfinite(effective_per_chain)) + or float(effective_per_chain.detach().cpu()) + < target_effective_samples_per_chain + ): + reasons.append( + "effective samples/chain=" + f"{float(effective_per_chain.detach().cpu()):.1f} " + f"< {target_effective_samples_per_chain:.1f}" + ) + reliable = not reasons + estimates[name] = TorchVMCConvergenceEstimate( + mean=mean, + variance=variance, + stderr=stderr, + r_hat=diagnostics.r_hat, + split_r_hat=split_r_hat, + integrated_autocorrelation_time=( + diagnostics.integrated_autocorrelation_time + ), + max_integrated_autocorrelation_time=max_tau, + effective_sample_size=effective_sample_size, + effective_samples_per_chain=effective_per_chain, + n_samples_per_chain=int(chain_values.shape[0]), + n_chains=n_chains, + recommended_sweep_size=recommended_sweep_size, + reliable=reliable, + reliability_reason=( + "reliable" if reliable else "; ".join(reasons) + ), + ) + return estimates + + rng_context = torch.random.fork_rng( + devices=fork_devices, + enabled=self.generator is None, + ) + try: + with rng_context: + temporary = TorchVMCDriver( + self.model, + self.graph, + self.configs.detach().clone(), + connection_fn=temporary_connection_fn, + terms=temporary_terms, + site_order=self.site_order, + connection_kwargs=temporary_connection_kwargs, + term_constant=self.term_constant, + amplitudes=self.amplitudes.detach().clone(), + proposal=self.proposal, + hopping_rate=self.hopping_rate, + spin_flip_rate=self.spin_flip_rate, + pair_toggle_rate=self.pair_toggle_rate, + encoding=self.encoding, + chunk_size=self.chunk_size, + compile_kernels=self.compile_kernels, + log_amplitude_fn=( + self.log_amplitude_fn + if self.log_amplitude_fn is not None + else False + ), + generator=make_generator(), + ) + for _ in range(burn_in): + update_progress(temporary.sample_sweep()) + + for step in range(1, max_chain_length + 1): + update_progress(temporary.sample_sweep()) + values = temporary.local_observables(observable_map) + for name, value in values.items(): + histories[name].append(value.detach()) + + should_check = ( + step >= min_chain_length + and ( + step == min_chain_length + or (step - min_chain_length) % check_interval == 0 + ) + ) + if should_check: + current = make_estimates() + if all(estimate.reliable for estimate in current.values()): + break + estimates = make_estimates() + finally: + if bar is not None: + bar.close() + + n_steps = len(next(iter(histories.values()))) + reliable = all(estimate.reliable for estimate in estimates.values()) + if reliable: + reliability_reason = ( + "all observables met the split-R-hat and effective-sample " + "targets" + ) + else: + details = " | ".join( + f"{name}: {estimate.reliability_reason}" + for name, estimate in estimates.items() + if not estimate.reliable + ) + reliability_reason = ( + "maximum chain length reached before convergence: " + details + ) + recommended_sweep_size = max( + estimate.recommended_sweep_size for estimate in estimates.values() + ) + elapsed = time.perf_counter() - start + return TorchVMCConvergenceReport( + estimates=estimates, + n_samples_per_chain=n_steps, + n_chains=n_chains, + burn_in=burn_in, + sweep_size=1, + n_sweeps=n_steps, + n_proposed=n_proposed, + n_accepted=n_accepted, + acceptance_rate=(n_accepted / n_proposed if n_proposed else 0.0), + elapsed_seconds=elapsed, + reliable=reliable, + reliability_reason=reliability_reason, + recommended_sweep_size=recommended_sweep_size, + min_chain_length=min_chain_length, + max_chain_length=max_chain_length, + target_effective_samples_per_chain=target_effective_samples_per_chain, + rhat_threshold=rhat_threshold, + ) + def measure_samples( self, samples, diff --git a/src/pepsy/vmc/torch/fermion.py b/src/pepsy/vmc/torch/fermion.py index 5f03fc2..450a3a0 100644 --- a/src/pepsy/vmc/torch/fermion.py +++ b/src/pepsy/vmc/torch/fermion.py @@ -827,6 +827,30 @@ def sample( ) return super().sample(sampling=sampling, **kwargs) + def check_mc_convergence( + self, + observables=None, + *, + contraction=None, + contraction_opts=None, + **kwargs, + ): + """Check energy/observable chain mixing without mutating VMC state. + + The fermionic wrapper compiles the requested native observable map and + then delegates to :meth:`TorchVMCDriver.check_mc_convergence`, which + runs a temporary raw-sweep sampler from the current walker positions. + """ + self._ensure_initialized( + contraction=contraction, + contraction_opts=contraction_opts, + ) + compiled = self._measurement_observables( + observables, + include_energy=True, + ) + return super().check_mc_convergence(compiled, **kwargs) + def measure( self, samples, @@ -1154,6 +1178,27 @@ def sample(self, sampling=None): ) return native.to_common() + def check_mc_convergence( + self, + *, + sampling=None, + contraction=None, + contraction_opts=None, + **kwargs, + ): + """Run the native non-mutating convergence diagnostic.""" + sampling = self.sampling if sampling is None else sampling + self.driver._ensure_initialized( + sampling=sampling, + contraction=contraction, + contraction_opts=contraction_opts, + ) + return self.driver.check_mc_convergence( + contraction=contraction, + contraction_opts=contraction_opts, + **kwargs, + ) + def _measurement_terms(self, observables): if observables is None: return dict(self.driver.observables) diff --git a/src/pepsy/vmc/torch/local_energy.py b/src/pepsy/vmc/torch/local_energy.py index 0b9fc3b..79188b8 100644 --- a/src/pepsy/vmc/torch/local_energy.py +++ b/src/pepsy/vmc/torch/local_energy.py @@ -704,14 +704,45 @@ def _weighted_energy_statistics(local_energies, weights): ) -def torch_chain_diagnostics(values, *, max_lag=None): +def _gelman_r_hat(values): + """Return the ordinary Gelman--Rubin R-hat for chain-shaped values.""" + torch = _require_torch() + n_steps, _ = (int(value) for value in values.shape) + chain_means = values.mean(dim=0) + within = values.var(dim=0, unbiased=True).mean() + between = n_steps * chain_means.var(unbiased=True) + variance_hat = ((n_steps - 1) * within + between) / n_steps + if bool(within == 0): + return torch.where( + between == 0, + torch.ones_like(variance_hat), + torch.full_like(variance_hat, float("inf")), + ) + return torch.sqrt(torch.clamp(variance_hat / within, min=1.0)) + + +def _split_chain_values(values): + """Split each chain in half for the more sensitive split R-hat.""" + n_steps = int(values.shape[0]) + if n_steps < 4: + return None + half = n_steps // 2 + return _require_torch().cat( + (values[:half], values[-half:]), + dim=1, + ) + + +def torch_chain_diagnostics(values, *, max_lag=None, split_rhat=False): """Return ``R-hat``, integrated autocorrelation time, and ESS. ``values`` must have shape ``(n_samples_per_chain, n_chains)``. The - implementation uses split-chain-independent Gelman--Rubin statistics and - an FFT autocorrelation estimate with an initial-positive-sequence cutoff. - Complex values are reduced to their real parts, as appropriate for a - Hermitian local observable. + The default R-hat is the ordinary Gelman--Rubin estimate used by the + legacy adaptive measurement loop. Set ``split_rhat=True`` to also return + a split-chain R-hat, which is more sensitive to non-stationary chains. + Autocorrelation uses an FFT estimate with an initial-positive-sequence + cutoff. Complex values are reduced to their real parts, as appropriate + for a Hermitian local observable. """ torch = _require_torch() values = torch.as_tensor(values) @@ -731,19 +762,11 @@ def torch_chain_diagnostics(values, *, max_lag=None): values = values.to(torch.float64) chain_means = values.mean(dim=0) - within = values.var(dim=0, unbiased=True).mean() - between = n_steps * chain_means.var(unbiased=True) - variance_hat = ( - (n_steps - 1) * within + between - ) / n_steps - if bool(within == 0): - r_hat = torch.where( - between == 0, - torch.ones_like(variance_hat), - torch.full_like(variance_hat, float("inf")), - ) - else: - r_hat = torch.sqrt(torch.clamp(variance_hat / within, min=1.0)) + r_hat = _gelman_r_hat(values) + split_values = _split_chain_values(values) if split_rhat else None + split_r_hat = ( + None if split_values is None else _gelman_r_hat(split_values) + ) if max_lag is None: max_lag = n_steps - 1 @@ -772,6 +795,17 @@ def torch_chain_diagnostics(values, *, max_lag=None): break tau = tau + 2 * rho[lag] tau = torch.clamp(tau, min=1.0) + chain_taus = torch.ones( + n_chains, + dtype=values.dtype, + device=values.device, + ) + for chain in range(n_chains): + for lag in range(1, max_lag + 1): + if bool(normalized[lag, chain] <= 0): + break + chain_taus[chain] = chain_taus[chain] + 2 * normalized[lag, chain] + chain_taus = torch.clamp(chain_taus, min=1.0) total_samples = n_steps * n_chains effective_sample_size = torch.as_tensor( total_samples, @@ -784,6 +818,8 @@ def torch_chain_diagnostics(values, *, max_lag=None): effective_sample_size=effective_sample_size, n_samples_per_chain=n_steps, n_chains=n_chains, + split_r_hat=split_r_hat, + max_integrated_autocorrelation_time=chain_taus.max(), ) diff --git a/src/pepsy/vmc/torch/results.py b/src/pepsy/vmc/torch/results.py index 81f4f9d..b5ccb00 100644 --- a/src/pepsy/vmc/torch/results.py +++ b/src/pepsy/vmc/torch/results.py @@ -118,7 +118,7 @@ class TorchMCMCSamples: proposal_stats: Any = None provenance: TorchSampleProvenance | None = None - def diagnostics(self, values=None, *, max_lag=None): + def diagnostics(self, values=None, *, max_lag=None, split_rhat=False): """Compute chain diagnostics for a scalar observable. If ``values`` is omitted, the sampled ``|psi|**2`` values are used as @@ -129,7 +129,11 @@ def diagnostics(self, values=None, *, max_lag=None): values = self.amplitudes.abs().square() from ._core import torch_chain_diagnostics - return torch_chain_diagnostics(values, max_lag=max_lag) + return torch_chain_diagnostics( + values, + max_lag=max_lag, + split_rhat=split_rhat, + ) def to_common(self): """Convert to the backend-neutral :class:`pepsy.vmc.VMCSamples`.""" @@ -204,6 +208,8 @@ class TorchChainDiagnostics: effective_sample_size: Any n_samples_per_chain: int n_chains: int + split_r_hat: Any = None + max_integrated_autocorrelation_time: Any = None @property def rhat(self): @@ -215,6 +221,101 @@ def tau(self): """Alias for :attr:`integrated_autocorrelation_time`.""" return self.integrated_autocorrelation_time + @property + def rhat_split(self): + """Alias for the optional split-chain R-hat estimate.""" + return self.split_r_hat + + @property + def tau_max(self): + """Alias for the optional maximum per-chain autocorrelation time.""" + return self.max_integrated_autocorrelation_time + + +@dataclass(frozen=True) +class TorchVMCConvergenceEstimate: + """Convergence summary for one observable measured along raw sweeps.""" + + mean: Any + variance: Any + stderr: Any + r_hat: Any + split_r_hat: Any + integrated_autocorrelation_time: Any + max_integrated_autocorrelation_time: Any + effective_sample_size: Any + effective_samples_per_chain: Any + n_samples_per_chain: int + n_chains: int + recommended_sweep_size: int + reliable: bool + reliability_reason: str + + @property + def rhat(self): + """Alias for :attr:`r_hat`.""" + return self.r_hat + + @property + def tau(self): + """Alias for :attr:`integrated_autocorrelation_time`.""" + return self.integrated_autorrelation_time + + @property + def tau_max(self): + """Alias for :attr:`max_integrated_autocorrelation_time`.""" + return self.max_integrated_autocorrelation_time + + +@dataclass(frozen=True) +class TorchVMCConvergenceReport: + """Non-mutating post-run MCMC convergence report. + + The report is produced from a temporary sampler whose retained samples + are separated by one raw Metropolis sweep. ``recommended_sweep_size`` + is therefore expressed in the same sweep units accepted by + :class:`SamplingConfig`. + """ + + estimates: Mapping[str, TorchVMCConvergenceEstimate] + n_samples_per_chain: int + n_chains: int + burn_in: int + sweep_size: int + n_sweeps: int + n_proposed: int + n_accepted: int + acceptance_rate: float + elapsed_seconds: float + reliable: bool + reliability_reason: str + recommended_sweep_size: int + min_chain_length: int + max_chain_length: int + target_effective_samples_per_chain: float + rhat_threshold: float | None + + def __post_init__(self): + object.__setattr__( + self, + "estimates", + MappingProxyType(dict(self.estimates)), + ) + + @property + def energy(self): + """Return the energy estimate when it was included in the report.""" + return self.estimates.get("energy") + + @property + def n_samples(self): + """Total number of retained scalar samples across all chains.""" + return self.n_samples_per_chain * self.n_chains + + def __getitem__(self, name): + """Return one named observable estimate.""" + return self.estimates[name] + @dataclass(frozen=True) class TorchVMCStepResult: @@ -543,6 +644,8 @@ def _accumulate_cache_profile(total, snapshot): "TorchMCMCSamples", "TorchMetropolisResult", "TorchSampleProvenance", + "TorchVMCConvergenceEstimate", + "TorchVMCConvergenceReport", "TorchVMCImportanceEstimate", "TorchVMCEnergyEstimate", "TorchVMCStepResult", diff --git a/tests/test_vmc_convergence.py b/tests/test_vmc_convergence.py new file mode 100644 index 0000000..1745f63 --- /dev/null +++ b/tests/test_vmc_convergence.py @@ -0,0 +1,61 @@ +"""Regression tests for native Torch VMC convergence diagnostics.""" + +import pytest + + +def test_chain_diagnostics_can_return_split_rhat_and_max_tau(): + torch = pytest.importorskip("torch") + from pepsy.vmc.torch import torch_chain_diagnostics + + values = torch.cat( + ( + torch.zeros((8, 2), dtype=torch.float64), + torch.full((8, 2), 10.0, dtype=torch.float64), + ), + dim=0, + ) + diagnostics = torch_chain_diagnostics(values, split_rhat=True) + + assert diagnostics.split_r_hat is not None + assert torch.isfinite(diagnostics.max_integrated_autocorrelation_time) + assert diagnostics.max_integrated_autocorrelation_time >= 1.0 + + +def test_convergence_check_uses_temporary_chains_and_rng_state(): + torch = pytest.importorskip("torch") + from pepsy.vmc import TorchVMCDriver + + class ProductAmplitude(torch.nn.Module): + def __init__(self): + super().__init__() + self.weights = torch.nn.Parameter( + torch.tensor([1.0, 2.0], dtype=torch.float64) + ) + + def forward(self, configs): + return self.weights[configs].prod(dim=1) + + generator = torch.Generator().manual_seed(23) + driver = TorchVMCDriver( + ProductAmplitude(), + [(0, 1)], + torch.tensor([[0, 1], [1, 0]], dtype=torch.long), + terms={0: torch.tensor([[0.0, 1.0], [1.0, 0.0]])}, + proposal="spin", + generator=generator, + ) + configs_before = driver.configs.clone() + rng_before = driver.generator.get_state().clone() + + report = driver.check_mc_convergence( + min_chain_length=4, + max_chain_length=6, + target_effective_samples_per_chain=1.0, + seed=7, + ) + + assert report.n_samples_per_chain >= 4 + assert report.energy is not None + assert report.energy.split_r_hat is not None + assert torch.equal(driver.configs, configs_before) + assert torch.equal(driver.generator.get_state(), rng_before) From d0f0c7ed356fec9a32d26d37b487b891d61567ba Mon Sep 17 00:00:00 2001 From: rezaquant Date: Tue, 28 Jul 2026 14:01:14 -0700 Subject: [PATCH 10/70] tree: support physical root across APIs --- .github/skills/tree-optimizer/SKILL.md | 13 +- .../references/performance-layout.md | 7 +- docs/api/optimizers/tree.md | 27 ++-- src/pepsy/__init__.py | 3 +- src/pepsy/optimizers/__init__.py | 1 + src/pepsy/optimizers/tree/layout.py | 20 ++- src/pepsy/optimizers/tree/optimizer.py | 6 + src/pepsy/sampling/tree.py | 138 +++++++++++------- src/pepsy/tensors/constructors.py | 69 +++++++-- tests/test_optimize_tree.py | 81 ++++++++++ tests/test_public_api.py | 5 +- tests/test_tree_sampler.py | 42 +++++- 12 files changed, 317 insertions(+), 95 deletions(-) diff --git a/.github/skills/tree-optimizer/SKILL.md b/.github/skills/tree-optimizer/SKILL.md index 5e9f71c..48d3bb7 100644 --- a/.github/skills/tree-optimizer/SKILL.md +++ b/.github/skills/tree-optimizer/SKILL.md @@ -196,7 +196,7 @@ one-node case. connected subtree via `ttn.subtree_span(nodes)` (union of tree paths from `nodes[0]`; generalises `steiner_nodes` to arbitrary internal nodes). - `ttn.canonize_around_qubits_(qubits)` is the qubit-level "range" entry point = - `canonize_subtree_(leaves_of(qubits), span=True)`. + `canonize_subtree_(nodes_of(qubits), span=True)`. - `ttn.is_subtree_canonical_form(nodes=None, span=False)` verifies every outside tensor is an inward isometry (defaults to the tracked region); `is_canonical_form` is its one-node case and delegates to it. @@ -211,11 +211,11 @@ This is the paper's accuracy point (Figs. 3-6) -- do not regress it. 1. SVD-split the gate into left/right factors joined by a virtual bond (`cutoff=0.0`, exact rank `k <= 4`). -2. Move the centre to leaf `a`, absorb the left factor into `a`. -3. Thread the virtual bond **exactly** along the geodesic to leaf `b` via +2. Move the centre to physical node `a`, absorb the left factor into `a`. +3. Thread the virtual bond **exactly** along the geodesic to physical node `b` via `_thread_hop` (economical **QR**, lossless, `absorb="right"`); the crossed bond grows transiently by at most `k <= 4`. -4. Absorb the right factor into leaf `b`. +4. Absorb the right factor into physical node `b`. 5. Only now run `_compress_path` -- a single canonical compression sweep back along the geodesic, truncating every touched bond to `chi`. @@ -302,7 +302,7 @@ interface use the dense `to_dense()` fallback and remain subject to when the gauge is unknown; native readout leaves the gauge untouched. Normalized native readout reuses a state-versioned norm denominator until a mutation invalidates it. -- `measure(q, outcome=None)`: move centre to the leaf, read exact Born +- `measure(q, outcome=None)`: move centre to the physical node, read exact Born probabilities from that one tensor (`w_i = sum_bond |t[i,bond]|^2`, normalise), sample via `self.rng.choice` or force `outcome`, project with a one-hot `apply_1q`, then `normalize()`. Returns the outcome bit. `reset(q)` = @@ -321,6 +321,9 @@ interface use the dense `to_dense()` fallback and remain subject to branch norm, and both can return support/span/bond/norm diagnostics. - `to_dense()` returns a host NumPy statevector in `k0, k1, ..., k(n-1)` order; it is a readout boundary, not evidence that a Torch/CuPy live state moved. +- `ps_to_ttn`, `hrs_to_ttn`, and `TreeSampler` resolve physical sites through + `node_of_qubit`, so an optional root site is constructed and sampled in the + same `q0..q(n-1)` order as leaf sites. - `run(progbar=True)` shows a tqdm replay bar with one-/two-/multi-qubit counts, current bond usage, norm, and a norm-based truncation proxy. Dense and native fermionic replay use the same `1 - (norm / reference_norm)^2` proxy; diff --git a/.github/skills/tree-optimizer/references/performance-layout.md b/.github/skills/tree-optimizer/references/performance-layout.md index ea47fed..992ce45 100644 --- a/.github/skills/tree-optimizer/references/performance-layout.md +++ b/.github/skills/tree-optimizer/references/performance-layout.md @@ -5,9 +5,10 @@ Tree Optimizer skill so the upload-facing `SKILL.md` stays concise. ## Performance and stability -- **BLAS thread cap is the biggest performance lever.** Tree tensors are tiny - (rank `<= 3`, bounded by `chi`), so multi-threaded BLAS/OpenMP is dominated - by thread launch/sync overhead. `threads=1` is the default; gate +- **BLAS thread cap is the biggest performance lever.** Tree tensors are + moderate-rank (set by local arity and an optional root physical leg, with + dimensions bounded by `chi`), so multi-threaded BLAS/OpenMP is dominated by + thread launch/sync overhead. `threads=1` is the default; gate application and heavy readouts run inside `self._thread_ctx()` using `threadpoolctl` when available. Only raise `threads` in a large-`chi` regime. - The self-healing tid cache (`_nid_to_tid`, `_tid`) validates cached tensor diff --git a/docs/api/optimizers/tree.md b/docs/api/optimizers/tree.md index dc15715..bd75c20 100644 --- a/docs/api/optimizers/tree.md +++ b/docs/api/optimizers/tree.md @@ -306,13 +306,15 @@ readout, `pepsy.TreeEnergyOptimizer` wraps this batch path and returns an For the package-level product-state constructor, matching `ps_to_mps`, use `pepsy.ps_to_ttn(n, theta=..., tree=...)`. It builds the requested tree, -initialises every leaf with `[cos(theta), sin(theta)]`, and optionally expands -the virtual bonds with `chi`. +initialises every physical site with `[cos(theta), sin(theta)]`, and optionally +expands the virtual bonds with `chi`. Pass `root_qubit=q` to build the plan +directly, or supply a matching root-site `TreePlan` through `tree=`. For a native Symmray fermionic state, pass a `Fermion` model and occupations: `pepsy.ps_to_ttn(n, tree=plan, fermion=fermion, occupations=..., chi=1)`. -Leaves then carry the model's physical charge/parity sectors, internal nodes -are neutral, and every tree edge uses conjugate Symmray virtual indices. +Physical sites then carry the model's charge/parity sectors, virtual-only +internal nodes are neutral, and every tree edge uses conjugate Symmray virtual +indices. The constructor selects a definite local Fock basis vector, not a random vector inside a degenerate charge sector. For spinful `U1`/`Z2`, a scalar occupation `1` selects the checkerboard `|up>, |down>, ...` representative; pass @@ -320,8 +322,13 @@ inside a degenerate charge sector. For spinful `U1`/`Z2`, a scalar occupation graded product tree is normalized by an exact graded norm contraction, so its represented norm is one rather than an arbitrary constructor scalar. `pepsy.hrs_to_ttn(..., chi=...)` creates the corresponding random symmetric -tree with the requested charge-sector bond dimension. These constructors keep -the Symmray arrays native; they do not materialize dense tensor data. +tree with the requested charge-sector bond dimension and accepts the same +`root_qubit=` option. These constructors keep the Symmray arrays native; they +do not materialize dense tensor data. + +`pepsy.TreeSampler(state)` samples every registered physical site, including +the optional root site. Its cached canonical arrays use parent, physical, then +child axes, so probabilities and amplitudes retain normal `q0..q(n-1)` order. `TreeTensorNetwork.show()` prints a top-down ASCII drawing of the tree -- the tree analogue of a quimb MPS `show()` -- with the root at the top, structural @@ -384,7 +391,8 @@ any arity, controlled by two knobs on `TreeLayoutFinder` / `TreePlan.from_order` bisection. Binary trees remain a valid special case (`max_arity=2`). A caller may bypass the finder entirely by passing an explicit `TreePlan` via -`TreeOptimizer(..., tree=plan)`. Build one with +`TreeOptimizer(..., tree=plan)`. `TreePlan` is exported from both `pepsy` and +`pepsy.optimizers.tree`. Build one with `TreePlan.from_order(order, weights=..., structure=..., max_arity=...)`, or -- for a fully hand-specified arbitrary-arity tree -- with `TreePlan.from_children(children, qubit_of_leaf)`, which validates that the @@ -664,7 +672,7 @@ For stream control events, `TreeOptimizer.measure_event`, `cap_event`, `reset_event`, and `measure_reset_event` build the same tuple forms as `MpsOptimizer`, including Pauli-basis measurement and reset. Their recorded results are `(pauli, where, outcome, probability)` in `measurements`. -`cap(q, vec)` contracts and removes one leaf, shifting the remaining labels +`cap(q, vec)` contracts and removes one physical site, shifting the remaining labels above `q` down by one unless stable labels are requested. `normalize()` rescales the represented state to unit norm and `max_bond()` reports the largest virtual bond. Truncation details are available through `truncation_report()`, `get_infidelities()`, and @@ -679,7 +687,8 @@ available through `truncation_report()`, `get_infidelities()`, and QR bond-threading and double-bond fusion of the general geodesic route and is the common case in a locality-aware layout. -- **Thread cap.** Tree tensors are small (rank `<= 3`, bounded by `chi`), so +- **Thread cap.** Tree tensors are moderate-rank (set by local arity and the + optional root physical leg, with dimensions bounded by `chi`), so multi-threaded BLAS/OpenMP linear algebra is dominated by thread launch and synchronisation overhead. `TreeOptimizer` caps threads to `1` around gate application and the heavy read-outs by default (`threads=1`), which makes diff --git a/src/pepsy/__init__.py b/src/pepsy/__init__.py index 22609e4..7d8a48f 100644 --- a/src/pepsy/__init__.py +++ b/src/pepsy/__init__.py @@ -147,6 +147,7 @@ "TreeEnergyOptimizer": ".optimizers", "TreeLayoutFinder": ".optimizers", "TreeOptimizer": ".optimizers", + "TreePlan": ".optimizers", "TreeStabOptimizer": ".optimizers", "TreeTensorNetwork": ".optimizers", "compile_stim_circuit": ".optimizers", @@ -346,7 +347,7 @@ def __getattr__(name): y, z, ) - from .optimizers import DeferredInjectionRecord, DeferredInjectionReport, DeferredProjectionRecord, GlobalOptimizer, ImmediateInjectionReport, ImmediateProjectionRecord, MeasurementRecord, MpoOptimizer, MpsEnergyOptimizer, MpsOptimizer, MpsStabOptimizer, NormEventRecord, PepsEnergyOptimizer, PepsOptimizer, STNState, StabilizerMpsSettingsAdvice, StabilizerMpsRunResult, StabilizerMpsSimulator, StabilizerTreeRunResult, StreamAnalysisRecord, SimpleUpdateGen, SymDMRG2, SweepOptimizer, TreeEnergyOptimizer, TreeLayoutFinder, TreeOptimizer, TreeStabOptimizer, run_stabilizer_mps_stream, run_stabilizer_tree_stream # noqa: F401 + from .optimizers import DeferredInjectionRecord, DeferredInjectionReport, DeferredProjectionRecord, GlobalOptimizer, ImmediateInjectionReport, ImmediateProjectionRecord, MeasurementRecord, MpoOptimizer, MpsEnergyOptimizer, MpsOptimizer, MpsStabOptimizer, NormEventRecord, PepsEnergyOptimizer, PepsOptimizer, STNState, StabilizerMpsSettingsAdvice, StabilizerMpsRunResult, StabilizerMpsSimulator, StabilizerTreeRunResult, StreamAnalysisRecord, SimpleUpdateGen, SymDMRG2, SweepOptimizer, TreeEnergyOptimizer, TreeLayoutFinder, TreeOptimizer, TreePlan, TreeStabOptimizer, run_stabilizer_mps_stream, run_stabilizer_tree_stream # noqa: F401 from .sampling import FermionConfigurationEncoding, MpsDiagonalEstimate, MpsBatchSampleResult, MpsSampleResult, MpsSampler, PEPSSampleResult, PepsBpSampler, TreeBatchSampleResult, TreeSampleResult, TreeSampler, VecSampler # noqa: F401 from .solvers import FDSolver # noqa: F401 from .tensors import ( # noqa: F401 diff --git a/src/pepsy/optimizers/__init__.py b/src/pepsy/optimizers/__init__.py index d451239..43e20cd 100644 --- a/src/pepsy/optimizers/__init__.py +++ b/src/pepsy/optimizers/__init__.py @@ -26,6 +26,7 @@ "SweepOptimizer": ".sweep", "TreeLayoutFinder": ".tree", "TreeOptimizer": ".tree", + "TreePlan": ".tree", "TreeStabOptimizer": ".tree_stabilizer", "TreeTensorNetwork": ".tree", "CoalescedMeasurementRecord": ".noise", diff --git a/src/pepsy/optimizers/tree/layout.py b/src/pepsy/optimizers/tree/layout.py index a2e4981..0cc0435 100644 --- a/src/pepsy/optimizers/tree/layout.py +++ b/src/pepsy/optimizers/tree/layout.py @@ -337,6 +337,11 @@ def __init__( self.root_qubit = ( None if root_qubit is None else int(root_qubit) ) + if self.root_qubit is not None and self.root in self.qubit_of_leaf: + raise ValueError( + "the root cannot carry both a leaf qubit and root_qubit; " + "insert a unary structural root above the leaf." + ) self.qubit_of_node = dict(self.qubit_of_leaf) if self.root_qubit is not None: self.qubit_of_node[self.root] = self.root_qubit @@ -560,6 +565,12 @@ def build(qs): if order: root = build(order) + if root_qubit is not None and root in qubit_of_leaf: + # With one non-root qubit, ``build`` returns that physical + # leaf itself. The top qubit needs its own tensor, so insert a + # unary structural root rather than putting two physical legs + # on the same node. + root = make_internal((root,)) else: root = new_node() children[root] = () @@ -789,10 +800,17 @@ def binary_subtree(nodes): # The blocking node is already a valid root. In particular, do # not add a unary wrapper for n=1 or n <= block_size: it adds a # useless bond and makes the fixed layered family less efficient. + # The exception is a physical root over one physical leaf: those + # two qubits require distinct tensors joined by one bond. + root_nid = block_nodes[0] + if root_qubit is not None and not children_map[root_nid]: + child = root_nid + root_nid = new_node() + children_map[root_nid] = (child,) return cls.from_children( children_map, qubit_of_leaf, - root=block_nodes[0], + root=root_nid, root_qubit=root_qubit, ) root_arity = min(cls.LAYERED_ROOT_ARITY, num_blocks) diff --git a/src/pepsy/optimizers/tree/optimizer.py b/src/pepsy/optimizers/tree/optimizer.py index efbad09..edb260f 100644 --- a/src/pepsy/optimizers/tree/optimizer.py +++ b/src/pepsy/optimizers/tree/optimizer.py @@ -617,6 +617,10 @@ def __init__(self, gates=None, n=None, *, chi=64, cutoff=1e-12, tree = self.layout_finder.run() if not isinstance(tree, TreePlan): raise TypeError("tree must be a TreePlan or None.") + if tree.n != self.n: + raise ValueError( + f"tree contains {tree.n} qubits, but n={self.n} was requested." + ) if root_qubit is not None and tree.root_qubit != root_qubit: raise ValueError( "root_qubit does not match the supplied tree/layout plan." @@ -1400,6 +1404,8 @@ def layout_report(self): if self.layout_finder is None: return { "n_qubits": self.n, + "root": self.plan.root, + "root_qubit": self.plan.root_qubit, "is_binary": self.plan.is_binary(), "max_arity": self.plan.max_arity(), } diff --git a/src/pepsy/sampling/tree.py b/src/pepsy/sampling/tree.py index 8b97bf7..36da03e 100644 --- a/src/pepsy/sampling/tree.py +++ b/src/pepsy/sampling/tree.py @@ -92,7 +92,7 @@ class TreeSampleResult: probs : list[float] Born probability ``||**2`` for each sample. nqubits : int - Number of qubit leaves in the tree. + Number of physical qubit sites in the tree. """ configs: list[list[int]] @@ -138,7 +138,7 @@ class TreeBatchSampleResult: probs : np.ndarray Born probabilities for ``configs`` with shape ``(n_samples,)``. nqubits : int - Number of qubit leaves in the tree. + Number of physical qubit sites in the tree. """ configs: np.ndarray @@ -237,7 +237,7 @@ def __init__(self, state, *, seed=None, threads: int | None = 1, fermion=None): self._nqubits = None self._root = None self._children = None - self._qubit_of_leaf = None + self._qubit_of_node = None self._arrays = None self._fermionic = False self._configuration_encoding = None @@ -292,11 +292,13 @@ def refresh(self, state=None): def _extract_arrays(self, tn): """Extract per-node dense arrays with a canonical axis order. - Leaf arrays have axes ``(parent_bond, phys)``; internal arrays have - axes ``(parent_bond, child_0, child_1, ...)``. The root is given a - dummy parent bond of size 1 so the sampling recursion is uniform. The - root array is L2-normalized in place, which normalizes the whole state - because every other tensor is isometric in this canonical form. + Every array starts with its parent bond, with the root receiving a + dummy parent bond of size 1. A physical site's axis comes next, followed + by its child bonds. Thus a leaf has ``(parent_bond, phys)``, a + virtual-only internal node has ``(parent_bond, child_0, ...)``, and a + physical root has ``(dummy_parent, phys, child_0, ...)``. The root array + is L2-normalized in place, which normalizes the whole state because + every other tensor is isometric in this canonical form. Bond indices are resolved by intersecting adjacent node tensors rather than by the deterministic ``_tb{lo}_{hi}`` names, because gate threading @@ -305,7 +307,7 @@ def _extract_arrays(self, tn): plan = tn.plan root = plan.root children = {nid: tuple(plan.children[nid]) for nid in plan.children} - qubit_of_leaf = dict(plan.qubit_of_leaf) + qubit_of_node = dict(plan.qubit_of_node) # Native Symmray fermionic trees keep block-sparse arrays; densify them # once. A graded-canonical tensor is also plain-isometric (its exchange @@ -333,21 +335,16 @@ def bond_between(a, b): t = tn.node_tensor(nid) is_root = nid == root parent = plan.parent.get(nid) - if not ch: # leaf - phys = tn.site_ind(qubit_of_leaf[nid]) - if is_root: # single-qubit tree - arr = to_arr(t.transpose(phys).data).reshape(1, -1) - else: - pbond = bond_between(nid, parent) - arr = to_arr(t.transpose(pbond, phys).data) - else: # internal node - cbonds = [bond_between(nid, c) for c in ch] - if is_root: - arr = to_arr(t.transpose(*cbonds).data) - arr = arr.reshape((1,) + arr.shape) - else: - pbond = bond_between(nid, parent) - arr = to_arr(t.transpose(pbond, *cbonds).data) + ordered_inds = [] + if not is_root: + ordered_inds.append(bond_between(nid, parent)) + q = qubit_of_node.get(nid) + if q is not None: + ordered_inds.append(tn.site_ind(q)) + ordered_inds.extend(bond_between(nid, child) for child in ch) + arr = to_arr(t.transpose(*ordered_inds).data) + if is_root: + arr = arr.reshape((1,) + arr.shape) arrays[nid] = arr # Normalize via the root array (state is canonical with centre = root). @@ -359,20 +356,20 @@ def bond_between(a, b): self._nqubits = int(plan.n) self._root = root self._children = children - self._qubit_of_leaf = qubit_of_leaf + self._qubit_of_node = qubit_of_node self._arrays = arrays self._fermionic = fermionic self._configuration_encoding = ( - self._build_configuration_encoding(tn, arrays, qubit_of_leaf) + self._build_configuration_encoding(tn, arrays, qubit_of_node) if fermionic else None ) - def _build_configuration_encoding(self, tn, arrays, qubit_of_leaf): + def _build_configuration_encoding(self, tn, arrays, qubit_of_node): """Build the dense-basis occupation decoder for a fermionic tree.""" - if not qubit_of_leaf: + if not qubit_of_node: return None - phys_dim = int(arrays[next(iter(qubit_of_leaf))].shape[-1]) + phys_dim = int(arrays[next(iter(qubit_of_node))].shape[1]) fermion = self._fermion if fermion is not None and hasattr(fermion, "spinful"): spinful = bool(fermion.spinful) @@ -393,7 +390,7 @@ def _build_configuration_encoding(self, tn, arrays, qubit_of_leaf): @property def nqubits(self) -> int: - """Number of qubit leaves in the tree.""" + """Number of physical qubit sites in the tree.""" return self._nqubits # -- sampling ------------------------------------------------------------ @@ -403,7 +400,7 @@ def _sample_arrays(self, n_samples, rng): B = int(n_samples) arrays = self._arrays children = self._children - qubit_of_leaf = self._qubit_of_leaf + qubit_of_node = self._qubit_of_node configs = np.zeros((B, self._nqubits), dtype=np.int64) prob = np.ones(B, dtype=np.float64) batch = np.arange(B) @@ -412,10 +409,13 @@ def visit(nid, rho): # rho: (B, d_par, d_par) reduced density on nid's parent bond. ch = children[nid] arr = arrays[nid] - if not ch: # leaf + q = qubit_of_node.get(nid) + if q is not None: # p[B, x] = Re sum_{a,a'} rho[a,a'] T[a,x] conj(T[a',x]). - tmp = np.einsum("BaA,ax->BAx", rho, arr) - p = np.einsum("BAx,Ax->Bx", tmp, arr.conj()).real + flat = arr.reshape(arr.shape[0], arr.shape[1], -1) + p = np.einsum( + "BaA,axF,AxF->Bx", rho, flat, flat.conj() + ).real p = np.clip(p, 0.0, None) total = p.sum(axis=1, keepdims=True) probs = p / np.where(total > 0.0, total, 1.0) @@ -423,22 +423,39 @@ def visit(nid, rho): cdf = np.cumsum(probs, axis=1) x = (draws[:, None] > cdf).sum(axis=1) x = np.minimum(x, probs.shape[1] - 1).astype(np.int64) - configs[:, qubit_of_leaf[nid]] = x + configs[:, q] = x prob[:] *= probs[batch, x] - # phi[B, a] = T[a, x] -- the collapsed subtree amplitude. - return arr[:, x].T - - par = arr.shape[0] - # First child: trace the (unbatched) future siblings to identity. - d0 = arr.shape[1] - F0 = int(np.prod(arr.shape[2:])) if len(ch) > 1 else 1 - ur = arr.reshape(par, d0, F0) - env = np.einsum("acF,AdF->acAd", ur, ur.conj()) - rho0 = np.einsum("BaA,acAd->Bcd", rho, env) - phi0 = visit(ch[0], rho0) - # Collapse child 0 into the node tensor -> batched remainder. - K = np.tensordot(phi0, arr, axes=([1], [1])) # (B, par, d1, ...) - for i in range(1, len(ch)): + # Selecting one physical value leaves a batched tensor over + # the parent and child bonds. A physical leaf has no remaining + # child axes and can return immediately. + selected = np.moveaxis(arr[:, x, ...], 1, 0) + if not ch: + return selected.reshape(B, arr.shape[0]) + + par = arr.shape[0] + K = selected + start = 0 + else: + if not ch: + raise ValueError( + f"virtual tree leaf {nid} has no physical qubit." + ) + + par = arr.shape[0] + # First child: trace the (unbatched) future siblings to + # identity. This avoids broadcasting every virtual-only node + # across the sample batch. + d0 = arr.shape[1] + F0 = int(np.prod(arr.shape[2:])) if len(ch) > 1 else 1 + ur = arr.reshape(par, d0, F0) + env = np.einsum("acF,AdF->acAd", ur, ur.conj()) + rho0 = np.einsum("BaA,acAd->Bcd", rho, env) + phi0 = visit(ch[0], rho0) + # Collapse child 0 into the node tensor -> batched remainder. + K = np.tensordot(phi0, arr, axes=([1], [1])) + start = 1 + + for i in range(start, len(ch)): di = K.shape[2] Fi = int(np.prod(K.shape[3:])) if K.ndim > 3 else 1 Kf = K.reshape(B, par, di, Fi) @@ -506,20 +523,31 @@ def _check_configs(self, configs): def _amplitudes(self, configs): arrays = self._arrays children = self._children - qubit_of_leaf = self._qubit_of_leaf + qubit_of_node = self._qubit_of_node B = configs.shape[0] def visit(nid): ch = children[nid] arr = arrays[nid] - if not ch: # leaf - x = configs[:, qubit_of_leaf[nid]] - return arr[:, x].T # (B, d_par) + q = qubit_of_node.get(nid) + if q is not None: + x = configs[:, q] + K = np.moveaxis(arr[:, x, ...], 1, 0) + if not ch: + return K.reshape(B, arr.shape[0]) + start = 0 + else: + if not ch: + raise ValueError( + f"virtual tree leaf {nid} has no physical qubit." + ) + K = None + start = 0 + par = arr.shape[0] - K = None for i, child in enumerate(ch): phi_c = visit(child) - if i == 0: + if K is None and i == start: K = np.tensordot(phi_c, arr, axes=([1], [1])) else: di = K.shape[2] diff --git a/src/pepsy/tensors/constructors.py b/src/pepsy/tensors/constructors.py index 16c7c16..0aed728 100644 --- a/src/pepsy/tensors/constructors.py +++ b/src/pepsy/tensors/constructors.py @@ -461,8 +461,8 @@ def _fermionic_product_fock_specs(fermion, n, occupations, site_charge): return specs, {site: charge for site, (charge, _) in specs.items()} -def _set_fermionic_product_leaf(tensor, physical_index, *, charge, basis_index): - """Replace a Symmray leaf by one selected Fock-basis vector in-place.""" +def _set_fermionic_product_site(tensor, physical_index, *, charge, basis_index): + """Set one Symmray physical tensor to a selected Fock-basis vector.""" data = tensor.data physical_axis = tensor.inds.index(physical_index) chargemap = data.indices[physical_axis].chargemap @@ -572,7 +572,7 @@ def ps_to_mps( to_backend=None, ) for site, (charge, basis_index) in fock_specs.items(): - _set_fermionic_product_leaf( + _set_fermionic_product_site( state.mps[site], state.mps.site_ind(site), charge=charge, basis_index=basis_index, ) @@ -616,6 +616,7 @@ def ps_to_ttn( *, tree=None, order=None, + root_qubit=None, structure="balanced", max_arity=2, community_frac=0.35, @@ -652,6 +653,10 @@ def ps_to_ttn( ``order`` (or ``range(n)``) using ``structure``. order : sequence of int, optional Leaf order used to build a plan when ``tree`` is not supplied. + When ``root_qubit`` is set, this contains every other qubit. + root_qubit : int, optional + Qubit carried by the top tensor rather than a structural leaf. When an + explicit ``tree`` is supplied, this must match its root site. structure, max_arity, community_frac, star_frac Forwarded to :meth:`TreePlan.from_order`. chi : int, optional @@ -701,15 +706,27 @@ def ps_to_ttn( if tree is not None and order is not None: raise ValueError("pass either tree= or order=, not both.") if tree is None: + if root_qubit is not None: + root_qubit = int(root_qubit) if order is None: - order = range(n) + order = ( + range(n) + if root_qubit is None + else (q for q in range(n) if q != root_qubit) + ) plan = TreePlan.from_order( order, structure=structure, max_arity=max_arity, community_frac=community_frac, star_frac=star_frac, + root_qubit=root_qubit, ) + if plan.n != n: + raise ValueError( + f"constructed tree contains {plan.n} qubits, " + f"but n={n} was requested." + ) else: if not isinstance(tree, TreePlan): raise TypeError("tree must be a TreePlan.") @@ -718,6 +735,10 @@ def ps_to_ttn( raise ValueError( f"tree contains {plan.n} qubits, but n={n} was requested." ) + if root_qubit is not None and int(root_qubit) != plan.root_qubit: + raise ValueError( + "root_qubit does not match the supplied tree plan." + ) if fermion is not None: from .symmetric import ( # pylint: disable=import-outside-toplevel @@ -760,8 +781,8 @@ def ps_to_ttn( node_tag_id=node_tag_id, ) for qubit, (charge, basis_index) in fock_specs.items(): - _set_fermionic_product_leaf( - ttn.node_tensor(ttn.leaf_of_qubit(qubit)), ttn.site_ind(qubit), + _set_fermionic_product_site( + ttn.node_tensor(ttn.node_of_qubit(qubit)), ttn.site_ind(qubit), charge=charge, basis_index=basis_index, ) _apply_to_tensor_network_arrays(ttn, to_backend) @@ -792,7 +813,7 @@ def ps_to_ttn( ) local_vec = np.array([math.cos(theta), math.sin(theta)], dtype=dtype) for q in range(n): - tensor = ttn.node_tensor(ttn.leaf_of_qubit(q)) + tensor = ttn.node_tensor(ttn.node_of_qubit(q)) phys_axis = tensor.inds.index(ttn.site_ind(q)) data = np.zeros_like(tensor.data, dtype=dtype) slicer = [0] * data.ndim @@ -816,6 +837,7 @@ def hrs_to_ttn( *, tree=None, order=None, + root_qubit=None, structure="balanced", max_arity=2, community_frac=0.35, @@ -832,11 +854,12 @@ def hrs_to_ttn( ): """Create a random product or charge-preserving Symmray TTN. - With ``fermion=`` the leaves receive the model's physical charge sectors, - while internal tree nodes are neutral and every virtual tree edge is a - conjugate pair of Symmray charge-sector indices. ``chi`` is the requested - total virtual-bond dimension. All block-sparse and fermionic operations - are delegated to Symmray/Quimb. + With ``fermion=`` the physical sites receive the model's charge sectors, + while virtual-only internal nodes are neutral and every virtual tree edge + is a conjugate pair of Symmray charge-sector indices. ``root_qubit`` places + one physical site on the top tensor. ``chi`` is the requested total + virtual-bond dimension. All block-sparse and fermionic operations are + delegated to Symmray/Quimb. """ from ..optimizers.tree import TreePlan, TreeTensorNetwork @@ -855,19 +878,37 @@ def hrs_to_ttn( if tree is not None and order is not None: raise ValueError("pass either tree= or order=, not both.") if tree is None: + if root_qubit is not None: + root_qubit = int(root_qubit) + if order is None: + order = ( + range(n) + if root_qubit is None + else (q for q in range(n) if q != root_qubit) + ) plan = TreePlan.from_order( - range(n) if order is None else order, + order, structure=structure, max_arity=max_arity, community_frac=community_frac, star_frac=star_frac, + root_qubit=root_qubit, ) + if plan.n != n: + raise ValueError( + f"constructed tree contains {plan.n} qubits, " + f"but n={n} was requested." + ) else: if not isinstance(tree, TreePlan): raise TypeError("tree must be a TreePlan.") if tree.n != n: raise ValueError(f"tree contains {tree.n} qubits, but n={n} was requested.") plan = tree + if root_qubit is not None and int(root_qubit) != plan.root_qubit: + raise ValueError( + "root_qubit does not match the supplied tree plan." + ) if fermion is not None: from .symmetric import ( # pylint: disable=import-outside-toplevel @@ -939,7 +980,7 @@ def hrs_to_ttn( [np.cos(theta / 2.0), np.exp(1j * phi) * np.sin(theta / 2.0)], dtype=dtype, ) - tensor = ttn.node_tensor(ttn.leaf_of_qubit(q)) + tensor = ttn.node_tensor(ttn.node_of_qubit(q)) physical_axis = tensor.inds.index(ttn.site_ind(q)) data = np.zeros_like(tensor.data, dtype=dtype) selector = [0] * data.ndim diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index 42678c8..83be5e7 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -371,6 +371,41 @@ def test_root_physical_qubit_is_first_class_tree_site(): assert state.validate(check_canonical=True) is state +@pytest.mark.parametrize("root_qubit", [0, 1]) +def test_two_qubit_tree_uses_distinct_root_and_leaf_sites(root_qubit): + """The smallest root-site tree uses a unary root over one physical leaf.""" + leaf_qubit = 1 - root_qubit + plan = TreePlan.from_order( + [leaf_qubit], structure="balanced", root_qubit=root_qubit, + ) + layered = TreePlan.build_layered( + [leaf_qubit], block_size=2, root_qubit=root_qubit, + ) + found = TreeLayoutFinder( + [], n=2, root_qubit=root_qubit, max_arity=2, + ).run() + automatic = TreeOptimizer( + None, n=2, root_qubit=root_qubit, max_arity=2, run=False, + ) + + for candidate in (plan, layered, found): + assert candidate.n == 2 + assert candidate.root_qubit == root_qubit + assert candidate.node_of_qubit[root_qubit] == candidate.root + assert candidate.node_of_qubit[leaf_qubit] != candidate.root + assert len(candidate.children[candidate.root]) == 1 + assert automatic.plan.node_of_qubit[root_qubit] == automatic.plan.root + assert automatic.tn.validate(check_canonical=True) is automatic.tn + + stream = [ + (pepsy.h(), root_qubit), + (pepsy.cnot(), (root_qubit, leaf_qubit)), + ] + opt = TreeOptimizer(stream, tree=plan, chi=8) + assert _fidelity(_exact_state(stream, 2), opt.to_dense()) > 1 - 1e-12 + assert opt.tn.validate(check_canonical=True) is opt.tn + + def test_root_physical_qubit_gate_and_submpo_replay_are_exact(): """Direct gates and a structured sub-MPO can target the top physical leg.""" plan = TreePlan.from_order( @@ -432,6 +467,7 @@ def test_root_physical_qubit_layout_and_cap_are_root_aware(): assert automatic.plan.root_qubit == 4 opt = TreeOptimizer(None, tree=plan, chi=16, run=False) + assert opt.layout_report()["root_qubit"] == 4 opt.apply_1q(pepsy.h(), 4) x = np.array([[0.0, 1.0], [1.0, 0.0]]) assert opt.tn.local_expectation(x, 4) == pytest.approx(1.0) @@ -443,6 +479,21 @@ def test_root_physical_qubit_layout_and_cap_are_root_aware(): assert opt.tn.validate(check_canonical=True) is opt.tn +def test_explicit_tree_rejects_mismatched_n(): + """Explicit plans enforce the same qubit-count invariant as finders.""" + plan = TreePlan.from_order( + range(4), structure="balanced", root_qubit=4, + ) + with pytest.raises( + ValueError, match=r"tree contains 5 qubits, but n=4" + ): + TreeOptimizer(None, n=4, tree=plan, run=False) + with pytest.raises( + ValueError, match=r"tree contains 5 qubits, but n=4" + ): + TreeOptimizer(None, n=4, layout=plan, run=False) + + def test_layout_finder_builds_valid_tree(): """With max_arity=2 the finder returns a rooted binary tree over all qubits.""" rng = np.random.default_rng(8) @@ -2024,6 +2075,36 @@ def test_ps_to_ttn_matches_product_state_constructor_api(): assert explicit.plan is plan +def test_product_ttn_constructors_support_a_physical_root_site(): + """Product/random public constructors resolve every physical node.""" + theta = 0.23 + local = np.array([np.cos(theta), np.sin(theta)], dtype="complex128") + expected = local + for _ in range(4): + expected = np.kron(expected, local) + + plan = TreePlan.from_order( + [0, 1, 3, 4], structure="balanced", root_qubit=2, + ) + explicit = pepsy.ps_to_ttn(5, tree=plan, theta=theta) + automatic = pepsy.ps_to_ttn(5, root_qubit=2, theta=theta) + smallest = pepsy.ps_to_ttn(2, root_qubit=1, theta=theta) + random = pepsy.hrs_to_ttn(5, root_qubit=2, seed=11) + + assert explicit.plan is plan + assert automatic.plan.root_qubit == 2 + assert smallest.plan.n == 2 + assert smallest.plan.root_qubit == 1 + assert random.plan.root_qubit == 2 + assert np.allclose(explicit.to_statevector(), expected) + assert np.allclose(automatic.to_statevector(), expected) + assert random.to_statevector().shape == (2**5,) + assert explicit.validate(check_canonical=True) is explicit + + with pytest.raises(ValueError, match="root_qubit does not match"): + pepsy.ps_to_ttn(5, tree=plan, root_qubit=3) + + def test_ttn_copy_preserves_geometry_and_type(): """copy() keeps the plan, ids, and class, with an independent tid cache.""" plan = TreePlan.from_order(range(6), structure="balanced") diff --git a/tests/test_public_api.py b/tests/test_public_api.py index 00d0b21..7c64c83 100644 --- a/tests/test_public_api.py +++ b/tests/test_public_api.py @@ -32,11 +32,12 @@ def test_namespace_exports_have_clear_core_and_advanced_groups(): def test_tree_optimizers_are_available_from_high_level_api(): """Tree layout and execution helpers resolve from ``import pepsy as py``.""" - from pepsy.optimizers.tree import TreeLayoutFinder, TreeOptimizer + from pepsy.optimizers.tree import TreeLayoutFinder, TreeOptimizer, TreePlan from pepsy.optimizers.tree_stabilizer import TreeStabOptimizer assert pepsy.TreeLayoutFinder is TreeLayoutFinder assert pepsy.TreeOptimizer is TreeOptimizer + assert pepsy.TreePlan is TreePlan assert pepsy.TreeStabOptimizer is TreeStabOptimizer @@ -59,6 +60,7 @@ def test_tree_optimizers_are_available_from_high_level_api(): "TreeEnergyOptimizer", "TreeLayoutFinder", "TreeOptimizer", + "TreePlan", "TreeStabOptimizer", "TreeTensorNetwork", "DeferredInjectionRecord", "DeferredInjectionReport", "DeferredProjectionRecord", @@ -132,6 +134,7 @@ def test_internal_symbols_not_exported(): "TreeEnergyOptimizer", "TreeLayoutFinder", "TreeOptimizer", + "TreePlan", "TreeStabOptimizer", "TreeTensorNetwork", "TreeSampler", "TreeBatchSampleResult", "TreeSampleResult", diff --git a/tests/test_tree_sampler.py b/tests/test_tree_sampler.py index e6ed239..72a1ab7 100644 --- a/tests/test_tree_sampler.py +++ b/tests/test_tree_sampler.py @@ -4,7 +4,7 @@ import pytest import pepsy -from pepsy.optimizers.tree import TreeOptimizer +from pepsy.optimizers.tree import TreeOptimizer, TreePlan from pepsy.sampling import ( TreeBatchSampleResult, TreeSampleResult, @@ -80,6 +80,34 @@ def test_amplitudes_match_statevector_up_to_phase(): assert np.allclose(sampler.probabilities(_all_configs(n)), np.abs(amps) ** 2) +def test_physical_root_probabilities_and_amplitudes_match_statevector(): + """Sampling treats the optional root physical leg as an ordinary site.""" + n = 5 + root_qubit = 2 + rng = np.random.default_rng(71) + plan = TreePlan.from_order( + [0, 1, 3, 4], structure="balanced", root_qubit=root_qubit, + ) + stream = _random_stream(n, 30, rng, two_qubit_frac=0.7) + opt = TreeOptimizer(stream, tree=plan, chi=128) + psi, exact = _exact_probs(opt, n) + configs = _all_configs(n) + + sampler = TreeSampler(opt, seed=0) + amplitudes = sampler.amplitudes(configs) + probabilities = sampler.probabilities(configs) + pivot = int(np.argmax(np.abs(psi))) + phase = psi[pivot] / amplitudes[pivot] + + assert np.max(np.abs(amplitudes * phase - psi)) < 1e-10 + assert np.max(np.abs(probabilities - exact)) < 1e-10 + result = sampler.sample_batch(32, seed=3) + assert result.configs.shape == (32, n) + assert np.allclose( + result.probs, sampler.probabilities(result.configs), atol=1e-12 + ) + + # -- empirical sampling ------------------------------------------------------- @@ -247,7 +275,7 @@ def test_public_api_exports_tree_sampler(): # -- fermionic tree sampling -------------------------------------------------- -def _fermionic_tree(*, chi=64, steps=4): +def _fermionic_tree(*, chi=64, steps=4, root_qubit=None): """Build a mildly entangled U1U1 spinful Fermi-Hubbard tree (L=4).""" from pepsy.optimizers.tree import TreeLayoutFinder @@ -269,7 +297,7 @@ def _fermionic_tree(*, chi=64, steps=4): plan = TreeLayoutFinder( [(fermion.hopping_gate(0.1, t=t, imaginary=False), e) for e in edges_1d], - n=L, chi=8, objective="hybrid", + n=L, chi=8, objective="hybrid", root_qubit=root_qubit, ).recommend_arities((2, 3, 4), seed=0)["plan"] seed_ttn = pepsy.ps_to_ttn( L, tree=plan, fermion=fermion, occupations=occupations, dtype=dtype @@ -301,9 +329,12 @@ def _all_base_d_configs(n, d): ) -def test_fermionic_tree_probabilities_match_statevector(): +@pytest.mark.parametrize("root_qubit", [None, 3]) +def test_fermionic_tree_probabilities_match_statevector(root_qubit): pytest.importorskip("symmray") - engine, fermion, target, L = _fermionic_tree(chi=64) + engine, fermion, target, L = _fermionic_tree( + chi=64, root_qubit=root_qubit, + ) sampler = TreeSampler(engine, fermion=fermion, seed=0) assert sampler._configuration_encoding is not None @@ -367,4 +398,3 @@ def test_non_fermionic_occupations_raises(): assert res.configuration_encoding is None with pytest.raises(ValueError, match="no fermion configuration encoding"): res.occupations() - From 7669a57e3ba49df75e541c3dba72e33056a92152 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Tue, 28 Jul 2026 15:19:07 -0700 Subject: [PATCH 11/70] tree: preserve routed isometry metadata --- .github/skills/tree-optimizer/SKILL.md | 10 +- .../references/performance-layout.md | 4 + docs/api/optimizers/tree.md | 14 +- src/pepsy/optimizers/tree/optimizer.py | 65 ++++++-- tests/test_optimize_tree.py | 155 ++++++++++++++++++ 5 files changed, 227 insertions(+), 21 deletions(-) diff --git a/.github/skills/tree-optimizer/SKILL.md b/.github/skills/tree-optimizer/SKILL.md index 48d3bb7..fb6f684 100644 --- a/.github/skills/tree-optimizer/SKILL.md +++ b/.github/skills/tree-optimizer/SKILL.md @@ -267,9 +267,13 @@ covering range then compressed (quimb's `gate_with_submpo` is `MatrixProductStat physical and exterior state legs, then contract its new state bond into the parent together with the old state/operator bonds. No dense state tensor for the whole Steiner subtree is formed; the last node is the hub. -5. Recover the hub centre by QR, then make one depth-first canonical SVD sweep: - every affected tree edge is truncated once, after the complete operator has - arrived. `renormalize=True` renormalises afterwards (for Kraus/projection). +5. Install every routed Q factor with its ``left_inds`` isometry metadata. + Dense trees can then recover the hub centre through the normal canonical + state machine without repeating those QRs; native fermionic trees retain + their explicit graded QR recovery. Finally make one depth-first canonical + SVD sweep: every affected tree edge is truncated once, after the complete + operator has arrived. `renormalize=True` renormalises afterwards (for + Kraus/projection). State bonds are always read from the live tensors because gate application can rename them. New state message bonds are fresh per-update names, while operator diff --git a/.github/skills/tree-optimizer/references/performance-layout.md b/.github/skills/tree-optimizer/references/performance-layout.md index 992ce45..f89064d 100644 --- a/.github/skills/tree-optimizer/references/performance-layout.md +++ b/.github/skills/tree-optimizer/references/performance-layout.md @@ -13,6 +13,10 @@ Tree Optimizer skill so the upload-facing `SKILL.md` stays concise. `threadpoolctl` when available. Only raise `threads` in a large-`chi` regime. - The self-healing tid cache (`_nid_to_tid`, `_tid`) validates cached tensor ids against `self.tn.tensor_map`; a stale entry is recomputed safely. +- Dense path and subtree routing preserve each QR-produced Q tensor's + `left_inds`. Canonical recovery therefore recognizes an already-isometric + routed branch without repeating its decomposition. Native fermionic routing + deliberately retains explicit graded QR recovery. - `copy()` shares the immutable `TreePlan`, owns `self.tn.copy()`, resets the tid cache, and derives a deterministic child seed for an independent RNG. diff --git a/docs/api/optimizers/tree.md b/docs/api/optimizers/tree.md index bd75c20..a5608cd 100644 --- a/docs/api/optimizers/tree.md +++ b/docs/api/optimizers/tree.md @@ -144,10 +144,12 @@ nodes. Application then proceeds recursively from subtree leaves to a hub: each local state/operator message is losslessly QR-split on one edge and absorbed by its parent, carrying every still-open operator virtual leg. No dense state tensor -for the whole Steiner subtree is formed. Once all MPO factors have arrived, the -tree is canonicalized about the hub and every touched edge is SVD-compressed -once. Thus every truncation sees the complete operator in an isometric -environment. +for the whole Steiner subtree is formed. Each dense routed Q tensor retains its +`left_inds` isometry metadata, so canonical recovery recognizes that it already +points toward the hub instead of repeating the same QR; native fermionic trees +retain explicit graded QR recovery. Once all MPO factors have arrived, every +touched edge is SVD-compressed once. Thus every truncation sees the complete +operator in an isometric environment. `op` acts on `len(where)` qubits: an array reshaped to `(2,) * 2k` with output indices first, `op[o_0..o_{k-1}, i_0..i_{k-1}]` (a `(2**k, 2**k)` matrix is @@ -701,6 +703,10 @@ available through `truncation_report()`, `get_infidelities()`, and orthogonality centre; `from_plan` records that centre on the network rather than recomputing it on the first gate. Native fermionic product trees are additionally normalized by their exact graded norm readout. +- **Routed isometry reuse.** Dense geodesic and subtree QR routing retains each + Q tensor's `left_inds`, allowing later canonical recovery to reuse the proven + isometry without repeating the decomposition. Native fermionic trees keep + their separate explicit graded QR path. - **State-owned centre.** The orthogonality centre lives on the `TreeTensorNetwork` (`orthogonality_center`, an `_EXTRA_PROPS` field), so the optimizer and the state cannot disagree and the centre is carried by diff --git a/src/pepsy/optimizers/tree/optimizer.py b/src/pepsy/optimizers/tree/optimizer.py index edb260f..b9e0e7c 100644 --- a/src/pepsy/optimizers/tree/optimizer.py +++ b/src/pepsy/optimizers/tree/optimizer.py @@ -2302,8 +2302,19 @@ def _thread_hop(self, u, v): get="tensors", ) merged_v = qtn.tensor_contract(carry, tv) - tu.modify(data=keep.data, inds=keep.inds) - tv.modify(data=merged_v.data, inds=merged_v.inds) + # ``keep`` is the exact Q factor pointing toward ``v``. Preserve that + # isometry metadata so a later canonical walk can recognize the tensor + # without repeating the same QR decomposition. + tu.modify( + data=keep.data, + inds=keep.inds, + left_inds=keep.left_inds, + ) + tv.modify( + data=merged_v.data, + inds=merged_v.inds, + left_inds=None, + ) def _fermionic_thread_hop(self, u, v): """QR-route the operator bond without leaving native graded arrays.""" @@ -2644,6 +2655,39 @@ def _route_subtree_messages( state_inds[v].add(new_bond) operator_inds[v] = set(local[v].inds) - state_inds[v] + def _install_routed_subtree(self, local, snodes, hub): + """Install routed tensors and recover their proven hub centre. + + Dense routing already QR-isometrizes every peeled non-hub tensor toward + ``hub``. Retaining each Q factor's ``left_inds`` lets Quimb's canonical + recovery walk short-circuit those decompositions while still advancing + the canonical-region state machine honestly. Native fermionic tensors + deliberately keep the prior behavior: their graded QR recovery remains + explicit inside :class:`TreeTensorNetwork`. + """ + dense = not self.tn.fermionic + for nid in snodes: + routed = local[nid] + modify_opts = { + "data": routed.data, + "inds": routed.inds, + } + if dense: + if nid == hub: + # The accumulated operator and state norm live here. + modify_opts["left_inds"] = None + else: + if routed.left_inds is None: + raise RuntimeError( + "dense subtree routing lost QR isometry metadata " + f"for non-hub node {nid}." + ) + modify_opts["left_inds"] = routed.left_inds + self.tn.tensor_map[self._tid(nid)].modify(**modify_opts) + + self.tn.canonical_region = frozenset(snodes) + self._move_center(hub) + # -- general multi-qubit / sub-MPO application ---------------------------- def apply_subtree_operator(self, op, where, *, max_bond=None, @@ -3077,14 +3121,11 @@ def _try_apply_native_submpo( "native sub-MPO application left open operator bonds; " "use an MPO with a closed tensor-network contraction." ) - for nid in snodes: - node_t = self.tn.tensor_map[self._tid(nid)] - node_t.modify(data=local[nid].data, inds=local[nid].inds) # The exterior remained isometric toward the updated Steiner subtree. - # Recover a single centre within it by QR, then truncate only now that - # every MPO virtual bond has reached its destination. - self.tn.canonical_region = frozenset(snodes) - self._move_center(hub) + # Dense routed Q tensors retain their isometry metadata, so recovering + # the hub centre is metadata-only; native graded trees keep their + # explicit QR recovery. Truncate only after every MPO bond has arrived. + self._install_routed_subtree(local, snodes, hub) self._compress_subtree( snodes, hub, max_bond=max_bond, cutoff=cutoff, ) @@ -3125,11 +3166,7 @@ def _apply_factorized_subtree_operator_impl( "factorized subtree operator left open operator bonds at its hub." ) - for nid in snodes: - node_t = self.tn.tensor_map[self._tid(nid)] - node_t.modify(data=local[nid].data, inds=local[nid].inds) - self.tn.canonical_region = frozenset(snodes) - self._move_center(hub) + self._install_routed_subtree(local, snodes, hub) self._compress_subtree( snodes, hub, max_bond=max_bond, cutoff=cutoff, ) diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index 83be5e7..13be8af 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -144,6 +144,29 @@ def test_tree_two_site_direct_and_mpo_modes_agree(): assert _fidelity(mpo.to_dense(), exact) > 1 - 1e-9 +def test_dense_path_thread_preserves_qr_isometry_metadata(monkeypatch): + """Every dense path-thread Q keeps its toward-destination isometry.""" + rng = np.random.default_rng(919) + opt = TreeOptimizer(None, n=8, chi=16, run=False) + checked = [] + compress_path = opt._compress_path + + def check_then_compress(path, **kwargs): + for node, toward_destination in zip(path, path[1:]): + tensor = opt.tn.node_tensor(node) + bond = opt.tn.bond(node, toward_destination) + assert tensor.left_inds is not None + assert set(tensor.left_inds) == set(tensor.inds) - {bond} + checked.append(node) + return compress_path(path, **kwargs) + + monkeypatch.setattr(opt, "_compress_path", check_then_compress) + opt.apply_2q(_rand_unitary(2, rng), 0, 7) + + assert checked + assert opt.tn.validate(check_canonical=True) is opt.tn + + def test_tree_mpo_mode_keeps_small_operator_schmidt_components(): """MPO lowering must not apply Quimb's default gate-SVD cutoff.""" x = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex) @@ -214,6 +237,57 @@ def test_tree_multisite_submpo_qr_routes_before_one_subtree_sweep(): assert all(event["max_bond"] == 1 for event in opt.truncation_history) +def test_dense_subtree_hub_recovery_reuses_routed_q_metadata(monkeypatch): + """Dense routed Q tensors recover the hub without another numerical QR.""" + import quimb.tensor.tensor_core as qtc + + rng = np.random.default_rng(52) + n = 8 + where = (0, 3, 7) + gate = _rand_unitary(3, rng) + opt = TreeOptimizer(None, n=n, chi=16, run=False) + expected = np.zeros(2**n, dtype=complex) + expected[0] = 1.0 + expected = _sv_apply_kq(expected, gate, where, n) + + qr_calls = [] + tensor_split = qtc.tensor_split + + def traced_tensor_split(*args, **kwargs): + if kwargs.get("method") == "qr": + qr_calls.append(args[0]) + return tensor_split(*args, **kwargs) + + recoveries = [] + move_center = opt._move_center + + def traced_move_center(target): + region = opt.canonical_region + if region is not None and len(region) > 1 and target in region: + for nid in region: + tensor = opt.tn.node_tensor(nid) + if nid == target: + assert tensor.left_inds is None + continue + toward_hub = opt.plan.node_path(nid, target)[1] + bond = opt.tn.bond(nid, toward_hub) + assert tensor.left_inds is not None + assert set(tensor.left_inds) == set(tensor.inds) - {bond} + before = len(qr_calls) + result = move_center(target) + recoveries.append(len(qr_calls) - before) + return result + return move_center(target) + + monkeypatch.setattr(qtc, "tensor_split", traced_tensor_split) + monkeypatch.setattr(opt, "_move_center", traced_move_center) + opt.apply_subtree_operator(gate, where) + + assert recoveries == [0] + assert _fidelity(expected, opt.to_dense()) > 1 - 1e-10 + assert opt.tn.validate(check_canonical=True) is opt.tn + + def test_tree_mode_is_construction_and_run_override(): """Tree gate implementation mode follows the MPS construction/run API.""" opt = TreeOptimizer(None, n=2, mode="direct", run=False) @@ -2166,6 +2240,7 @@ def test_tree_torch_state_stays_native_across_public_operations(): } assert opt.expectation_pauli("ZZ", (0, 1)) == pytest.approx(1.0) opt.apply_subtree_operator(to_backend(np.eye(8, dtype=complex)), (0, 1, 2)) + assert opt.tn.validate(check_canonical=True) is opt.tn opt.project_pauli("ZZ", (0, 1), +1) assert opt.measure(2, outcome=0) == 0 assert opt.reset(2) == 0 @@ -3032,6 +3107,86 @@ def build_stream(): assert float(tensors.tn_fidelity(engine.p, mps_exact.p)) > 1 - 1e-8 +def test_native_fermionic_submpo_keeps_graded_hub_recovery(monkeypatch): + """Dense isometry metadata must not replace native graded subtree QR.""" + pytest.importorskip("symmray") + L = 4 + fermion = pepsy.Fermion( + spinful=True, + symmetry="U1U1", + t=1.0, + U=8.0, + mu=0.0, + dtype="complex128", + ) + occupations = ((1, 0), (0, 1), (1, 0), (0, 1)) + plan = TreePlan.from_order(range(L), structure="balanced") + seed = pepsy.ps_to_ttn( + L, + tree=plan, + fermion=fermion, + occupations=occupations, + dtype="complex128", + ) + sites = (0, 1, 2) + local_ops = [ + fermion.onsite_gate( + 0.01, site=site, U=8.0, mu=0.0, imaginary=False + ) + for site in sites + ] + submpo = qtn.MPO_product_operator( + local_ops, + sites=sites, + L=L, + upper_ind_id="k{}", + lower_ind_id="b{}", + ) + candidate = TreeOptimizer( + None, + n=L, + tree=plan, + state=seed.copy(), + chi=64, + cutoff=0.0, + run=False, + ) + reference = TreeOptimizer( + None, + n=L, + tree=plan, + state=seed.copy(), + chi=64, + cutoff=0.0, + run=False, + ) + installs = [] + install_routed = candidate._install_routed_subtree + + def traced_install(local, snodes, hub): + installs.append((frozenset(snodes), hub)) + assert candidate.tn.fermionic + return install_routed(local, snodes, hub) + + monkeypatch.setattr(candidate, "_install_routed_subtree", traced_install) + candidate.apply_submpo(submpo, sites) + for site, op in zip(sites, local_ops): + reference.apply_1q(op, site) + + def dense_vector(opt): + tensor = opt.tn.contract(all, optimize="greedy").transpose( + *(opt.tn.site_ind(q) for q in range(L)) + ) + return np.asarray(tensor.data.to_dense()).reshape(-1) + + assert installs + assert ( + _fidelity(dense_vector(candidate), dense_vector(reference)) + > 1 - 1e-10 + ) + assert candidate.tn.validate(check_canonical=True) is candidate.tn + + def test_tree_stable_labels_route_submpo_by_payload_sites(monkeypatch): """Stable logical labels do not disable native structured MPO routing.""" x = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex) From c50e127a1b5aff7ab3ca8b42fff8850f09f03b7c Mon Sep 17 00:00:00 2001 From: rezaquant Date: Tue, 28 Jul 2026 16:03:48 -0700 Subject: [PATCH 12/70] tree: use one-sided subtree compression --- src/pepsy/optimizers/tree/optimizer.py | 7 +- src/pepsy/optimizers/tree/ttn.py | 8 +- tests/test_optimize_tree.py | 101 +++++++++++++++++++++++++ 3 files changed, 113 insertions(+), 3 deletions(-) diff --git a/src/pepsy/optimizers/tree/optimizer.py b/src/pepsy/optimizers/tree/optimizer.py index b9e0e7c..1475cff 100644 --- a/src/pepsy/optimizers/tree/optimizer.py +++ b/src/pepsy/optimizers/tree/optimizer.py @@ -2534,7 +2534,7 @@ def _split_with_diagnostics( return left, right def _compress_edge_with_diagnostics( - self, u, v, *, max_bond=None, cutoff=None, + self, u, v, *, max_bond=None, cutoff=None, reduced=True, ): """Compress one live tree edge and record its truncation diagnostics.""" max_bond = ( @@ -2572,7 +2572,8 @@ def _compress_edge_with_diagnostics( # Keep the live canonical-region metadata in one place: the TTN edge # wrapper performs the compression and advances its tracked centre. self.tn.compress_edge_( - u, v, max_bond=max_bond, cutoff=cutoff, absorb="right" + u, v, max_bond=max_bond, cutoff=cutoff, absorb="right", + reduced=reduced, ) bond_after = self.tn.bond(u, v) after_bond = int(self.tn.ind_size(bond_after)) @@ -2607,6 +2608,7 @@ def _compress_subtree(self, snodes, hub, *, max_bond=None, cutoff=None): """ snodes = frozenset(snodes) self._move_center(hub) + forward_reduced = True if self.tn.fermionic else "left" def descend(node, parent): children = sorted( @@ -2617,6 +2619,7 @@ def descend(node, parent): for child in children: self._compress_edge_with_diagnostics( node, child, max_bond=max_bond, cutoff=cutoff, + reduced=forward_reduced, ) descend(child, node) self.tn.canonize_edge_(child, node, absorb="right") diff --git a/src/pepsy/optimizers/tree/ttn.py b/src/pepsy/optimizers/tree/ttn.py index 93a3683..e08d484 100644 --- a/src/pepsy/optimizers/tree/ttn.py +++ b/src/pepsy/optimizers/tree/ttn.py @@ -1059,13 +1059,18 @@ def canonize_edge_(self, a, b, absorb="right"): return self def compress_edge_(self, a, b, *, max_bond=None, cutoff=1e-12, - absorb="right"): + absorb="right", reduced=True): """Compress the tree edge ``a -> b`` in place. Dense/nonfermionic trees delegate to Quimb's ``compress_between``. Native fermionic trees explicitly SVD the complete two-node tensor. The tracked :attr:`orthogonality_center` advances as for :meth:`canonize_edge_`. + + ``reduced`` is forwarded only on the dense path. Quimb's one-sided + ``"left"`` mode is exact when node ``b`` is already isometric on its + non-shared legs. Native fermionic compression ignores this option and + retains its explicit graded split. """ previous = self.orthogonality_center if self.fermionic: @@ -1083,6 +1088,7 @@ def compress_edge_(self, a, b, *, max_bond=None, cutoff=1e-12, max_bond=max_bond, cutoff=cutoff, absorb=absorb, + reduced=reduced, ) self._invalidate_norm_cache() self._track_edge_center(a, b, absorb, previous=previous) diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index 13be8af..5e578de 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -288,6 +288,88 @@ def traced_move_center(target): assert opt.tn.validate(check_canonical=True) is opt.tn +def test_dense_subtree_uses_proven_one_sided_compression(monkeypatch): + """A routed gate ladder matches full reduction while skipping child QRs.""" + rng = np.random.default_rng(53) + n = 8 + plan = TreePlan.from_order(range(n), structure="balanced") + seed = TreeTensorNetwork.rand(plan, D=2, seed=53) + optimized = TreeOptimizer( + None, + tree=plan, + state=seed.copy(), + chi=2, + cutoff=1e-12, + run=False, + ) + reference = TreeOptimizer( + None, + tree=plan, + state=seed.copy(), + chi=2, + cutoff=1e-12, + run=False, + ) + + left_reductions = [] + compress_edge = TreeTensorNetwork.compress_edge_ + + def traced_compress_edge(tn, a, b, **kwargs): + if tn is optimized.tn: + reduced = kwargs.get("reduced", True) + if reduced == "left": + child = tn.node_tensor(b) + bond = tn.bond(a, b) + assert child.left_inds is not None + assert set(child.left_inds) == set(child.inds) - {bond} + left_reductions.append((a, b)) + return compress_edge(tn, a, b, **kwargs) + + full_compress = reference._compress_edge_with_diagnostics + + def force_full_reduction( + u, v, *, max_bond=None, cutoff=None, reduced=True, + ): + del reduced + return full_compress( + u, + v, + max_bond=max_bond, + cutoff=cutoff, + reduced=True, + ) + + monkeypatch.setattr( + TreeTensorNetwork, "compress_edge_", traced_compress_edge, + ) + monkeypatch.setattr( + reference, "_compress_edge_with_diagnostics", force_full_reduction, + ) + + exact = seed.to_statevector() + ladder_supports = ( + (0, 2, 4), + (1, 3, 5), + (2, 4, 6), + (3, 5, 7), + ) + for where in ladder_supports: + gate = _rand_unitary(3, rng) + optimized.apply_subtree_operator(gate, where) + reference.apply_subtree_operator(gate, where) + exact = _sv_apply_kq(exact, gate, where, n) + + assert left_reductions + assert _fidelity(optimized.to_dense(), reference.to_dense()) > 1 - 1e-10 + assert any(event["truncated"] for event in optimized.truncation_history) + assert _fidelity(optimized.to_dense(), exact) == pytest.approx( + _fidelity(reference.to_dense(), exact), + rel=1e-10, + abs=1e-12, + ) + assert optimized.tn.validate(check_canonical=True) is optimized.tn + + def test_tree_mode_is_construction_and_run_override(): """Tree gate implementation mode follows the MPS construction/run API.""" opt = TreeOptimizer(None, n=2, mode="direct", run=False) @@ -3162,13 +3244,30 @@ def test_native_fermionic_submpo_keeps_graded_hub_recovery(monkeypatch): ) installs = [] install_routed = candidate._install_routed_subtree + compressions = [] + compress_edge = candidate._compress_edge_with_diagnostics def traced_install(local, snodes, hub): installs.append((frozenset(snodes), hub)) assert candidate.tn.fermionic return install_routed(local, snodes, hub) + def traced_compress_edge( + u, v, *, max_bond=None, cutoff=None, reduced=True, + ): + compressions.append(reduced) + return compress_edge( + u, + v, + max_bond=max_bond, + cutoff=cutoff, + reduced=reduced, + ) + monkeypatch.setattr(candidate, "_install_routed_subtree", traced_install) + monkeypatch.setattr( + candidate, "_compress_edge_with_diagnostics", traced_compress_edge, + ) candidate.apply_submpo(submpo, sites) for site, op in zip(sites, local_ops): reference.apply_1q(op, site) @@ -3180,6 +3279,8 @@ def dense_vector(opt): return np.asarray(tensor.data.to_dense()).reshape(-1) assert installs + assert compressions + assert all(reduced is True for reduced in compressions) assert ( _fidelity(dense_vector(candidate), dense_vector(reference)) > 1 - 1e-10 From a7633ea1b5107b4840eaa2e8bb9789243cc34d1c Mon Sep 17 00:00:00 2001 From: rezaquant Date: Tue, 28 Jul 2026 17:42:12 -0600 Subject: [PATCH 13/70] Add qMERA fermion schedules and lightcones --- docs/api/optimizers/mera.md | 327 +++++++- docs/examples.md | 14 +- examples/qmera_fermion_hubbard_2d.py | 106 +++ examples/qmera_fermion_hubbard_4x4_pbc.py | 73 ++ examples/qmera_majorana_2d.py | 96 +++ examples/qmera_scale_plan_6x6.py | 42 + src/pepsy/optimizers/__init__.py | 4 + src/pepsy/optimizers/mera/__init__.py | 29 +- src/pepsy/optimizers/mera/builders.py | 304 ++++++- src/pepsy/optimizers/mera/cache.py | 44 +- src/pepsy/optimizers/mera/fermions.py | 223 ++++- src/pepsy/optimizers/mera/gates.py | 8 + src/pepsy/optimizers/mera/lightcones.py | 493 ++++++++++- src/pepsy/optimizers/mera/schedules.py | 962 +++++++++++++++++++--- src/pepsy/optimizers/mera/schematics.py | 207 ++++- src/pepsy/tensors/symmetric.py | 522 ++++++++---- tests/test_optimize_mera.py | 556 ++++++++++++- 17 files changed, 3679 insertions(+), 331 deletions(-) create mode 100644 examples/qmera_fermion_hubbard_2d.py create mode 100644 examples/qmera_fermion_hubbard_4x4_pbc.py create mode 100644 examples/qmera_majorana_2d.py create mode 100644 examples/qmera_scale_plan_6x6.py diff --git a/docs/api/optimizers/mera.md b/docs/api/optimizers/mera.md index 2c65c5b..0e6cffe 100644 --- a/docs/api/optimizers/mera.md +++ b/docs/api/optimizers/mera.md @@ -46,6 +46,26 @@ energy = builder.parametric_loss( ) ``` +For a fixed MERA-like state, `lightcone_energy(...)` exposes the same local +contraction primitive directly. It selects only the reverse cone, applies each +local operator with `.gate`, and can reuse one Pepsy/cotengra path per cone +topology: + +```python +from pepsy.optimizers.mera import lightcone_energy + +path_cache = builder.contraction_path_cache(max_repeats=16) +energy = lightcone_energy( + mera, + {(0, 1): h2, (2, 3): h2}, + energy_per_site=False, + path_cache=path_cache, +) +``` + +This is the fixed-state analogue of `builder.parametric_loss(...)`; the latter +rebuilds the qMERA cone from a parameter dictionary on every evaluation. + For repeated optimization, compile static local-cone contractions once and reuse them from NumPy, Torch, or JAX-compatible parameter dictionaries: @@ -76,11 +96,34 @@ param_opt = builder.parametric_optimizer( result = param_opt.run(solver="torch-adam", n_steps=10, compiled=True) ``` +### qMERA schematics + +`QMeraSchedule.draw_schematic()` uses Quimb's manual `schematic.Drawing` +primitives. The default `style="clean"` view separates the input sites, +disentangler (`D`), isometry (`W`), and coarse-output stages, with colored +patches and arrows for the RG flow: + +```python +drawing = schedule.draw_schematic( + style="clean", # or "register" for the low-level wiring view + figsize=(14, 5), + label_sites=True, + label_blocks=True, + scale_figsize=False, +) +``` + +The clean view is intended for explaining a schedule or a fermionic block +layout; `schedule.schematic_blocks()` remains the machine-readable placement +audit. + Native Symmray fermion helpers are available under this module, but the -fermion convention is explicit. Use `QMeraGeometry(site_modes=("up", "down"))` -and `qmera_symmray_fermi_hubbard_terms(...)` for the native graded-array path; -do not mix it silently with dense spin or Jordan-Wigner local operators. The -same model can supply both site-native MPS terms and qMERA mode terms: +fermion convention is explicit. The `Fermion` helper can now be supplied to +`QMeraBuilder`; it infers the canonical spinful mode pair and mode-major +register order, then extracts qMERA mode terms internally. Do not mix this +native graded-array path silently with dense spin or Jordan-Wigner local +operators. The same model can supply both site-native MPS terms and qMERA mode +terms: ```python import pepsy @@ -101,26 +144,286 @@ geometry = QMeraGeometry(shape=3, site_modes=("up", "down")) qmera_terms = fermion.local_terms(geometry, layout="qmera") ``` -When the builder owns the geometry, it can perform the same conversion and -construct the optimizer directly: +For the normal spinful Hubbard workflow, let the builder own the mode +expansion and conversion: ```python +from pepsy.optimizers.mera import ( + QMeraBuilder, + QMeraSymmrayFermionBackend, + symmray_fermion_gate_registry, +) + +backend = QMeraSymmrayFermionBackend.from_fermion(fermion) builder = QMeraBuilder( - geometry=geometry, - site_modes=("up", "down"), - mode_order="mode-major", - # use a Symmray fermion gate registry and product-state factory here + shape=3, + fermion=fermion, + gate_registry=symmray_fermion_gate_registry(backend=backend), + gate_family="symmray-fsim", + product_state_factory=backend.product_state, ) -terms = builder.fermion_terms(fermion) +geometry = builder.geometry # inferred from the model +terms = builder.fermion_terms() # inferred from the model optimizer = builder.fermion_parametric_optimizer( - fermion, energy_per_site=False, ) ``` +Passing `site_modes` or `mode_order` remains useful when testing a custom +register convention. A site-layout object such as +`fermion.hamiltonian(edges)` is still a valid native MPS/PEPS Hamiltonian, but +it does not by itself specify qMERA's explicit mode registers or RG schedule; +use `fermion.local_terms(geometry, layout="qmera")` or the builder shortcut +above for that conversion. + For a larger native Torch workflow, use the corresponding examples maintained in the separate `pepsy_examples` repository; the package API is demonstrated by the builder flow above. +## 2D multimode RG schedules + +For a two-dimensional geometry, `shape` describes physical lattice sites and +`site_modes` expands each site into explicit register modes. For example, +`site_modes=("up", "down")` gives two modes per site. The RG schedule blocks +physical sites first, then retains every mode on the representative coarse +site. Spatial gates are grouped by mode, so the schedule never silently turns +an `up`/`down` pair into a spatial fermion gate. `mode_order="mode-major"` +selects registers as +`((site_0, "up"), ..., (site_n, "up"), (site_0, "down"), ...)`; the default +`"site-major"` interleaves modes at each site. + +Use `QMeraScaleSpec` when the RG geometry changes from one scale to the next. +For example, this generic 6x6 periodic plan reduces 6x6 to 3x3 with 2x2 +covering blocks, then reduces 3x3 to one site with a 3x3 covering block and +vertical 3-site internal disentangler strips: + +```python +from pepsy.optimizers.mera import ( + QMeraBuilder, + QMeraDisentanglerSpec, + QMeraIsometrySpec, + QMeraScaleSpec, +) + +scale_plan = ( + QMeraScaleSpec( + isometry=QMeraIsometrySpec(block_shape=(2, 2)), + disentangler=QMeraDisentanglerSpec( + block_shape=(2, 2), + placement="boundary-square", + ), + ), + QMeraScaleSpec( + isometry=QMeraIsometrySpec(block_shape=(3, 3)), + disentangler=QMeraDisentanglerSpec( + block_shape=3, + orientation="vertical", + placement="within-block", + circuit_depth=3, + ), + ), +) +builder = QMeraBuilder( + shape=(6, 6), + boundary="periodic", + scales=scale_plan, +) +schedule = builder.build_schedule() # 36 -> 9 -> 1 active sites +``` + +`orientation="vertical"` resolves an integer strip length to a 1x3 block in +the geometry's `(x, y)` convention. The three circuit rounds cover the three +nearest-neighbor edges of an odd periodic 3-site line without overlapping +gates. `placement="within-block"` is for internal disentanglers; the default +`"boundary-faces"` and `"boundary-square"` placements remain the +inter-isometry-boundary choices. + +The schedule is independent of the operator representation. A native +spinful Fermi--Hubbard workflow therefore uses `U1U1` and the Symmray FSIM +registry: + +```python +from pepsy.optimizers.mera import ( + QMeraBuilder, + QMeraGeometry, + QMeraSymmrayFermionBackend, + symmray_fermion_gate_registry, +) + +geometry = QMeraGeometry( + shape=(4, 4), + site_modes=("up", "down"), + mode_order="mode-major", +) +backend = QMeraSymmrayFermionBackend(symmetry="U1U1") +builder = QMeraBuilder( + geometry=geometry, + gate_registry=symmray_fermion_gate_registry(backend=backend), + gate_family="symmray-fsim", + isometry={"block_size": (2, 2), "gate_family": "symmray-fsim"}, + product_state_factory=backend.product_state, +) +``` + +For a 1D multimode geometry, use `mode_order="mode-major"` when spatial +brickwall gates should connect equal flavors. The scheduler partitions every +1D block by explicit mode before pairing, so native `U1U1` gates cannot +silently become spin-changing `up`--`down` gates. A one-site two-mode block +therefore receives no invalid gate rather than raising during gate materialization. + +For the explicit 4x4 periodic construction, use a 2x2 square disentangler +around every inter-block face and a 2x2 covering unitary for each RG block: + +```python +from pepsy.optimizers.mera import ( + QMeraBuilder, + QMeraDisentanglerSpec, + QMeraGeometry, + QMeraIsometrySpec, + QMeraSymmrayFermionBackend, + QMeraUnitarySpec, + symmray_fermion_gate_registry, +) + +geometry = QMeraGeometry( + shape=(4, 4), + boundary="periodic", + site_modes=("up", "down"), + mode_order="mode-major", +) +backend = QMeraSymmrayFermionBackend( + symmetry="U1U1", + site_modes=("up", "down"), + mode_order="mode-major", +) +hubbard_unitary = QMeraUnitarySpec( + gate_family="symmray-hubbard", + family="fermion", + arity_kind="mode", + symmetry="U1U1", + preserves_parity=True, + metadata={"model": "fermi-hubbard", "term": "hopping"}, +) +builder = QMeraBuilder( + geometry=geometry, + gate_registry=symmray_fermion_gate_registry(backend=backend), + disentangler=QMeraDisentanglerSpec( + block_shape=(2, 2), + unitary=hubbard_unitary, + placement="boundary-square", + circuit_depth=2, + periodic_wrap=True, + ), + isometry=QMeraIsometrySpec( + block_shape=(2, 2), + unitary=hubbard_unitary, + circuit_depth=2, + implementation="unitary-completion", + ), + max_layers=2, +) +schedule = builder.build_schedule() # 4x4 -> 2x2 -> 1, including PBC wraps +``` + +Set `parameter_sharing` on `QMeraUnitarySpec` to choose the parameter scope: + +- `"per-placement"`: every scheduled gate has independent parameters. +- `"per-block"`: gates in one RG block share parameters across brickwall + rounds; different blocks remain independent. +- `"per-scale"`: all placements of one stage at one RG scale share parameters. +- `"per-axis"`: one parameter set per scale and spatial axis. +- `"shared"`: one parameter set is shared across all scales for that stage. + +The default is `"per-placement"`. + +`QMeraIsometrySpec` currently means a square local unitary circuit whose +output representative is retained by the RG schedule; it does not yet build a +rectangular tensor with an exact isometry constraint. Set +`implementation="true-isometry"` only when that backend is implemented. +`symmray-hubbard` is the native two-mode, number-conserving Hubbard hopping +layer used by this schedule (`symmray-fsim` is its lower-level alias). The +onsite `U n_up n_down - mu n` terms remain part of the +Fermi--Hubbard Hamiltonian supplied through `Fermion.local_terms(...)`; this +keeps gate topology and Hamiltonian terms separate while preserving the native +Symmray `U1U1` grading. + +Use `convert_terms=False` when the Hamiltonian terms are already native +Symmray arrays. This preserves their graded contraction behavior. + +## Grouped cones and direct validation + +`builder.parametric_loss(...)` groups terms with the same input support and +scheduled gate topology. `builder.contraction_path_cache(...)` creates a lazy +cache with one reusable contraction optimizer per topology: + +```python +path_cache = builder.contraction_path_cache(max_repeats=16) +energy = builder.parametric_loss( + params, + terms, + schedule=schedule, + convert_terms=False, + path_cache=path_cache, +) +``` + +For debugging a new schedule, compare this local-cone result with +`builder.direct_parametric_loss(...)`. The latter constructs the complete +direct-gate qMERA tensor network and is intentionally a validation oracle, +not the optimization path. Agreement should be checked for representative +1D, 2D multimode, native Hubbard, and native Majorana cases. + +## Majorana and pairing convention + +The implemented true-Majorana convention is one spinless complex mode per +physical site with native `Z2` fermion parity: + +```python +import pepsy as py +from pepsy.optimizers.mera import ( + QMeraGeometry, + qmera_symmray_majorana_terms, + symmray_majorana_gate_registry, +) + +fermion = py.Fermion(spinful=False, symmetry="Z2") +gamma_x = fermion.majorana_operator("x", site=0) # c + c^dag +gamma_y = fermion.majorana_operator("y", site=0) # -i (c - c^dag) +pairing = fermion.pairing_operator((0, 1), phase=0.2) +geometry = QMeraGeometry(shape=(2, 2), site_modes=("mode",)) +terms = qmera_symmray_majorana_terms( + geometry, + fermion=fermion, + coupling=0.4, + pairing=0.2, +) +registry = symmray_majorana_gate_registry() +``` + +Individual Majoranas are parity odd (`charge=1`), while Majorana bilinears, +pairing operators, and their gates are neutral (`charge=0`). This is why the +Majorana path uses `Z2`: a single Majorana is not homogeneous under particle- +number `U1`, and generic pairing is not compatible with the charge-conserving +`U1U1` route. `Z2Z2` remains the natural future extension for two explicit +flavors with separately tracked parity, but it is not currently presented as +the default Majorana convention. + +The implementation does not introduce independent real Majorana sites or a +separate BdG/Nambu symmetry. Nambu doubling is useful as a quadratic-model +calculation basis, but here the physical representation remains one complex +mode per site and the native Symmray `Z2` graded algebra carries the signs. +This keeps operator construction, parity-preserving gates, mode ordering, and +2D fermionic sign validation on one representation path. + +Runnable versions of these workflows are in +`examples/qmera_fermion_hubbard_2d.py`, +`examples/qmera_fermion_hubbard_4x4_pbc.py`, and +`examples/qmera_majorana_2d.py`. For a first Torch energy comparison against +the U1U1 SymDMRG2 reference, see +`examples/qmera_fermion_hubbard_vs_symdmrg.py`. The comparison is variational: +the shallow qMERA energy is expected to remain above the better-converged DMRG +energy, while both calculations use the same physical Hubbard terms and +particle-number sector. + > API details are maintained as handwritten Markdown in this page. diff --git a/docs/examples.md b/docs/examples.md index a927f97..df60280 100644 --- a/docs/examples.md +++ b/docs/examples.md @@ -20,10 +20,20 @@ It demonstrates: - [How-To: Choose Parameters](howto/choose_parameters.md) - [How-To: Tune Sweep Solvers](howto/solver_tuning.md) +## qMERA examples + +- `examples/qmera_scale_plan_6x6.py` — generic heterogeneous 6x6 PBC RG + schedule with scale-specific blocks. +- `examples/qmera_fermion_hubbard_2d.py` — native `U1U1` 2D multimode + Fermi--Hubbard schedule, grouped cones, and direct-state validation. +- `examples/qmera_fermion_hubbard_4x4_pbc.py` — explicit 4x4 PBC square + disentangler/isometry RG schedule. +- `examples/qmera_majorana_2d.py` — native `Z2` Majorana and pairing gates. + ## Extended examples -The package repository keeps only the lightweight example notebook under -`examples/`. Larger runnable workflows and experiment scripts are maintained +The package repository keeps lightweight runnable examples under `examples/`. +Larger runnable workflows and experiment scripts are maintained in the separate `pepsy_examples` repository. The tests in this repository keep the deleted Relay-BP examples' numerical coverage without depending on local example files. diff --git a/examples/qmera_fermion_hubbard_2d.py b/examples/qmera_fermion_hubbard_2d.py new file mode 100644 index 0000000..37270f3 --- /dev/null +++ b/examples/qmera_fermion_hubbard_2d.py @@ -0,0 +1,106 @@ +"""Native Symmray 2D spinful Fermi--Hubbard qMERA example. + +This example keeps the physical lattice two-dimensional while using an +explicit register mode for each ``(site, spin)`` pair. The native +``U1U1`` path conserves up and down particle number independently. +""" + +import pepsy as py +from pepsy.optimizers.mera import ( + QMeraBuilder, + QMeraGeometry, + QMeraSymmrayFermionBackend, + symmray_fermion_gate_registry, +) + + +def main(): + try: + import symmray # noqa: F401 # pylint: disable=import-outside-toplevel + except ImportError as exc: # pragma: no cover - example dependency + raise SystemExit("Install Symmray to run this example.") from exc + + geometry = QMeraGeometry( + shape=(2, 2), + site_modes=("up", "down"), + mode_order="mode-major", + ) + backend = QMeraSymmrayFermionBackend( + symmetry="U1U1", + site_modes=("up", "down"), + mode_order="mode-major", + ) + registry = symmray_fermion_gate_registry(backend=backend) + + def product_state_factory(schedule, sites, **kwargs): + # One particle per site in a checkerboard spin pattern. + occupations = { + site: int( + (sum(schedule.geometry.to_site(site)) % 2 == 0) + == (schedule.geometry.to_mode(site)[1] == "up") + ) + for site in sites + } + return backend.product_state( + schedule, + sites, + occupations=occupations, + **kwargs, + ) + + builder = QMeraBuilder( + geometry=geometry, + gate_registry=registry, + gate_family="symmray-fsim", + disentangler={ + "block_size": 2, + "circuit_depth": 1, + "gate_family": "symmray-fsim", + }, + isometry={ + "block_size": (2, 2), + "circuit_depth": 1, + "gate_family": "symmray-fsim", + }, + max_layers=1, + seed=7, + param_scale=0.02, + product_state_factory=product_state_factory, + ) + fermion = py.Fermion( + spinful=True, + symmetry="U1U1", + t=0.2, + U=4.0, + mu=0.1, + ) + schedule = builder.build_schedule() + parameters = builder.initialize_parameters(schedule) + terms = builder.fermion_terms(fermion) + + # Symmray operators are already native; preserve them with + # convert_terms=False. The cache reuses paths for repeated local cones. + path_cache = builder.contraction_path_cache(max_repeats=8) + lightcone_energy = builder.parametric_loss( + parameters, + terms, + schedule=schedule, + convert_terms=False, + energy_per_site=False, + path_cache=path_cache, + ) + direct_energy = builder.direct_parametric_loss( + parameters, + terms, + schedule=schedule, + convert_terms=False, + energy_per_site=False, + path_cache=path_cache, + ) + print("2D U1U1 qMERA lightcone energy:", lightcone_energy) + print("2D U1U1 qMERA direct energy: ", direct_energy) + print("cached cone paths:", path_cache.num_cached_paths) + + +if __name__ == "__main__": + main() diff --git a/examples/qmera_fermion_hubbard_4x4_pbc.py b/examples/qmera_fermion_hubbard_4x4_pbc.py new file mode 100644 index 0000000..d5c9b6b --- /dev/null +++ b/examples/qmera_fermion_hubbard_4x4_pbc.py @@ -0,0 +1,73 @@ +"""Explicit 4x4 PBC spinful Fermi--Hubbard qMERA RG schedule. + +The first layer has four 2x2 covering blocks and square disentanglers on all +horizontal and vertical block interfaces, including the periodic wraps. The +second layer combines the four coarse sites into one 2x2 block. +""" + +import pepsy as py +from pepsy.optimizers.mera import ( + QMeraBuilder, + QMeraDisentanglerSpec, + QMeraGeometry, + QMeraIsometrySpec, + QMeraSymmrayFermionBackend, + QMeraUnitarySpec, + symmray_fermion_gate_registry, +) + + +def main(): + geometry = QMeraGeometry( + shape=(4, 4), + boundary="periodic", + site_modes=("up", "down"), + mode_order="mode-major", + ) + backend = QMeraSymmrayFermionBackend( + symmetry="U1U1", + site_modes=("up", "down"), + mode_order="mode-major", + ) + unitary = QMeraUnitarySpec( + gate_family="symmray-hubbard", + family="fermion", + arity_kind="mode", + symmetry="U1U1", + preserves_parity=True, + metadata={"model": "fermi-hubbard", "term": "hopping"}, + ) + builder = QMeraBuilder( + geometry=geometry, + gate_registry=symmray_fermion_gate_registry(backend=backend), + disentangler=QMeraDisentanglerSpec( + block_shape=(2, 2), + unitary=unitary, + placement="boundary-square", + circuit_depth=2, + periodic_wrap=True, + ), + isometry=QMeraIsometrySpec( + block_shape=(2, 2), + unitary=unitary, + circuit_depth=2, + implementation="unitary-completion", + ), + max_layers=2, + ) + schedule = builder.build_schedule() + + # The Fermion helper supplies the onsite U and chemical-potential terms; + # the qMERA unitary above supplies the native number-conserving hopping + # layer. Keeping these roles separate makes the U1U1 convention explicit. + fermion = py.Fermion(spinful=True, symmetry="U1U1", t=0.2, U=4.0, mu=0.1) + terms = builder.fermion_terms(fermion) + + print("RG register sizes:", [len(layer.input_sites) for layer in schedule.layers], "->", len(schedule.top_sites)) + print("first-layer isometry blocks:", len(schedule.layers[0].isometry_blocks)) + print("first-layer square disentanglers:", len(schedule.layers[0].disentangler_blocks)) + print("Fermi--Hubbard terms:", len(terms)) + + +if __name__ == "__main__": + main() diff --git a/examples/qmera_majorana_2d.py b/examples/qmera_majorana_2d.py new file mode 100644 index 0000000..ce06414 --- /dev/null +++ b/examples/qmera_majorana_2d.py @@ -0,0 +1,96 @@ +"""Native Z2 Majorana/pairing 2D qMERA example. + +The convention uses one complex spinless mode per physical site: +``gamma_x = c + c^dag`` and ``gamma_y = -i (c - c^dag)``. Individual +Majoranas are parity odd, while bilinears and pairing gates are Z2 neutral. +""" + +import pepsy as py +from pepsy.optimizers.mera import ( + QMeraBuilder, + QMeraGeometry, + QMeraSymmrayFermionBackend, + symmray_majorana_gate_registry, +) + + +def main(): + try: + import symmray # noqa: F401 # pylint: disable=import-outside-toplevel + except ImportError as exc: # pragma: no cover - example dependency + raise SystemExit("Install Symmray to run this example.") from exc + + geometry = QMeraGeometry( + shape=(2, 2), + site_modes=("mode",), + mode_order="mode-major", + ) + backend = QMeraSymmrayFermionBackend( + symmetry="Z2", + site_modes=("mode",), + ) + registry = symmray_majorana_gate_registry(backend=backend) + + def product_state_factory(schedule, sites, **kwargs): + occupations = { + site: sum(schedule.geometry.to_site(site)) % 2 + for site in sites + } + return backend.product_state( + schedule, + sites, + occupations=occupations, + **kwargs, + ) + + builder = QMeraBuilder( + geometry=geometry, + gate_registry=registry, + gate_family="symmray-majorana", + disentangler={ + "block_size": 2, + "circuit_depth": 1, + "gate_family": "symmray-majorana", + }, + isometry={ + "block_size": (2, 2), + "circuit_depth": 1, + "gate_family": "symmray-majorana", + }, + max_layers=1, + seed=11, + param_scale=0.02, + product_state_factory=product_state_factory, + ) + fermion = py.Fermion(spinful=False, symmetry="Z2") + schedule = builder.build_schedule() + parameters = builder.initialize_parameters(schedule) + terms = builder.majorana_terms( + fermion, + coupling=0.4, + pairing=0.2, + phase=0.1, + ) + + # Generic Majorana pairing is not a U1U1 charge-conserving operation, so + # this example deliberately uses the native parity-preserving Z2 route. + lightcone_energy = builder.parametric_loss( + parameters, + terms, + schedule=schedule, + convert_terms=False, + energy_per_site=False, + ) + direct_energy = builder.direct_parametric_loss( + parameters, + terms, + schedule=schedule, + convert_terms=False, + energy_per_site=False, + ) + print("2D Z2 Majorana lightcone energy:", lightcone_energy) + print("2D Z2 Majorana direct energy: ", direct_energy) + + +if __name__ == "__main__": + main() diff --git a/examples/qmera_scale_plan_6x6.py b/examples/qmera_scale_plan_6x6.py new file mode 100644 index 0000000..09858a9 --- /dev/null +++ b/examples/qmera_scale_plan_6x6.py @@ -0,0 +1,42 @@ +"""Generic heterogeneous 6x6 periodic qMERA scale-plan example.""" + +from pepsy.optimizers.mera import ( + QMeraBuilder, + QMeraDisentanglerSpec, + QMeraIsometrySpec, + QMeraScaleSpec, +) + + +def main(): + scales = ( + QMeraScaleSpec( + isometry=QMeraIsometrySpec(block_shape=(2, 2)), + disentangler=QMeraDisentanglerSpec( + block_shape=(2, 2), + placement="boundary-square", + ), + ), + QMeraScaleSpec( + isometry=QMeraIsometrySpec(block_shape=(3, 3)), + disentangler=QMeraDisentanglerSpec( + block_shape=3, + orientation="vertical", + placement="within-block", + circuit_depth=3, + ), + ), + ) + schedule = QMeraBuilder( + shape=(6, 6), + boundary="periodic", + scales=scales, + ).build_schedule() + + print("active sites:", [len(layer.input_sites) for layer in schedule.layers], "->", len(schedule.top_sites)) + print("isometry blocks:", [len(layer.isometry_blocks) for layer in schedule.layers]) + print("disentangler blocks:", [len(layer.disentangler_blocks) for layer in schedule.layers]) + + +if __name__ == "__main__": + main() diff --git a/src/pepsy/optimizers/__init__.py b/src/pepsy/optimizers/__init__.py index d451239..2480706 100644 --- a/src/pepsy/optimizers/__init__.py +++ b/src/pepsy/optimizers/__init__.py @@ -15,8 +15,12 @@ "GlobalOptimizer": ".global_opt", "MeraEnergyOptimizer": ".mera", "QMeraBuilder": ".mera", + "QMeraDisentanglerSpec": ".mera", "QMeraGeometry": ".mera", + "QMeraIsometrySpec": ".mera", "QMeraParametricEnergyOptimizer": ".mera", + "QMeraScaleSpec": ".mera", + "QMeraUnitarySpec": ".mera", "build_qmera_contraction_optimizer": ".mera", "MpoOptimizer": ".mpo", "MpsOptimizer": ".mps", diff --git a/src/pepsy/optimizers/mera/__init__.py b/src/pepsy/optimizers/mera/__init__.py index a6a1815..080cf45 100644 --- a/src/pepsy/optimizers/mera/__init__.py +++ b/src/pepsy/optimizers/mera/__init__.py @@ -1,7 +1,7 @@ """MERA and qMERA energy optimization helpers.""" from .builders import QMeraAnsatz, QMeraBuilder -from .cache import build_qmera_contraction_optimizer +from .cache import QMeraContractionPathCache, build_qmera_contraction_optimizer from .compiled import ( QMeraCompiledLightconeChunk, compile_qmera_parametric_lightcone, @@ -12,7 +12,9 @@ from .fermions import ( QMeraSymmrayFermionBackend, qmera_symmray_fermi_hubbard_terms, + qmera_symmray_majorana_terms, symmray_fermion_gate_registry, + symmray_majorana_gate_registry, ) from .gates import ( GateRegistry, @@ -25,14 +27,21 @@ from .lightcones import ( LightconeChunk, QMeraLightconeTN, + QMeraLightconeGroup, QMeraParametricLightconeChunk, build_lightcone_chunks, build_qmera_lightcone_chunks, build_qmera_parametric_lightcone_chunks, + contract_qmera_lightcone_group, contract_qmera_lightcone_tn, + group_qmera_parametric_lightcone_chunks, + lightcone_energy, local_qmera_parametric_lightcone_expectation, local_lightcone_expectation, qmera_parametric_energy, + qmera_direct_parametric_energy, + qmera_parametric_lightcone_group_state, + qmera_parametric_state, qmera_parametric_lightcone_state, qmera_parametric_lightcone_tn, select_lightcone, @@ -42,9 +51,13 @@ from .parametric import QMeraParametricEnergyOptimizer from .schedules import ( QMeraBlockSpec, + QMeraDisentanglerSpec, QMeraGatePlacement, + QMeraIsometrySpec, QMeraLayerSpec, + QMeraScaleSpec, QMeraSchedule, + QMeraUnitarySpec, build_qmera_schedule, ) from .schematics import ( @@ -64,15 +77,21 @@ "QMeraBlockSpec", "QMeraBuilder", "QMeraCompiledLightconeChunk", + "QMeraContractionPathCache", + "QMeraDisentanglerSpec", "QMeraGatePlacement", "QMeraGeometry", + "QMeraIsometrySpec", "QMeraLayerSpec", "QMeraLightconeTN", + "QMeraLightconeGroup", "QMeraParametricLightconeChunk", "QMeraParametricEnergyOptimizer", "QMeraSchedule", + "QMeraScaleSpec", "QMeraSchematicBlock", "QMeraSymmrayFermionBackend", + "QMeraUnitarySpec", "UserGateFamily", "build_lightcone_chunks", "build_qmera_contraction_optimizer", @@ -82,6 +101,7 @@ "compile_qmera_parametric_lightcone", "compile_qmera_parametric_lightcones", "contract_qmera_lightcone_tn", + "contract_qmera_lightcone_group", "default_gate_registry", "draw_qmera_schedule", "local_qmera_compiled_lightcone_expectation", @@ -89,13 +109,20 @@ "local_lightcone_expectation", "normalize_local_terms", "qmera_compiled_parametric_energy", + "qmera_direct_parametric_energy", "qmera_parametric_energy", + "qmera_parametric_lightcone_group_state", "qmera_parametric_lightcone_state", + "qmera_parametric_state", "qmera_parametric_lightcone_tn", "qmera_schematic_blocks", "qmera_symmray_fermi_hubbard_terms", + "qmera_symmray_majorana_terms", "resolve_gate_spec", + "group_qmera_parametric_lightcone_chunks", + "lightcone_energy", "select_lightcone", "site_tags_for_where", "symmray_fermion_gate_registry", + "symmray_majorana_gate_registry", ] diff --git a/src/pepsy/optimizers/mera/builders.py b/src/pepsy/optimizers/mera/builders.py index 043ad92..450a3d1 100644 --- a/src/pepsy/optimizers/mera/builders.py +++ b/src/pepsy/optimizers/mera/builders.py @@ -25,10 +25,20 @@ from .geometry import QMeraGeometry from .lightcones import ( build_qmera_parametric_lightcone_chunks, + group_qmera_parametric_lightcone_chunks, + qmera_direct_parametric_energy, qmera_parametric_energy, qmera_parametric_lightcone_tn, ) -from .schedules import QMeraBlockSpec, QMeraSchedule, build_qmera_schedule +from .schedules import ( + QMeraBlockSpec, + QMeraDisentanglerSpec, + QMeraIsometrySpec, + QMeraScaleSpec, + QMeraSchedule, + QMeraUnitarySpec, + build_qmera_schedule, +) from .terms import convert_local_terms, normalize_local_terms __all__ = ["QMeraAnsatz", "QMeraBuilder"] @@ -84,16 +94,32 @@ def _coerce_geometry( boundary="open", mapper=None, site_modes=None, - mode_order="site-major", + mode_order=None, ): if isinstance(geometry, QMeraGeometry): - return geometry + # A model-aware builder may need to add the explicit local modes to a + # geometry that was created as a plain spatial lattice. Preserve an + # already explicit geometry unless the caller supplied an override. + requested_modes = geometry.site_modes if site_modes is None else site_modes + requested_order = geometry.mode_order if mode_order is None else mode_order + if ( + tuple(requested_modes or ()) == tuple(geometry.site_modes or ()) + and requested_order == geometry.mode_order + ): + return geometry + return QMeraGeometry( + geometry.shape, + boundary=geometry.boundary, + site_labels=geometry.site_labels, + site_modes=requested_modes, + mode_order=requested_order, + ) if geometry is not None: if isinstance(geometry, Mapping): opts = dict(geometry) if site_modes is not None: opts.setdefault("site_modes", site_modes) - if mode_order != "site-major": + if mode_order is not None: opts.setdefault("mode_order", mode_order) return QMeraGeometry(**opts) if shape is not None: @@ -103,7 +129,7 @@ def _coerce_geometry( boundary=boundary, mapper=mapper, site_modes=site_modes, - mode_order=mode_order, + mode_order="site-major" if mode_order is None else mode_order, ) if shape is None: raise TypeError("QMeraBuilder requires geometry or shape.") @@ -112,20 +138,49 @@ def _coerce_geometry( boundary=boundary, mapper=mapper, site_modes=site_modes, - mode_order=mode_order, + mode_order="site-major" if mode_order is None else mode_order, ) +def _infer_fermion_site_modes(fermion): + """Infer qMERA's explicit mode labels from a ``Fermion`` helper.""" + if getattr(fermion, "spinful", False): + return ("up", "down") + return ("mode",) + + def _coerce_block_spec(value, *, kind, gate_family): + if kind == "disentangler" and isinstance(value, QMeraDisentanglerSpec): + return value.to_block_spec(default_gate_family=gate_family) + if kind == "isometry" and isinstance(value, QMeraIsometrySpec): + return value.to_block_spec(default_gate_family=gate_family) if isinstance(value, QMeraBlockSpec): if value.kind != kind: raise ValueError(f"{kind} block spec has kind={value.kind!r}.") return value opts = dict(value or {}) + unitary = opts.pop("unitary", None) + if unitary is not None: + unitary = QMeraUnitarySpec.coerce( + unitary, + default_gate_family=gate_family, + ) + opts.setdefault("gate_family", unitary.gate_family) + opts.setdefault("unitary_spec", unitary) + elif "unitary_spec" in opts and opts["unitary_spec"] is not None: + opts["unitary_spec"] = QMeraUnitarySpec.coerce( + opts["unitary_spec"], + default_gate_family=opts.get("gate_family", gate_family), + ) opts.setdefault("gate_family", gate_family) return QMeraBlockSpec(kind=kind, **opts) +def _normalize_gate_token(value): + key = str(value).strip().lower().replace("_", "-") + return {"fermionic": "fermion", "qubit": "spin"}.get(key, key) + + class QMeraBuilder: """Build a schedule-first qMERA ansatz from explicit Pepsy objects.""" @@ -137,10 +192,12 @@ def __init__( boundary="open", mapper=None, site_modes=None, - mode_order="site-major", + mode_order=None, + fermion=None, physical_dim: int = 2, disentangler=None, isometry=None, + scales=None, gate_family: str = "rxx", isometry_gate_family: str | None = None, gate_registry: GateRegistry | None = None, @@ -152,14 +209,55 @@ def __init__( parameter_backend=None, product_state_factory=None, ): + self.fermion = fermion + if fermion is not None and not callable( + getattr(fermion, "local_terms", None) + ): + raise TypeError("fermion must provide local_terms(...).") + + inferred_modes = ( + _infer_fermion_site_modes(fermion) + if fermion is not None + else site_modes + ) + requested_modes = ( + site_modes + if site_modes is not None + else ( + geometry.site_modes + if isinstance(geometry, QMeraGeometry) + and geometry.site_modes is not None + else inferred_modes + ) + ) + # A mode-major register is the natural default for a spinful qMERA: + # each spatial layer contains one complete up/down mode register. The + # generic builder retains its historical site-major default. + inferred_order = ( + mode_order + if mode_order is not None + else ( + geometry.mode_order + if isinstance(geometry, QMeraGeometry) + else ("mode-major" if fermion is not None else None) + ) + ) self.geometry = _coerce_geometry( geometry, shape=shape, boundary=boundary, mapper=mapper, - site_modes=site_modes, - mode_order=mode_order, + site_modes=requested_modes, + mode_order=inferred_order, ) + if fermion is not None: + expected_modes = tuple(_infer_fermion_site_modes(fermion)) + actual_modes = tuple(self.geometry.site_modes or ()) + if actual_modes != expected_modes: + raise ValueError( + "The qMERA geometry's site_modes must match the Fermion " + f"helper: expected {expected_modes!r}, got {actual_modes!r}." + ) self.physical_dim = int(physical_dim) if self.physical_dim != 2: raise NotImplementedError("QMeraBuilder currently supports qubits only.") @@ -178,6 +276,17 @@ def __init__( kind="isometry", gate_family=isometry_gate_family or gate_family, ) + self._validate_unitary_spec(self.disentangler) + self._validate_unitary_spec(self.isometry) + if scales is not None: + self.scales = tuple(scales) + for scale in self.scales: + if not isinstance(scale, (QMeraScaleSpec, Mapping)): + raise TypeError( + "scales must contain QMeraScaleSpec or mapping objects." + ) + else: + self.scales = None self.max_layers = max_layers self.top_size = top_size self.seed = seed @@ -186,15 +295,54 @@ def __init__( self.parameter_backend = parameter_backend self.product_state_factory = product_state_factory + def _validate_unitary_spec(self, block_spec): + """Validate explicit unitary metadata against the selected registry.""" + unitary = block_spec.unitary_spec + if unitary is None: + return + if _normalize_gate_token(unitary.gate_family) != _normalize_gate_token( + block_spec.gate_family + ): + raise ValueError( + f"{block_spec.kind} unitary gate_family={unitary.gate_family!r} " + f"does not match block gate_family={block_spec.gate_family!r}." + ) + spec = resolve_gate_spec(block_spec.gate_family, self.gate_registry) + checks = ( + ("family", unitary.family, spec.family), + ("arity_kind", unitary.arity_kind, spec.arity_kind), + ("symmetry", unitary.symmetry, getattr(spec, "symmetry", None)), + ("preserves_parity", unitary.preserves_parity, spec.preserves_parity), + ) + for name, expected, actual in checks: + if expected is None: + continue + if name in {"family", "arity_kind"}: + expected = _normalize_gate_token(expected) + actual = _normalize_gate_token(actual) + elif name == "symmetry": + expected = str(expected).upper() + actual = None if actual is None else str(actual).upper() + if actual != expected: + raise ValueError( + f"{block_spec.kind} unitary requires {name}={expected!r}, " + f"but gate family {spec.name!r} provides {actual!r}." + ) + def build_schedule(self): """Build the static qMERA schedule.""" - return build_qmera_schedule( + schedule = build_qmera_schedule( self.geometry, disentangler=self.disentangler, isometry=self.isometry, + scales=self.scales, max_layers=self.max_layers, top_size=self.top_size, ) + for scale in schedule.scale_specs: + self._validate_unitary_spec(scale.disentangler) + self._validate_unitary_spec(scale.isometry) + return schedule def schematic_blocks(self, *, layer=None): """Return display-oriented disentangler/isometry blocks.""" @@ -208,6 +356,12 @@ def contraction_optimizer(self, **kwargs): """Build a reusable contraction optimizer for repeated local cones.""" return build_qmera_contraction_optimizer(**kwargs) + def contraction_path_cache(self, **kwargs): + """Create a lazy topology-aware contraction-path cache.""" + from .cache import QMeraContractionPathCache + + return QMeraContractionPathCache(optimizer_options=kwargs) + def parametric_lightcone_chunks( self, hamiltonian, @@ -226,14 +380,31 @@ def parametric_lightcone_chunks( terms = convert_local_terms(terms, backend) return build_qmera_parametric_lightcone_chunks(schedule, terms) - def fermion_terms(self, fermion, **params): + def parametric_lightcone_groups(self, hamiltonian, schedule=None, **kwargs): + """Group local terms sharing one reverse-lightcone topology.""" + chunks = self.parametric_lightcone_chunks( + hamiltonian, + schedule=schedule, + **kwargs, + ) + return group_qmera_parametric_lightcone_chunks(chunks) + + def fermion_terms(self, fermion=None, **params): """Return qMERA mode terms from a unified :class:`Fermion` helper. qMERA's fermionic path represents each physical site as the explicit ``("up", "down")`` pair of two-state modes. Keeping this conversion - on the builder makes the representation choice visible while allowing - the regular local-term and optimizer machinery to handle the result. + on the builder keeps the representation choice in one place while + allowing the regular local-term and optimizer machinery to handle the + result. If the builder was created with ``fermion=...``, the argument + can be omitted. """ + fermion = self.fermion if fermion is None else fermion + if fermion is None: + raise TypeError( + "Provide fermion=... to QMeraBuilder or pass a Fermion helper " + "to fermion_terms(...)." + ) if not callable(getattr(fermion, "local_terms", None)): raise TypeError("fermion must provide local_terms(...).") if not getattr(fermion, "spinful", False): @@ -251,16 +422,30 @@ def fermion_terms(self, fermion, **params): ) ) + def majorana_terms(self, fermion=None, **params): + """Return native parity-preserving Majorana terms for this geometry.""" + fermion = self.fermion if fermion is None else fermion + if fermion is None: + raise TypeError( + "Provide fermion=... to QMeraBuilder or pass a Fermion helper " + "to majorana_terms(...)." + ) + if not callable(getattr(fermion, "majorana_terms", None)): + raise TypeError("fermion must provide majorana_terms(...).") + return tuple(fermion.majorana_terms(self.geometry, **params)) + def fermion_parametric_loss( self, - fermion, - parameters, + fermion=None, + parameters=None, schedule=None, *, term_params=None, **loss_kwargs, ): """Evaluate qMERA energy directly from a unified ``Fermion`` model.""" + if parameters is None: + raise TypeError("parameters must be supplied for fermion_parametric_loss.") schedule = self.build_schedule() if schedule is None else schedule terms = self.fermion_terms(fermion, **dict(term_params or {})) return self.parametric_loss( @@ -272,7 +457,7 @@ def fermion_parametric_loss( def fermion_parametric_optimizer( self, - fermion, + fermion=None, *, schedule=None, term_params=None, @@ -287,6 +472,44 @@ def fermion_parametric_optimizer( **optimizer_kwargs, ) + def majorana_parametric_loss( + self, + fermion=None, + parameters=None, + schedule=None, + *, + term_params=None, + **loss_kwargs, + ): + """Evaluate a native ``Z2`` Majorana qMERA energy.""" + if parameters is None: + raise TypeError("parameters must be supplied for majorana_parametric_loss.") + schedule = self.build_schedule() if schedule is None else schedule + terms = self.majorana_terms(fermion, **dict(term_params or {})) + return self.parametric_loss( + parameters, + terms, + schedule=schedule, + **loss_kwargs, + ) + + def majorana_parametric_optimizer( + self, + fermion=None, + *, + schedule=None, + term_params=None, + **optimizer_kwargs, + ): + """Create a qMERA optimizer for native ``Z2`` Majorana terms.""" + schedule = self.build_schedule() if schedule is None else schedule + terms = self.majorana_terms(fermion, **dict(term_params or {})) + return self.parametric_optimizer( + terms, + schedule=schedule, + **optimizer_kwargs, + ) + def parametric_loss( self, parameters, @@ -304,6 +527,8 @@ def parametric_loss( simplify=False, gate_contract=True, contract_opts=None, + group_terms=True, + path_cache=None, ): """Evaluate qMERA energy from params by rebuilding local cones only.""" schedule = self.build_schedule() if schedule is None else schedule @@ -326,6 +551,49 @@ def parametric_loss( gate_contract=gate_contract, contract_opts=contract_opts, product_state_factory=self.product_state_factory, + group_terms=group_terms, + path_cache=path_cache, + ) + + def direct_parametric_loss( + self, + parameters, + hamiltonian=None, + schedule=None, + *, + chunks=None, + array_backend=None, + gate_array_backend=None, + convert_terms=True, + normalized=True, + energy_per_site=True, + real=True, + contraction_opt="auto-hq", + contract_opts=None, + group_terms=True, + path_cache=None, + ): + """Evaluate the full direct-gate TN as a validation oracle.""" + schedule = self.build_schedule() if schedule is None else schedule + backend = self.array_backend if array_backend is None else array_backend + return qmera_direct_parametric_energy( + schedule, + parameters, + hamiltonian, + chunks=chunks, + gate_registry=self.gate_registry, + array_backend=backend, + gate_array_backend=gate_array_backend, + convert_terms=convert_terms, + physical_dim=self.physical_dim, + optimize=contraction_opt, + normalized=normalized, + energy_per_site=energy_per_site, + real=real, + contract_opts=contract_opts, + group_terms=group_terms, + path_cache=path_cache, + product_state_factory=self.product_state_factory, ) def _parameter_converter(self): @@ -339,13 +607,15 @@ def _parameter_converter(self): return get_default_array_backend() def initialize_parameters(self, schedule=None, *, seed=None, scale=None): - """Initialize one parameter vector per scheduled gate.""" + """Initialize one parameter vector per unique sharing key.""" schedule = self.build_schedule() if schedule is None else schedule rng = np.random.default_rng(self.seed if seed is None else seed) scale = self.param_scale if scale is None else float(scale) converter = self._parameter_converter() params = {} for placement in schedule.placements: + if placement.param_key in params: + continue spec = resolve_gate_spec(placement.gate_family, self.gate_registry) if spec.num_params == 0: values = np.empty((0,), dtype=np.float64) diff --git a/src/pepsy/optimizers/mera/cache.py b/src/pepsy/optimizers/mera/cache.py index b00ba41..c497c6c 100644 --- a/src/pepsy/optimizers/mera/cache.py +++ b/src/pepsy/optimizers/mera/cache.py @@ -2,11 +2,51 @@ from __future__ import annotations -from typing import Any +from dataclasses import dataclass, field +from typing import Any, Mapping from ...tensors import build_optimizer -__all__ = ["build_qmera_contraction_optimizer"] +__all__ = [ + "QMeraContractionPathCache", + "build_qmera_contraction_optimizer", +] + + +@dataclass +class QMeraContractionPathCache: + """Lazily build reusable cotengra optimizers per local-cone topology. + + A qMERA schedule produces a small number of repeated cone topologies. The + cache keeps one :class:`cotengra.ReusableHyperOptimizer` per topology key, + so numerator and norm contractions for repeated terms reuse the same + searched paths. Passing a directory in ``optimizer_options`` additionally + persists cotengra's reusable paths across processes. + """ + + optimizer_options: Mapping[str, Any] = field(default_factory=dict) + _optimizers: dict[Any, Any] = field(default_factory=dict, init=False, repr=False) + + def optimizer_for(self, key=None): + """Return the reusable optimizer associated with ``key``.""" + key = "default" if key is None else key + try: + return self._optimizers[key] + except KeyError: + optimizer = build_qmera_contraction_optimizer(**dict(self.optimizer_options)) + self._optimizers[key] = optimizer + return optimizer + + @property + def num_cached_paths(self): + """Number of topology-specific reusable optimizers created so far.""" + return len(self._optimizers) + + def resolve(self, optimize, *, key=None): + """Resolve an ``optimize`` setting, reusing paths for auto settings.""" + if optimize is None or str(optimize).lower() in {"auto", "auto-hq"}: + return self.optimizer_for(key) + return optimize def build_qmera_contraction_optimizer( diff --git a/src/pepsy/optimizers/mera/fermions.py b/src/pepsy/optimizers/mera/fermions.py index 9e7f9ab..cd6a024 100644 --- a/src/pepsy/optimizers/mera/fermions.py +++ b/src/pepsy/optimizers/mera/fermions.py @@ -16,7 +16,9 @@ __all__ = [ "QMeraSymmrayFermionBackend", "qmera_symmray_fermi_hubbard_terms", + "qmera_symmray_majorana_terms", "symmray_fermion_gate_registry", + "symmray_majorana_gate_registry", ] @@ -102,9 +104,28 @@ class QMeraSymmrayFermionBackend: to_backend: Any = None flat: bool = False mode_order: str = "site-major" - zero_charge: Any = (0, 0) + zero_charge: Any = None _mode_charges: dict[Any, tuple[Any, Any]] = field(init=False, repr=False) + @classmethod + def from_fermion(cls, fermion, **kwargs): + """Construct the mode backend from a unified :class:`Fermion` model. + + Spinful models use the canonical ``("up", "down")`` pair while a + spinless model uses one ``"mode"`` register. Callers can still + override backend-specific options, including ``mode_order``. + """ + spinful = bool(getattr(fermion, "spinful", False)) + site_modes = ("up", "down") if spinful else ("mode",) + kwargs.setdefault("symmetry", getattr(fermion, "symmetry", "U1U1")) + kwargs.setdefault("site_modes", site_modes) + kwargs.setdefault("mode_order", "mode-major") + if "dtype" not in kwargs and hasattr(fermion, "dtype"): + kwargs["dtype"] = fermion.dtype + if "to_backend" not in kwargs and hasattr(fermion, "to_backend"): + kwargs["to_backend"] = fermion.to_backend + return cls(**kwargs) + def __post_init__(self): modes = tuple(self.site_modes) if not modes: @@ -121,6 +142,9 @@ def __post_init__(self): "_mode_charges", {mode: tuple(charges[mode]) for mode in modes}, ) + if self.zero_charge is None: + zero_charge = (0, 0) if self.symmetry in {"U1U1", "Z2Z2"} else 0 + object.__setattr__(self, "zero_charge", zero_charge) object.__setattr__(self, "dtype", np.dtype(self.dtype)) @staticmethod @@ -472,6 +496,97 @@ def qmera_symmray_fermi_hubbard_terms(geometry, *, fermion=None, **kwargs): return backend.fermi_hubbard_terms(geometry, **kwargs) +def qmera_symmray_majorana_terms(geometry, *, fermion=None, **kwargs): + """Return native parity-preserving Majorana terms for qMERA. + + The first qMERA Majorana convention is one spinless complex mode per + lattice site, represented with ``site_modes=("mode",)`` and conserved + ``Z2`` fermion parity. The physical Majoranas are the two quadratures of + each complex mode; they are not separate Hilbert-space sites. + + Parameters + ---------- + geometry : QMeraGeometry + A 1D or 2D geometry with one explicit ``"mode"`` per site. + fermion : pepsy.Fermion, optional + A spinless ``Fermion(symmetry="Z2")`` helper. + coupling : scalar, mapping, or callable, optional + Coefficient of ``i gamma_{j,y} gamma_{k,x}`` on each nearest-neighbor + edge. + pairing : scalar, mapping, or callable, optional + Coefficient of the Hermitian ``c_j^dag c_k^dag + h.c.`` term. + """ + if tuple(geometry.site_modes or ()) != ("mode",): + raise ValueError( + "Majorana qMERA requires geometry site_modes=('mode',), one " + "complex fermion mode per physical site." + ) + if fermion is None: + from ...tensors import Fermion # pylint: disable=import-outside-toplevel + + fermion = Fermion(spinful=False, symmetry="Z2") + if getattr(fermion, "spinful", True): + raise ValueError("Majorana qMERA requires a spinless Fermion helper.") + if str(getattr(fermion, "symmetry", "")) != "Z2": + raise ValueError( + "Majorana qMERA currently uses the native Z2 parity convention." + ) + + coupling = kwargs.pop("coupling", kwargs.pop("t", 1.0)) + pairing = kwargs.pop("pairing", 0.0) + left_component = kwargs.pop("left_component", 1) + right_component = kwargs.pop("right_component", 0) + phase = kwargs.pop("phase", 0.0) + if kwargs: + unknown = ", ".join(sorted(kwargs)) + raise TypeError(f"Unknown Majorana qMERA term option(s): {unknown}.") + + terms = [] + for left, right in geometry.nearest_neighbor_edges(): + left_mode = geometry.mode_label(left, "mode") + right_mode = geometry.mode_label(right, "mode") + coupling_edge = _edge_parameter(coupling, left, right) + if coupling_edge != 0: + terms.append( + LocalTerm( + where=(left_mode, right_mode), + operator=fermion.majorana_bilinear_operator( + (left_mode, right_mode), + left_component=left_component, + right_component=right_component, + coefficient=coupling_edge, + ), + metadata={ + "kind": "majorana-bilinear", + "edge": (left, right), + "fermionic": True, + "symmetry": "Z2", + "convention": "i-gamma-y-gamma-x", + }, + ) + ) + pairing_edge = _edge_parameter(pairing, left, right) + if pairing_edge != 0: + terms.append( + LocalTerm( + where=(left_mode, right_mode), + operator=fermion.pairing_operator( + (left_mode, right_mode), + coefficient=pairing_edge, + phase=_edge_angle_parameter(phase, left, right), + ), + metadata={ + "kind": "majorana-pairing", + "edge": (left, right), + "fermionic": True, + "symmetry": "Z2", + "convention": "creation-pair-plus-hc", + }, + ) + ) + return tuple(terms) + + def _null_context_gate(_params): raise ValueError("This qMERA gate family requires placement context.") @@ -488,26 +603,116 @@ def symmray_fermion_gate_registry(backend=None, *, base_registry=None): def fsim_context(params, *, placement=None, schedule=None, array_backend=None): _ = array_backend if placement is None or schedule is None: - raise ValueError("symmray-fsim requires qMERA placement and schedule.") + raise ValueError( + "native Symmray fermion gates require qMERA placement and schedule." + ) left, right = ( schedule.geometry.to_mode(register_site) for register_site in placement.where ) return backend.fsim_gate(left, right, theta=params[0], phi=params[1]) + common = dict( + arity=2, + num_params=2, + generator=_null_context_gate, + family="fermion", + arity_kind="mode", + preserves_parity=True, + mode_order="register", + symmetry=backend.symmetry, + contextual_generator=fsim_context, + ) registry.register( GateSpec( "symmray-fsim", + convention="symmray-fermionic-mode", + default_tags=("SYMMRAY_FSIM", "FSIM"), + **common, + ) + ) + registry.register( + GateSpec( + "symmray-hubbard", + convention="symmray-fermionic-hubbard-hopping", + default_tags=("SYMMRAY_HUBBARD", "HOPPING"), + **common, + ) + ) + return registry + + +def symmray_majorana_gate_registry(backend=None, *, base_registry=None): + """Return a native ``Z2`` parity-preserving Majorana gate registry.""" + from ...tensors import Fermion # pylint: disable=import-outside-toplevel + + backend = ( + QMeraSymmrayFermionBackend(symmetry="Z2", site_modes=("mode",)) + if backend is None + else backend + ) + if str(backend.symmetry) != "Z2" or tuple(backend.site_modes) != ("mode",): + raise ValueError( + "The Majorana gate registry requires a Z2 backend with " + "site_modes=('mode',)." + ) + fermion = Fermion( + spinful=False, + symmetry="Z2", + dtype=backend.dtype, + to_backend=backend.to_backend, + ) + registry = default_gate_registry() if base_registry is None else base_registry.copy() + + def _mode_pair(placement, schedule): + return tuple(schedule.geometry.to_mode(site) for site in placement.where) + + def majorana_context(params, *, placement=None, schedule=None, array_backend=None): + _ = array_backend + if placement is None or schedule is None: + raise ValueError("symmray-majorana requires qMERA placement and schedule.") + left, right = _mode_pair(placement, schedule) + return fermion.majorana_gate( + params[0], + edge=(left, right), + left_component=1, + right_component=0, + ) + + def pairing_context(params, *, placement=None, schedule=None, array_backend=None): + _ = array_backend + if placement is None or schedule is None: + raise ValueError("symmray-pairing requires qMERA placement and schedule.") + left, right = _mode_pair(placement, schedule) + return fermion.pairing_gate(params[0], edge=(left, right)) + + common = dict( + family="fermion", + convention="symmray-z2-majorana", + arity_kind="mode", + preserves_parity=True, + mode_order="register", + ) + registry.register( + GateSpec( + "symmray-majorana", 2, + 1, + _null_context_gate, + default_tags=("MAJORANA", "PARITY"), + contextual_generator=majorana_context, + **common, + ) + ) + registry.register( + GateSpec( + "symmray-pairing", 2, + 1, _null_context_gate, - family="fermion", - convention="symmray-fermionic-mode", - default_tags=("SYMMRAY_FSIM", "FSIM"), - arity_kind="mode", - preserves_parity=True, - mode_order="register", - contextual_generator=fsim_context, + default_tags=("PAIRING", "PARITY"), + contextual_generator=pairing_context, + **common, ) ) return registry diff --git a/src/pepsy/optimizers/mera/gates.py b/src/pepsy/optimizers/mera/gates.py index a544077..c831919 100644 --- a/src/pepsy/optimizers/mera/gates.py +++ b/src/pepsy/optimizers/mera/gates.py @@ -89,6 +89,7 @@ class GateSpec: preserves_parity: bool | None = None mode_order: str | None = None contextual_generator: Callable[..., Any] | None = None + symmetry: str | None = None def __post_init__(self): arity = int(self.arity) @@ -111,6 +112,11 @@ def __post_init__(self): None if self.preserves_parity is None else bool(self.preserves_parity), ) object.__setattr__(self, "mode_order", _normalize_mode_order(self.mode_order)) + object.__setattr__( + self, + "symmetry", + None if self.symmetry is None else str(self.symmetry), + ) @property def is_fermionic(self): @@ -166,6 +172,7 @@ class UserGateFamily: preserves_parity: bool | None = None mode_order: str | None = None contextual_generator: Callable[..., Any] | None = None + symmetry: str | None = None def to_gate_spec(self): """Convert to a registry-ready :class:`GateSpec`.""" @@ -181,6 +188,7 @@ def to_gate_spec(self): preserves_parity=self.preserves_parity, mode_order=self.mode_order, contextual_generator=self.contextual_generator, + symmetry=self.symmetry, ) diff --git a/src/pepsy/optimizers/mera/lightcones.py b/src/pepsy/optimizers/mera/lightcones.py index ec77106..a3792d6 100644 --- a/src/pepsy/optimizers/mera/lightcones.py +++ b/src/pepsy/optimizers/mera/lightcones.py @@ -14,12 +14,19 @@ __all__ = [ "LightconeChunk", + "QMeraLightconeGroup", "QMeraLightconeTN", "QMeraParametricLightconeChunk", "build_lightcone_chunks", "build_qmera_lightcone_chunks", "build_qmera_parametric_lightcone_chunks", + "group_qmera_parametric_lightcone_chunks", "contract_qmera_lightcone_tn", + "contract_qmera_lightcone_group", + "lightcone_energy", + "qmera_direct_parametric_energy", + "qmera_parametric_state", + "qmera_parametric_lightcone_group_state", "local_qmera_parametric_lightcone_expectation", "local_lightcone_expectation", "qmera_parametric_energy", @@ -96,6 +103,26 @@ def num_gates(self): return len(self.schedule_placement_ids) +@dataclass(frozen=True) +class QMeraLightconeGroup: + """Local terms sharing one qMERA reverse-lightcone topology.""" + + key: tuple + chunks: tuple[QMeraParametricLightconeChunk, ...] + input_sites: tuple[int, ...] + schedule_placement_ids: tuple[str, ...] + + @property + def num_terms(self): + """Number of local Hamiltonian terms in this group.""" + return len(self.chunks) + + @property + def num_gates(self): + """Number of gates in the shared local cone.""" + return len(self.schedule_placement_ids) + + @dataclass(frozen=True) class QMeraLightconeTN: """Explicit tensor networks for one scheduled qMERA local term.""" @@ -300,6 +327,23 @@ def build_qmera_parametric_lightcone_chunks(schedule, terms): ) +def group_qmera_parametric_lightcone_chunks(chunks): + """Group qMERA chunks that share input support and gate topology.""" + groups = {} + for chunk in tuple(chunks): + key = (tuple(chunk.input_sites), tuple(chunk.schedule_placement_ids)) + groups.setdefault(key, []).append(chunk) + return tuple( + QMeraLightconeGroup( + key=key, + chunks=tuple(group_chunks), + input_sites=key[0], + schedule_placement_ids=key[1], + ) + for key, group_chunks in groups.items() + ) + + def _maybe_real(value): try: return ar.do("real", value) @@ -322,6 +366,27 @@ def _site_ind(site): return f"k{site}" +def _apply_local_gate(state, operator, where, *, contract, inplace): + """Apply a local operator through the state's native TN gate API.""" + gate = getattr(state, "gate", None) + if callable(gate): + return gate( + operator, + where, + contract=contract, + inplace=inplace, + ) + gate_inds = getattr(state, "gate_inds", None) + if callable(gate_inds): + return gate_inds( + operator, + inds=_site_inds_for_where(state, where), + contract=contract, + inplace=inplace, + ) + raise TypeError("state must provide gate() or gate_inds().") + + def _product_state_on_sites(sites, *, physical_dim=2, array_backend=None): tensors = [] base = np.zeros((int(physical_dim),), dtype=np.complex128) @@ -363,6 +428,15 @@ def _placements_by_id(schedule): return by_id if isinstance(by_id, dict) else dict(by_id) +def _resolve_contraction_opt(optimize, path_cache, *, key): + if path_cache is None: + return optimize + resolver = getattr(path_cache, "resolve", None) + if not callable(resolver): + raise TypeError("path_cache must provide resolve(optimize, key=...).") + return resolver(optimize, key=key) + + def _gate_for_placement( placement, parameters, @@ -392,6 +466,46 @@ def _gate_for_placement( ) +def qmera_parametric_state( + schedule, + parameters, + *, + gate_registry=None, + array_backend=None, + gate_array_backend=None, + physical_dim=2, + contract=False, + product_state_factory=None, +): + """Build the complete direct-gate qMERA state from a parameter map.""" + gate_registry = default_gate_registry() if gate_registry is None else gate_registry + array_backend = _resolve_array_backend(array_backend) + placements = tuple(schedule.placements) + state = _product_state_for_schedule( + schedule, + schedule.geometry.register_sites, + physical_dim=physical_dim, + array_backend=array_backend, + product_state_factory=product_state_factory, + ) + for placement in placements: + gate = _gate_for_placement( + placement, + parameters, + gate_registry=gate_registry, + array_backend=gate_array_backend, + schedule=schedule, + ) + state = state.gate_inds( + gate, + inds=tuple(_site_ind(site) for site in placement.where), + contract=contract, + tags=placement.tags, + inplace=False, + ) + return state + + def qmera_parametric_lightcone_state( schedule, chunk: QMeraParametricLightconeChunk, @@ -434,6 +548,39 @@ def qmera_parametric_lightcone_state( return state +def qmera_parametric_lightcone_group_state( + schedule, + group: QMeraLightconeGroup, + parameters, + *, + gate_registry=None, + array_backend=None, + gate_array_backend=None, + physical_dim=2, + contract=False, + product_state_factory=None, +): + """Build one shared ket for all terms in a lightcone group.""" + representative = QMeraParametricLightconeChunk( + term=group.chunks[0].term, + tags=group.chunks[0].tags, + input_sites=group.input_sites, + schedule_placement_ids=group.schedule_placement_ids, + schedule_width_by_scale=group.chunks[0].schedule_width_by_scale, + ) + return qmera_parametric_lightcone_state( + schedule, + representative, + parameters, + gate_registry=gate_registry, + array_backend=array_backend, + gate_array_backend=gate_array_backend, + physical_dim=physical_dim, + contract=contract, + product_state_factory=product_state_factory, + ) + + def qmera_parametric_lightcone_tn( schedule, chunk: QMeraParametricLightconeChunk, @@ -483,8 +630,14 @@ def contract_qmera_lightcone_tn( normalized=True, real=True, contract_opts=None, + path_cache=None, ): """Contract an explicit qMERA local-cone TN with a cotengra optimizer.""" + optimize = _resolve_contraction_opt( + optimize, + path_cache, + key=(lightcone.chunk.input_sites, lightcone.chunk.schedule_placement_ids), + ) value = _contract( lightcone.numerator, optimize=optimize, @@ -508,6 +661,40 @@ def contract_qmera_lightcone_tn( return value +def contract_qmera_lightcone_group( + group: QMeraLightconeGroup, + ket, + *, + optimize="auto-hq", + normalized=True, + real=True, + contract_opts=None, + path_cache=None, + gate_contract=True, +): + """Contract all terms in ``group`` while reusing its ket and norm.""" + optimize = _resolve_contraction_opt(optimize, path_cache, key=group.key) + denominator = None + if normalized: + denominator = _contract(ket.H & ket, optimize=optimize, contract_opts=contract_opts) + values = [] + for chunk in group.chunks: + ket_g = ket.gate_inds( + chunk.term.operator, + inds=tuple(_site_ind(site) for site in chunk.term.where), + contract=gate_contract, + inplace=False, + ) + value = _contract(ket.H & ket_g, optimize=optimize, contract_opts=contract_opts) + if normalized: + value = value / denominator + if chunk.term.weight != 1.0: + value = value * chunk.term.weight + values.append(value) + value = sum(values[1:], values[0]) if values else 0.0 + return _maybe_real(value) if real else value + + def local_qmera_parametric_lightcone_expectation( schedule, chunk: QMeraParametricLightconeChunk, @@ -524,6 +711,7 @@ def local_qmera_parametric_lightcone_expectation( gate_contract=True, contract_opts=None, product_state_factory=None, + path_cache=None, ): """Contract one qMERA local term by rebuilding only its scheduled cone.""" lightcone = qmera_parametric_lightcone_tn( @@ -544,6 +732,7 @@ def local_qmera_parametric_lightcone_expectation( normalized=normalized, real=real, contract_opts=contract_opts, + path_cache=path_cache, ) @@ -566,6 +755,8 @@ def qmera_parametric_energy( gate_contract=True, contract_opts=None, product_state_factory=None, + group_terms=True, + path_cache=None, ): """Evaluate a qMERA energy by rebuilding each scheduled lightcone only.""" array_backend = _resolve_array_backend(array_backend) @@ -577,24 +768,148 @@ def qmera_parametric_energy( terms = convert_local_terms(terms, array_backend) chunks = build_qmera_parametric_lightcone_chunks(schedule, terms) value = None - for chunk in chunks: - term_value = local_qmera_parametric_lightcone_expectation( - schedule, - chunk, - parameters, - gate_registry=gate_registry, - array_backend=array_backend, - gate_array_backend=gate_array_backend, - physical_dim=physical_dim, - optimize=optimize, - normalized=normalized, - real=False, - simplify=simplify, - gate_contract=gate_contract, - contract_opts=contract_opts, - product_state_factory=product_state_factory, + # ``simplify`` is implemented by the per-chunk TN builder. Fall back to + # that path when requested so grouping never changes this public option. + if group_terms and simplify: + group_terms = False + if group_terms: + for group in group_qmera_parametric_lightcone_chunks(chunks): + ket = qmera_parametric_lightcone_group_state( + schedule, + group, + parameters, + gate_registry=gate_registry, + array_backend=array_backend, + gate_array_backend=gate_array_backend, + physical_dim=physical_dim, + product_state_factory=product_state_factory, + ) + group_value = contract_qmera_lightcone_group( + group, + ket, + optimize=optimize, + normalized=normalized, + real=False, + contract_opts=contract_opts, + path_cache=path_cache, + gate_contract=gate_contract, + ) + value = group_value if value is None else value + group_value + else: + for chunk in chunks: + term_value = local_qmera_parametric_lightcone_expectation( + schedule, + chunk, + parameters, + gate_registry=gate_registry, + array_backend=array_backend, + gate_array_backend=gate_array_backend, + physical_dim=physical_dim, + optimize=optimize, + normalized=normalized, + real=False, + simplify=simplify, + gate_contract=gate_contract, + contract_opts=contract_opts, + product_state_factory=product_state_factory, + path_cache=path_cache, + ) + value = term_value if value is None else value + term_value + if value is None: + raise ValueError("hamiltonian contains no local terms.") + if energy_per_site: + value = value / schedule.geometry.num_sites + if real: + value = _maybe_real(value) + return value + + +def qmera_direct_parametric_energy( + schedule, + parameters, + hamiltonian=None, + *, + chunks=None, + gate_registry=None, + array_backend=None, + gate_array_backend=None, + convert_terms=True, + physical_dim=2, + optimize="auto-hq", + normalized=True, + energy_per_site=True, + real=True, + contract_opts=None, + group_terms=True, + path_cache=None, + product_state_factory=None, +): + """Evaluate qMERA energy from the complete direct-gate tensor network. + + This is intentionally a validation/debugging oracle for the + schedule-first local-cone path, not the primary optimization route. + """ + array_backend = _resolve_array_backend(array_backend) + if chunks is None: + if hamiltonian is None: + raise ValueError("qmera_direct_parametric_energy requires hamiltonian or chunks.") + terms = normalize_local_terms(hamiltonian) + if convert_terms: + terms = convert_local_terms(terms, array_backend) + chunks = build_qmera_parametric_lightcone_chunks(schedule, terms) + state = qmera_parametric_state( + schedule, + parameters, + gate_registry=gate_registry, + array_backend=array_backend, + gate_array_backend=gate_array_backend, + physical_dim=physical_dim, + product_state_factory=product_state_factory, + ) + groups = group_qmera_parametric_lightcone_chunks(chunks) if group_terms else ( + QMeraLightconeGroup( + key=(chunk.input_sites, chunk.schedule_placement_ids), + chunks=(chunk,), + input_sites=chunk.input_sites, + schedule_placement_ids=chunk.schedule_placement_ids, + ) + for chunk in chunks + ) + value = None + for group in groups: + # Direct-state validation uses the full state, while the grouping still + # shares the norm/path topology for terms with the same local support. + denominator = None + group_optimize = _resolve_contraction_opt( + optimize, + path_cache, + key=group.key, ) - value = term_value if value is None else value + term_value + if normalized: + denominator = _contract( + state.H & state, + optimize=group_optimize, + contract_opts=contract_opts, + ) + group_value = 0.0 + for chunk in group.chunks: + state_g = state.gate_inds( + chunk.term.operator, + inds=tuple(_site_ind(site) for site in chunk.term.where), + contract=True, + inplace=False, + ) + term_value = _contract( + state.H & state_g, + optimize=group_optimize, + contract_opts=contract_opts, + ) + if normalized: + term_value = term_value / denominator + if chunk.term.weight != 1.0: + term_value = term_value * chunk.term.weight + group_value = group_value + term_value + value = group_value if value is None else value + group_value if value is None: raise ValueError("hamiltonian contains no local terms.") if energy_per_site: @@ -614,11 +929,18 @@ def local_lightcone_expectation( simplify=False, gate_contract=True, contract_opts=None, + path_cache=None, ): """Contract one local expectation value over a MERA reverse lightcone.""" term = chunk.term ket = select_lightcone(state, tags=chunk.tags, validate=False) - ket_g = ket.gate( + optimize = _resolve_contraction_opt( + optimize, + path_cache, + key=(tuple(chunk.tags), tuple(chunk.physical_outer_inds)), + ) + ket_g = _apply_local_gate( + ket, term.operator, term.where, contract=gate_contract, @@ -641,3 +963,138 @@ def local_lightcone_expectation( if real: value = _maybe_real(value) return value + + +def _lightcone_state_and_schedule(state, schedule): + """Unwrap an ansatz payload while retaining an optional qMERA schedule.""" + if schedule is None: + schedule = getattr(state, "schedule", None) + candidate = getattr(state, "state", state) + if hasattr(candidate, "select") and ( + hasattr(candidate, "gate") or hasattr(candidate, "gate_inds") + ): + return candidate, schedule + raise TypeError( + "state must be a MERA-like TensorNetwork with select() and gate(), " + "or an ansatz object exposing .state." + ) + + +def _lightcone_num_sites(state, schedule=None): + """Infer the physical-site count used for an energy-per-site result.""" + if schedule is not None: + geometry = getattr(schedule, "geometry", None) + num_sites = getattr(geometry, "num_sites", None) + if num_sites is not None: + return int(num_sites() if callable(num_sites) else num_sites) + for name in ("num_sites", "sites", "L"): + value = getattr(state, name, None) + if value is None: + continue + return int(len(tuple(value)) if name == "sites" else value() if callable(value) else value) + raise ValueError("Could not infer the number of MERA physical sites.") + + +def lightcone_energy( + state, + hamiltonian=None, + *, + chunks=None, + schedule=None, + array_backend=None, + convert_terms=True, + optimize="auto-hq", + normalized=True, + energy_per_site=True, + real=True, + simplify=False, + gate_contract=True, + contract_opts=None, + group_terms=True, + path_cache=None, +): + """Evaluate a local energy by contracting only reverse-lightcone TNs. + + This is the generic, fixed-state counterpart to + :func:`qmera_parametric_energy`. Each term selects its local cone, applies + the operator with ``TensorNetwork.gate`` (or the native indexed equivalent + for graded networks), and contracts the numerator and norm with a reusable + topology-specific optimizer when ``path_cache`` is supplied. ``schedule`` + may be supplied for a qMERA state so the selector follows schedule-derived + reverse-lightcone tags. + + Native Symmray operators and tensors are passed through unchanged when + ``convert_terms=False``. Consequently the graded contraction and + fermionic signs remain owned by Symmray; this function does not introduce + Jordan--Wigner strings or dense sign corrections. + """ + state, schedule = _lightcone_state_and_schedule(state, schedule) + if chunks is None: + if hamiltonian is None: + raise ValueError("lightcone_energy requires hamiltonian or chunks.") + terms = normalize_local_terms(hamiltonian) + if convert_terms: + terms = convert_local_terms(terms, array_backend) + if schedule is None: + chunks = build_lightcone_chunks(state, terms) + else: + chunks = build_qmera_lightcone_chunks(state, schedule, terms) + else: + chunks = tuple(chunks) + if not chunks: + raise ValueError("lightcone chunks cannot be empty.") + + value = None + if group_terms: + groups = {} + for chunk in chunks: + key = (tuple(chunk.tags), tuple(chunk.physical_outer_inds)) + groups.setdefault(key, []).append(chunk) + grouped_chunks = tuple(groups.items()) + else: + grouped_chunks = tuple( + (((tuple(chunk.tags), tuple(chunk.physical_outer_inds))), [chunk]) + for chunk in chunks + ) + + for key, group in grouped_chunks: + ket = select_lightcone(state, tags=group[0].tags, validate=False) + group_optimize = _resolve_contraction_opt( + optimize, + path_cache, + key=("lightcone", key), + ) + denominator = None + if normalized: + denominator = _contract( + _maybe_simplify(ket.H & ket, simplify), + optimize=group_optimize, + contract_opts=contract_opts, + ) + for chunk in group: + term = chunk.term + ket_g = _apply_local_gate( + ket, + term.operator, + term.where, + contract=gate_contract, + inplace=False, + ) + term_value = _contract( + _maybe_simplify(ket.H & ket_g, simplify), + optimize=group_optimize, + contract_opts=contract_opts, + ) + if normalized: + term_value = term_value / denominator + if term.weight != 1.0: + term_value = term_value * term.weight + value = term_value if value is None else value + term_value + + if value is None: + raise ValueError("hamiltonian contains no local terms.") + if energy_per_site: + value = value / _lightcone_num_sites(state, schedule) + if real: + value = _maybe_real(value) + return value diff --git a/src/pepsy/optimizers/mera/schedules.py b/src/pepsy/optimizers/mera/schedules.py index 57e348a..3f6f87b 100644 --- a/src/pepsy/optimizers/mera/schedules.py +++ b/src/pepsy/optimizers/mera/schedules.py @@ -8,13 +8,19 @@ from __future__ import annotations +from collections.abc import Mapping from dataclasses import dataclass from itertools import product +from typing import Any from .geometry import QMeraGeometry __all__ = [ "QMeraBlockSpec", + "QMeraUnitarySpec", + "QMeraDisentanglerSpec", + "QMeraIsometrySpec", + "QMeraScaleSpec", "QMeraGatePlacement", "QMeraLayerSpec", "QMeraSchedule", @@ -38,6 +44,194 @@ def _normalize_structure(structure): return key +def _normalize_layer_placement(placement, kind): + key = str(placement).strip().lower().replace("_", "-") + if key in {"auto", "default"}: + return "boundary-faces" if kind == "disentangler" else "covering" + if kind == "disentangler": + if key in { + "boundary", + "boundary-face", + "boundary-faces", + "inter-block-boundary", + "inter-block-faces", + }: + return "boundary-faces" + if key in { + "boundary-square", + "inter-block-square", + "square", + }: + return "boundary-square" + if key in {"within-block", "internal", "inside-block"}: + return "within-block" + raise ValueError( + "disentangler placement must be 'boundary-faces', " + "'boundary-square', or 'within-block'." + ) + if key in {"covering", "inside-block", "within-block"}: + return "covering" + raise ValueError("isometry placement must be 'covering'.") + + +def _normalize_corner_policy(policy): + key = str(policy).strip().lower().replace("_", "-") + if key in {"include", "included", "keep"}: + return "include" + if key in {"exclude", "excluded", "remove"}: + return "exclude" + raise ValueError("corner_policy must be 'include' or 'exclude'.") + + +def _normalize_orientation(orientation): + if orientation is None: + return None + key = str(orientation).strip().lower().replace("_", "-") + if key in {"x", "horizontal", "row", "rows"}: + return "x" + if key in {"y", "vertical", "column", "columns"}: + return "y" + raise ValueError("orientation must be 'horizontal'/'x' or 'vertical'/'y'.") + + +def _normalize_isometry_implementation(implementation): + key = str(implementation).strip().lower().replace("_", "-") + if key in {"unitary-completion", "unitary", "circuit"}: + return "unitary-completion" + if key in {"true", "true-isometry", "rectangular"}: + return "true-isometry" + raise ValueError( + "isometry implementation must be 'unitary-completion' or " + "'true-isometry'." + ) + + +@dataclass(frozen=True) +class QMeraUnitarySpec: + """Metadata for the local unitary used by a qMERA layer. + + The registry still owns tensor construction. This object makes the + intended representation explicit at the layer boundary and allows the + builder to validate that a fermionic layer is not accidentally paired + with a dense spin gate family. + """ + + gate_family: str = "rxx" + family: str | None = None + arity_kind: str | None = None + symmetry: str | None = None + preserves_parity: bool | None = None + parameter_sharing: str = "per-placement" + metadata: dict[str, Any] | None = None + + def __post_init__(self): + sharing = str(self.parameter_sharing).strip().lower().replace("_", "-") + if sharing not in { + "per-placement", + "per-block", + "per-scale", + "per-axis", + "shared", + }: + raise ValueError( + "parameter_sharing must be 'per-placement', 'per-block', " + "'per-scale', 'per-axis', or 'shared'." + ) + object.__setattr__(self, "gate_family", str(self.gate_family)) + object.__setattr__(self, "family", None if self.family is None else str(self.family)) + object.__setattr__(self, "arity_kind", None if self.arity_kind is None else str(self.arity_kind)) + object.__setattr__(self, "symmetry", None if self.symmetry is None else str(self.symmetry)) + object.__setattr__(self, "parameter_sharing", sharing) + object.__setattr__( + self, + "preserves_parity", + None if self.preserves_parity is None else bool(self.preserves_parity), + ) + object.__setattr__( + self, + "metadata", + {} if self.metadata is None else dict(self.metadata), + ) + + @classmethod + def coerce(cls, value, *, default_gate_family="rxx"): + """Normalize a gate-family name or unitary metadata object.""" + if value is None: + return cls(gate_family=default_gate_family) + if isinstance(value, cls): + return value + if isinstance(value, str): + return cls(gate_family=value) + if isinstance(value, Mapping): + options = dict(value) + options.setdefault("gate_family", default_gate_family) + return cls(**options) + raise TypeError("unitary must be a gate-family name or QMeraUnitarySpec.") + + +@dataclass(frozen=True) +class QMeraDisentanglerSpec: + """Placement and local-circuit policy for boundary disentanglers.""" + + block_shape: int | tuple[int, ...] = 2 + unitary: QMeraUnitarySpec | str | None = None + circuit_depth: int = 1 + structure: str = "brickwall" + placement: str = "boundary-faces" + corner_policy: str = "include" + periodic_wrap: bool = True + orientation: str | None = None + + def to_block_spec(self, *, default_gate_family="rxx"): + """Convert to the schedule's normalized block specification.""" + unitary = QMeraUnitarySpec.coerce( + self.unitary, + default_gate_family=default_gate_family, + ) + return QMeraBlockSpec( + kind="disentangler", + block_size=self.block_shape, + circuit_depth=self.circuit_depth, + structure=self.structure, + gate_family=unitary.gate_family, + placement=self.placement, + corner_policy=self.corner_policy, + periodic_wrap=self.periodic_wrap, + orientation=self.orientation, + unitary_spec=unitary, + ) + + +@dataclass(frozen=True) +class QMeraIsometrySpec: + """Placement and implementation policy for covering isometry blocks.""" + + block_shape: int | tuple[int, ...] = 2 + unitary: QMeraUnitarySpec | str | None = None + circuit_depth: int = 1 + structure: str = "brickwall" + implementation: str = "unitary-completion" + orientation: str | None = None + + def to_block_spec(self, *, default_gate_family="rxx"): + """Convert to the schedule's normalized block specification.""" + unitary = QMeraUnitarySpec.coerce( + self.unitary, + default_gate_family=default_gate_family, + ) + return QMeraBlockSpec( + kind="isometry", + block_size=self.block_shape, + circuit_depth=self.circuit_depth, + structure=self.structure, + gate_family=unitary.gate_family, + placement="covering", + implementation=self.implementation, + orientation=self.orientation, + unitary_spec=unitary, + ) + + def _tag_token(value): chars = [] for char in str(value).upper().replace("-", "_"): @@ -54,6 +248,12 @@ class QMeraBlockSpec: circuit_depth: int = 1 structure: str = "brickwall" gate_family: str = "rxx" + placement: str = "auto" + corner_policy: str = "include" + periodic_wrap: bool = True + implementation: str = "unitary-completion" + orientation: str | None = None + unitary_spec: QMeraUnitarySpec | None = None def __post_init__(self): kind = _normalize_stage(self.kind) @@ -63,8 +263,8 @@ def __post_init__(self): block_size = int(self.block_size) circuit_depth = int(self.circuit_depth) sizes = block_size if isinstance(block_size, tuple) else (block_size,) - if any(size < 2 for size in sizes): - raise ValueError("block_size entries must be >= 2.") + if any(size < 1 for size in sizes): + raise ValueError("block_size entries must be >= 1.") if circuit_depth < 0: raise ValueError("circuit_depth must be >= 0.") object.__setattr__(self, "kind", kind) @@ -72,6 +272,95 @@ def __post_init__(self): object.__setattr__(self, "circuit_depth", circuit_depth) object.__setattr__(self, "structure", _normalize_structure(self.structure)) object.__setattr__(self, "gate_family", str(self.gate_family)) + object.__setattr__(self, "placement", _normalize_layer_placement(self.placement, kind)) + object.__setattr__(self, "corner_policy", _normalize_corner_policy(self.corner_policy)) + object.__setattr__(self, "periodic_wrap", bool(self.periodic_wrap)) + object.__setattr__(self, "orientation", _normalize_orientation(self.orientation)) + object.__setattr__( + self, + "implementation", + _normalize_isometry_implementation(self.implementation), + ) + if self.unitary_spec is not None and not isinstance(self.unitary_spec, QMeraUnitarySpec): + raise TypeError("unitary_spec must be a QMeraUnitarySpec or None.") + if kind == "isometry" and self.implementation == "true-isometry": + raise NotImplementedError( + "true-isometry tensors are not implemented yet; use " + "implementation='unitary-completion'." + ) + + +@dataclass(frozen=True) +class QMeraScaleSpec: + """User-authored configuration for one bottom-to-top RG scale. + + ``None`` uses the builder's default operation for that stage. Supplying a + scale plan is the convenient way to use different block shapes at + different scales, for example 2x2 followed by 3x3 on a 6x6 lattice. + """ + + isometry: Any | None = None + disentangler: Any | None = None + name: str | None = None + + def __post_init__(self): + if self.name is not None: + object.__setattr__(self, "name", str(self.name)) + + def to_block_specs( + self, + *, + default_disentangler, + default_isometry, + ): + """Return normalized disentangler and isometry block specs.""" + disentangler = ( + default_disentangler + if self.disentangler is None + else _coerce_schedule_block( + self.disentangler, + kind="disentangler", + default_gate_family=default_disentangler.gate_family, + ) + ) + isometry = ( + default_isometry + if self.isometry is None + else _coerce_schedule_block( + self.isometry, + kind="isometry", + default_gate_family=default_isometry.gate_family, + ) + ) + return disentangler, isometry + + +def _coerce_schedule_block(value, *, kind, default_gate_family): + """Normalize one schedule-plan block without requiring a builder.""" + if kind == "disentangler" and isinstance(value, QMeraDisentanglerSpec): + return value.to_block_spec(default_gate_family=default_gate_family) + if kind == "isometry" and isinstance(value, QMeraIsometrySpec): + return value.to_block_spec(default_gate_family=default_gate_family) + if isinstance(value, QMeraBlockSpec): + if value.kind != kind: + raise ValueError(f"{kind} block spec has kind={value.kind!r}.") + return value + options = dict(value or {}) + unitary = options.pop("unitary", None) + if unitary is not None: + unitary = QMeraUnitarySpec.coerce( + unitary, + default_gate_family=default_gate_family, + ) + options.setdefault("gate_family", unitary.gate_family) + options.setdefault("unitary_spec", unitary) + elif options.get("unitary_spec") is not None: + options["unitary_spec"] = QMeraUnitarySpec.coerce( + options["unitary_spec"], + default_gate_family=options.get("gate_family", default_gate_family), + ) + options.setdefault("gate_family", default_gate_family) + return QMeraBlockSpec(kind=kind, **options) @dataclass(frozen=True) @@ -106,6 +395,8 @@ class QMeraLayerSpec: isometry_blocks: tuple[tuple[int, ...], ...] disentanglers: tuple[QMeraGatePlacement, ...] isometries: tuple[QMeraGatePlacement, ...] + disentangler_spec: QMeraBlockSpec | None = None + isometry_spec: QMeraBlockSpec | None = None @property def placements(self): @@ -122,6 +413,7 @@ class QMeraSchedule: disentangler: QMeraBlockSpec isometry: QMeraBlockSpec top_sites: tuple[int, ...] + scale_specs: tuple[QMeraScaleSpec, ...] = () @property def placements(self): @@ -232,8 +524,13 @@ def _block_ranges(blocks): def _pairs_for_round(block, round_index, *, periodic=False): + if len(block) < 2: + return () if len(block) == 2: return ((block[0], block[1]),) + if periodic and len(block) % 2: + edge = round_index % len(block) + return ((block[edge], block[(edge + 1) % len(block)]),) start = round_index % 2 pairs = [ (block[idx], block[idx + 1]) @@ -244,6 +541,35 @@ def _pairs_for_round(block, round_index, *, periodic=False): return tuple(pairs) +def _pairs_for_1d_block( + block, + round_index, + *, + mode_by_site=None, + mode_order=None, + periodic=False, +): + """Return 1D brickwall pairs without mixing explicit fermion modes.""" + if mode_by_site is None: + return _pairs_for_round(block, round_index, periodic=periodic) + by_mode = {} + for site in block: + by_mode.setdefault(mode_by_site[site], []).append(site) + pairs = [] + for mode in sorted( + by_mode, + key=lambda value: _mode_sort_key(value, mode_order), + ): + pairs.extend( + _pairs_for_round( + tuple(by_mode[mode]), + round_index, + periodic=periodic, + ) + ) + return tuple(pairs) + + def _block_shape(block_size, ndim): if isinstance(block_size, (tuple, list)): shape = tuple(int(size) for size in block_size) @@ -254,6 +580,26 @@ def _block_shape(block_size, ndim): return shape +def _oriented_block_shape(block_size, ndim, orientation=None): + """Resolve a block shape with an optional semantic long-axis direction.""" + if ndim == 2 and orientation is not None and not isinstance( + block_size, + (tuple, list), + ): + size = int(block_size) + return (size, 1) if orientation == "x" else (1, size) + shape = _block_shape(block_size, ndim) + if ndim != 2 or orientation is None: + return shape + long_size = max(shape) + short_size = min(shape) + return ( + (long_size, short_size) + if orientation == "x" + else (short_size, long_size) + ) + + def _placement_tags(gate_id, *, scale, stage, round_index, block, gate_family, axis=None): stage_tag = "DISENTANGLER" if stage == "disentangler" else "ISOMETRY" tags = [ @@ -269,27 +615,100 @@ def _placement_tags(gate_id, *, scale, stage, round_index, block, gate_family, a return tuple(tags) -def _pairs_for_2d_block(block, round_index, coords_by_site): - axis = "x" if (round_index % 2 == 0) else "y" +def _parameter_key(gate_id, *, scale, block, stage_spec, axis): + """Return the parameter key selected by a unitary sharing policy.""" + sharing = ( + "per-placement" + if stage_spec.unitary_spec is None + else stage_spec.unitary_spec.parameter_sharing + ) + if sharing == "per-placement": + return gate_id + short = "DIS" if stage_spec.kind == "disentangler" else "ISO" + if sharing == "shared": + return f"{short}_SHARED" + if sharing == "per-scale": + return f"L{scale}_{short}_SHARED" + if sharing == "per-block": + return f"L{scale}_{short}_B{block:04d}_SHARED" + axis_token = "NONE" if axis is None else str(axis).upper() + return f"L{scale}_{short}_{axis_token}_SHARED" + + +def _mode_sort_key(mode, mode_order=None): + if mode_order is None: + return (0, repr(mode)) + try: + return (mode_order.index(mode), repr(mode)) + except ValueError: + return (len(mode_order), repr(mode)) + + +def _pairs_for_2d_block( + block, + round_index, + coords_by_site, + *, + mode_by_site=None, + mode_order=None, + periodic=False, +): + """Return spatial nearest-neighbor pairs inside a 2D block. + + With multiple modes per physical site, each spatial line is partitioned by + mode before pairing. This keeps native fermionic gates on like modes and + makes the mode-blocking convention explicit instead of silently pairing + ``up`` with ``down``. + """ + requested_axis = "x" if (round_index % 2 == 0) else "y" block_set = set(block) - pairs = [] - if axis == "x": - ys = sorted({coords_by_site[site][1] for site in block}) - for y in ys: - row = sorted( - (site for site in block_set if coords_by_site[site][1] == y), - key=lambda site: coords_by_site[site][0], - ) - pairs.extend((row[idx], row[idx + 1]) for idx in range(0, len(row) - 1, 2)) - else: - xs = sorted({coords_by_site[site][0] for site in block}) - for x in xs: - col = sorted( - (site for site in block_set if coords_by_site[site][0] == x), - key=lambda site: coords_by_site[site][1], + + def pairs_for_axis(axis): + pairs = [] + if axis == "x": + line_values = sorted({coords_by_site[site][1] for site in block}) + coordinate_axis = 0 + else: + line_values = sorted({coords_by_site[site][0] for site in block}) + coordinate_axis = 1 + for line_value in line_values: + line = ( + site + for site in block_set + if coords_by_site[site][1 - coordinate_axis] == line_value ) - pairs.extend((col[idx], col[idx + 1]) for idx in range(0, len(col) - 1, 2)) - return tuple(pairs), axis + by_mode = {} + for site in line: + mode = None if mode_by_site is None else mode_by_site[site] + by_mode.setdefault(mode, []).append(site) + for mode in sorted( + by_mode, + key=lambda value: _mode_sort_key(value, mode_order), + ): + line_sites = sorted( + by_mode[mode], + key=lambda site: coords_by_site[site][coordinate_axis], + ) + pairs.extend( + _pairs_for_round( + line_sites, + round_index, + periodic=periodic, + ) + ) + return tuple(pairs) + + pairs = pairs_for_axis(requested_axis) + if not pairs: + # Coarse grids can be sparse along the requested brickwall axis after + # one RG step (e.g. a 2x4 grid reduced to two sites separated in y). + # Use the populated axis so every nontrivial block still receives its + # intended isometry/disentangler gate. + alternate_axis = "y" if requested_axis == "x" else "x" + alternate_pairs = pairs_for_axis(alternate_axis) + if alternate_pairs: + return alternate_pairs, alternate_axis + return pairs, requested_axis def _stage_placements( @@ -299,8 +718,11 @@ def _stage_placements( stage_spec, counter_start, coords_by_site=None, + mode_by_site=None, + mode_order=None, boundary_pairs_by_block=None, block_axes=None, + periodic=False, ): placements = [] counter = counter_start @@ -312,16 +734,35 @@ def _stage_placements( pairs = boundary_pairs_by_block[block_index] axis = None if block_axes is None else block_axes[block_index] elif coords_by_site is not None: - pairs, axis = _pairs_for_2d_block(block, round_index, coords_by_site) + pairs, axis = _pairs_for_2d_block( + block, + round_index, + coords_by_site, + mode_by_site=mode_by_site, + mode_order=mode_order, + periodic=periodic, + ) else: - pairs = _pairs_for_round(block, round_index, periodic=False) + pairs = _pairs_for_1d_block( + block, + round_index, + mode_by_site=mode_by_site, + mode_order=mode_order, + periodic=periodic, + ) axis = None for pair in pairs: gate_id = f"L{scale}_{short}_{counter:04d}" placements.append( QMeraGatePlacement( gate_id=gate_id, - param_key=gate_id, + param_key=_parameter_key( + gate_id, + scale=scale, + block=block_index, + stage_spec=stage_spec, + axis=axis, + ), where=tuple(pair), scale=scale, stage=stage, @@ -360,7 +801,22 @@ def _active_coords_by_site(geometry, active): def _nonoverlapping_blocks_2d(active, geometry, block_shape): coords_by_site = _active_coords_by_site(geometry, active) - site_by_coord = {coo: site for site, coo in coords_by_site.items()} + mode_by_site = { + site: ( + geometry.to_mode(site)[-1] + if geometry.has_explicit_modes + and isinstance(geometry.to_mode(site), tuple) + else None + ) + for site in active + } + # Map coordinates to physical site labels, then expand each block back to + # every active mode on those physical sites. This also handles ordinary + # one-mode geometries where the register label is an integer. + site_by_coord = { + coo: geometry.to_site(register_site) + for register_site, coo in coords_by_site.items() + } xs = sorted({coo[0] for coo in site_by_coord}) ys = sorted({coo[1] for coo in site_by_coord}) bx_size, by_size = block_shape @@ -370,15 +826,50 @@ def _nonoverlapping_blocks_2d(active, geometry, block_shape): x_vals = xs[x_start : x_start + bx_size] for biy, y_start in enumerate(range(0, len(ys), by_size)): y_vals = ys[y_start : y_start + by_size] - block = tuple( + physical_block = tuple( site_by_coord[(x, y)] for x, y in product(x_vals, y_vals) if (x, y) in site_by_coord ) + block = tuple( + register_site + for physical_site in physical_block + for register_site in geometry.site_to_registers[physical_site] + if register_site in active + ) if block: block_grid[(bix, biy)] = block blocks.append(block) - return tuple(blocks), block_grid, coords_by_site + return tuple(blocks), block_grid, coords_by_site, mode_by_site + + +def _within_blocks_2d(isometry_blocks, geometry, block_shape): + """Tile each covering isometry block with internal dis-entangler blocks.""" + blocks = [] + for isometry_block in isometry_blocks: + internal, _, _, _ = _nonoverlapping_blocks_2d( + isometry_block, + geometry, + block_shape, + ) + blocks.extend(internal) + return tuple(blocks) + + +def _coarse_grain_2d(isometry_blocks, active, geometry): + """Keep every mode on one representative site per 2D RG block.""" + active_set = set(active) + output = [] + for block in isometry_blocks: + if not block: + continue + representative_site = geometry.to_site(block[0]) + output.extend( + register_site + for register_site in geometry.site_to_registers[representative_site] + if register_site in active_set + ) + return tuple(output) def _slab_sites(block, coords_by_site, *, axis, side, depth=1): @@ -402,17 +893,87 @@ def _slab_sites(block, coords_by_site, *, axis, side, depth=1): ) -def _face_pairs(left_face, right_face, coords_by_site, *, axis): +def _face_pairs( + left_face, + right_face, + coords_by_site, + *, + axis, + mode_by_site=None, + mode_order=None, +): match_axis = 1 if axis == "x" else 0 - left_by_coord = {coords_by_site[site][match_axis]: site for site in left_face} - right_by_coord = {coords_by_site[site][match_axis]: site for site in right_face} + left_by_coord = { + ( + None if mode_by_site is None else mode_by_site[site], + coords_by_site[site][match_axis], + ): site + for site in left_face + } + right_by_coord = { + ( + None if mode_by_site is None else mode_by_site[site], + coords_by_site[site][match_axis], + ): site + for site in right_face + } pairs = [] - for value in sorted(set(left_by_coord).intersection(right_by_coord)): + keys = sorted( + set(left_by_coord).intersection(right_by_coord), + key=lambda value: (_mode_sort_key(value[0], mode_order), value[1]), + ) + for value in keys: pairs.append((left_by_coord[value], right_by_coord[value])) return tuple(pairs) -def _boundary_blocks_2d(block_grid, coords_by_site, *, boundary, width=2): +def _trim_boundary_corners(sites, coords_by_site, *, policy): + """Optionally remove the corner sites from a boundary support.""" + if policy == "include" or not sites: + return tuple(sites) + xs = {coords_by_site[site][0] for site in sites} + ys = {coords_by_site[site][1] for site in sites} + if len(xs) < 2 or len(ys) < 2: + return tuple(sites) + x_edges = {min(xs), max(xs)} + y_edges = {min(ys), max(ys)} + return tuple( + site + for site in sites + if not ( + coords_by_site[site][0] in x_edges + and coords_by_site[site][1] in y_edges + ) + ) + + +def _prepare_boundary_faces(left_face, right_face, coords_by_site, *, policy): + """Apply the corner policy to the complete two-face square support.""" + left_face = tuple(left_face) + right_face = tuple(right_face) + allowed = set( + _trim_boundary_corners( + (*left_face, *right_face), + coords_by_site, + policy=policy, + ) + ) + return ( + tuple(site for site in left_face if site in allowed), + tuple(site for site in right_face if site in allowed), + ) + + +def _boundary_blocks_2d( + block_grid, + coords_by_site, + *, + boundary, + width=2, + mode_by_site=None, + mode_order=None, + corner_policy="include", +): depth = max(1, int(width) // 2) blocks = [] pairs_by_block = [] @@ -426,9 +987,20 @@ def _boundary_blocks_2d(block_grid, coords_by_site, *, boundary, width=2): right = block_grid.get((bx + 1, by)) if left is None or right is None: continue - left_face = _slab_sites(left, coords_by_site, axis="x", side="right", depth=depth) - right_face = _slab_sites(right, coords_by_site, axis="x", side="left", depth=depth) - pairs = _face_pairs(left_face, right_face, coords_by_site, axis="x") + left_face, right_face = _prepare_boundary_faces( + _slab_sites(left, coords_by_site, axis="x", side="right", depth=depth), + _slab_sites(right, coords_by_site, axis="x", side="left", depth=depth), + coords_by_site, + policy=corner_policy, + ) + pairs = _face_pairs( + left_face, + right_face, + coords_by_site, + axis="x", + mode_by_site=mode_by_site, + mode_order=mode_order, + ) if pairs: blocks.append(tuple(left_face + right_face)) pairs_by_block.append(pairs) @@ -440,16 +1012,27 @@ def _boundary_blocks_2d(block_grid, coords_by_site, *, boundary, width=2): top = block_grid.get((bx, by + 1)) if bottom is None or top is None: continue - bottom_face = _slab_sites(bottom, coords_by_site, axis="y", side="top", depth=depth) - top_face = _slab_sites(top, coords_by_site, axis="y", side="bottom", depth=depth) - pairs = _face_pairs(bottom_face, top_face, coords_by_site, axis="y") + bottom_face, top_face = _prepare_boundary_faces( + _slab_sites(bottom, coords_by_site, axis="y", side="top", depth=depth), + _slab_sites(top, coords_by_site, axis="y", side="bottom", depth=depth), + coords_by_site, + policy=corner_policy, + ) + pairs = _face_pairs( + bottom_face, + top_face, + coords_by_site, + axis="y", + mode_by_site=mode_by_site, + mode_order=mode_order, + ) if pairs: blocks.append(tuple(bottom_face + top_face)) pairs_by_block.append(pairs) axes.append("y") if boundary == "periodic": - if len(bx_values) > 2: + if len(bx_values) >= 2: bx_left = bx_values[-1] bx_right = bx_values[0] for by in by_values: @@ -457,14 +1040,25 @@ def _boundary_blocks_2d(block_grid, coords_by_site, *, boundary, width=2): right = block_grid.get((bx_right, by)) if left is None or right is None: continue - left_face = _slab_sites(left, coords_by_site, axis="x", side="right", depth=depth) - right_face = _slab_sites(right, coords_by_site, axis="x", side="left", depth=depth) - pairs = _face_pairs(left_face, right_face, coords_by_site, axis="x") + left_face, right_face = _prepare_boundary_faces( + _slab_sites(left, coords_by_site, axis="x", side="right", depth=depth), + _slab_sites(right, coords_by_site, axis="x", side="left", depth=depth), + coords_by_site, + policy=corner_policy, + ) + pairs = _face_pairs( + left_face, + right_face, + coords_by_site, + axis="x", + mode_by_site=mode_by_site, + mode_order=mode_order, + ) if pairs: blocks.append(tuple(left_face + right_face)) pairs_by_block.append(pairs) axes.append("x") - if len(by_values) > 2: + if len(by_values) >= 2: by_bottom = by_values[-1] by_top = by_values[0] for bx in bx_values: @@ -472,9 +1066,20 @@ def _boundary_blocks_2d(block_grid, coords_by_site, *, boundary, width=2): top = block_grid.get((bx, by_top)) if bottom is None or top is None: continue - bottom_face = _slab_sites(bottom, coords_by_site, axis="y", side="top", depth=depth) - top_face = _slab_sites(top, coords_by_site, axis="y", side="bottom", depth=depth) - pairs = _face_pairs(bottom_face, top_face, coords_by_site, axis="y") + bottom_face, top_face = _prepare_boundary_faces( + _slab_sites(bottom, coords_by_site, axis="y", side="top", depth=depth), + _slab_sites(top, coords_by_site, axis="y", side="bottom", depth=depth), + coords_by_site, + policy=corner_policy, + ) + pairs = _face_pairs( + bottom_face, + top_face, + coords_by_site, + axis="y", + mode_by_site=mode_by_site, + mode_order=mode_order, + ) if pairs: blocks.append(tuple(bottom_face + top_face)) pairs_by_block.append(pairs) @@ -488,34 +1093,80 @@ def _build_qmera_schedule_1d( *, disentangler, isometry, + scale_specs=None, max_layers, top_size, ): - isometry_size = _block_shape(isometry.block_size, 1)[0] - disentangler_size = _block_shape(disentangler.block_size, 1)[0] - active = geometry.register_sites + mode_by_site = None + mode_order = None + if geometry.has_explicit_modes: + mode_by_site = { + site: geometry.to_mode(site)[-1] + for site in active + } + mode_order = geometry.site_modes layers = [] gate_counter = 0 scale = 0 - while len(active) > top_size and (max_layers is None or scale < max_layers): - isometry_blocks = _nonoverlapping_blocks(active, isometry_size) - disentangler_blocks = _boundary_blocks( - isometry_blocks, - disentangler_size, - periodic=geometry.boundary == "periodic", - ) - dis, gate_counter = _stage_placements( - disentangler_blocks, - scale=scale, - stage_spec=disentangler, - counter_start=gate_counter, + while ( + len( + {geometry.to_site(register_site) for register_site in active} + if geometry.has_explicit_modes + else active ) + > top_size + and (max_layers is None or scale < max_layers) + ): + if scale_specs is not None: + if scale >= len(scale_specs): + raise ValueError( + "qMERA scale plan ended before the geometry reached top_size." + ) + disentangler, isometry = scale_specs[scale] + isometry_size = _block_shape(isometry.block_size, 1)[0] + disentangler_size = _block_shape(disentangler.block_size, 1)[0] + isometry_blocks = _nonoverlapping_blocks(active, isometry_size) + if disentangler.placement == "within-block": + disentangler_blocks = tuple( + internal + for block in isometry_blocks + for internal in _nonoverlapping_blocks(block, disentangler_size) + ) + dis, gate_counter = _stage_placements( + disentangler_blocks, + scale=scale, + stage_spec=disentangler, + counter_start=gate_counter, + mode_by_site=mode_by_site, + mode_order=mode_order, + periodic=( + geometry.boundary == "periodic" and disentangler.periodic_wrap + ), + ) + else: + disentangler_blocks = _boundary_blocks( + isometry_blocks, + disentangler_size, + periodic=( + geometry.boundary == "periodic" and disentangler.periodic_wrap + ), + ) + dis, gate_counter = _stage_placements( + disentangler_blocks, + scale=scale, + stage_spec=disentangler, + counter_start=gate_counter, + mode_by_site=mode_by_site, + mode_order=mode_order, + ) iso, gate_counter = _stage_placements( _placement_blocks(isometry_blocks), scale=scale, stage_spec=isometry, counter_start=gate_counter, + mode_by_site=mode_by_site, + mode_order=mode_order, ) output_sites = _coarse_grain(isometry_blocks) if output_sites == active: @@ -529,6 +1180,8 @@ def _build_qmera_schedule_1d( isometry_blocks=isometry_blocks, disentanglers=dis, isometries=iso, + disentangler_spec=disentangler, + isometry_spec=isometry, ) ) active = output_sites @@ -542,53 +1195,118 @@ def _build_qmera_schedule_2d( *, disentangler, isometry, + scale_specs=None, max_layers, top_size, ): if top_size != 1: raise NotImplementedError("2D qMERA schedules currently support top_size=1.") - isometry_shape = _block_shape(isometry.block_size, 2) - if isinstance(disentangler.block_size, tuple): - disentangler_width = max(disentangler.block_size) - else: - disentangler_width = _block_shape(disentangler.block_size, 1)[0] active = geometry.register_sites layers = [] gate_counter = 0 scale = 0 - while len(active) > top_size and (max_layers is None or scale < max_layers): - isometry_blocks, block_grid, coords_by_site = _nonoverlapping_blocks_2d( - active, - geometry, - isometry_shape, + while ( + len({geometry.to_site(site) for site in active}) > top_size + and (max_layers is None or scale < max_layers) + ): + if scale_specs is not None: + if scale >= len(scale_specs): + raise ValueError( + "qMERA scale plan ended before the geometry reached top_size." + ) + disentangler, isometry = scale_specs[scale] + isometry_shape = _oriented_block_shape( + isometry.block_size, + 2, + isometry.orientation, + ) + disentangler_shape = _oriented_block_shape( + disentangler.block_size, + 2, + disentangler.orientation, ) + disentangler_width = max(disentangler_shape) ( - disentangler_blocks, - dis_pairs_by_block, - dis_axes, - ) = _boundary_blocks_2d( + isometry_blocks, block_grid, coords_by_site, - boundary=geometry.boundary, - width=disentangler_width, - ) - dis, gate_counter = _stage_placements( - disentangler_blocks, - scale=scale, - stage_spec=disentangler, - counter_start=gate_counter, - boundary_pairs_by_block=dis_pairs_by_block, - block_axes=dis_axes, + mode_by_site, + ) = _nonoverlapping_blocks_2d( + active, + geometry, + isometry_shape, ) + dis_pairs_by_block = None + dis_axes = None + if disentangler.placement == "within-block": + disentangler_blocks = _within_blocks_2d( + isometry_blocks, + geometry, + disentangler_shape, + ) + dis, gate_counter = _stage_placements( + disentangler_blocks, + scale=scale, + stage_spec=disentangler, + counter_start=gate_counter, + coords_by_site=coords_by_site, + mode_by_site=mode_by_site, + mode_order=geometry.site_modes, + periodic=( + geometry.boundary == "periodic" and disentangler.periodic_wrap + ), + ) + else: + ( + disentangler_blocks, + dis_pairs_by_block, + dis_axes, + ) = _boundary_blocks_2d( + block_grid, + coords_by_site, + boundary=( + geometry.boundary + if disentangler.periodic_wrap + else "open" + ), + width=disentangler_width, + mode_by_site=mode_by_site, + mode_order=geometry.site_modes, + corner_policy=disentangler.corner_policy, + ) + if disentangler.placement == "boundary-square": + dis, gate_counter = _stage_placements( + disentangler_blocks, + scale=scale, + stage_spec=disentangler, + counter_start=gate_counter, + coords_by_site=coords_by_site, + mode_by_site=mode_by_site, + mode_order=geometry.site_modes, + ) + elif disentangler.placement != "within-block": + dis, gate_counter = _stage_placements( + disentangler_blocks, + scale=scale, + stage_spec=disentangler, + counter_start=gate_counter, + boundary_pairs_by_block=dis_pairs_by_block, + block_axes=dis_axes, + mode_order=geometry.site_modes, + ) iso, gate_counter = _stage_placements( _placement_blocks(isometry_blocks), scale=scale, stage_spec=isometry, counter_start=gate_counter, coords_by_site=coords_by_site, + mode_by_site=mode_by_site, + mode_order=geometry.site_modes, ) - output_sites = _coarse_grain(isometry_blocks) + output_sites = _coarse_grain_2d(isometry_blocks, active, geometry) + if output_sites == active: + break layers.append( QMeraLayerSpec( scale=scale, @@ -598,6 +1316,8 @@ def _build_qmera_schedule_2d( isometry_blocks=isometry_blocks, disentanglers=dis, isometries=iso, + disentangler_spec=disentangler, + isometry_spec=isometry, ) ) active = output_sites @@ -611,26 +1331,58 @@ def build_qmera_schedule( *, disentangler=None, isometry=None, + scales=None, max_layers=None, top_size=1, ): """Build a deterministic brickwall qMERA schedule.""" geometry = geometry if isinstance(geometry, QMeraGeometry) else QMeraGeometry(geometry) - disentangler = ( - disentangler - if isinstance(disentangler, QMeraBlockSpec) - else QMeraBlockSpec(kind="disentangler", **dict(disentangler or {})) + disentangler = _coerce_schedule_block( + disentangler, + kind="disentangler", + default_gate_family="rxx", ) - isometry = ( - isometry - if isinstance(isometry, QMeraBlockSpec) - else QMeraBlockSpec(kind="isometry", **dict(isometry or {})) + isometry = _coerce_schedule_block( + isometry, + kind="isometry", + default_gate_family="rxx", ) if disentangler.kind != "disentangler": raise ValueError("disentangler spec must have kind='disentangler'.") if isometry.kind != "isometry": raise ValueError("isometry spec must have kind='isometry'.") + normalized_scales = None + if scales is not None: + normalized_scales = [] + for scale in tuple(scales): + if not isinstance(scale, QMeraScaleSpec): + if not isinstance(scale, Mapping): + raise TypeError( + "scales must contain QMeraScaleSpec or mapping objects." + ) + scale = QMeraScaleSpec(**dict(scale)) + scale_disentangler, scale_isometry = scale.to_block_specs( + default_disentangler=disentangler, + default_isometry=isometry, + ) + normalized_scales.append( + QMeraScaleSpec( + disentangler=scale_disentangler, + isometry=scale_isometry, + name=scale.name, + ) + ) + normalized_scales = tuple(normalized_scales) + if not normalized_scales: + raise ValueError("scales must contain at least one scale specification.") + scale_pairs = tuple( + (scale.disentangler, scale.isometry) + for scale in normalized_scales + ) + else: + scale_pairs = None + top_size = int(top_size) if top_size < 1: raise ValueError("top_size must be >= 1.") @@ -643,19 +1395,16 @@ def build_qmera_schedule( geometry, disentangler=disentangler, isometry=isometry, + scale_specs=scale_pairs, max_layers=max_layers, top_size=top_size, ) elif geometry.ndim == 2: - if geometry.num_modes != geometry.num_sites: - raise NotImplementedError( - "2D qMERA schedules with multiple modes per lattice site need an " - "explicit mode-blocking convention; use 1D mode schedules for now." - ) layers, top_sites, _ = _build_qmera_schedule_2d( geometry, disentangler=disentangler, isometry=isometry, + scale_specs=scale_pairs, max_layers=max_layers, top_size=top_size, ) @@ -665,7 +1414,18 @@ def build_qmera_schedule( return QMeraSchedule( geometry=geometry, layers=tuple(layers), - disentangler=disentangler, - isometry=isometry, + disentangler=( + normalized_scales[0].disentangler + if normalized_scales is not None + else disentangler + ), + isometry=( + normalized_scales[0].isometry + if normalized_scales is not None + else isometry + ), top_sites=top_sites, + scale_specs=() + if normalized_scales is None + else normalized_scales, ) diff --git a/src/pepsy/optimizers/mera/schematics.py b/src/pepsy/optimizers/mera/schematics.py index 4ae61e1..789b3d4 100644 --- a/src/pepsy/optimizers/mera/schematics.py +++ b/src/pepsy/optimizers/mera/schematics.py @@ -184,6 +184,173 @@ def _patch_block(drawing, coos, *, block, label_blocks): drawing.text((x, y), block.stage_label, preset="block_label") +def _unique_physical_sites(geometry, register_sites): + """Return physical sites represented by possibly repeated mode registers.""" + sites = [] + seen = set() + for register_site in register_sites: + site = geometry.to_site(register_site) + if site not in seen: + seen.add(site) + sites.append(site) + return tuple(sites) + + +def _clean_stage_positions(geometry, sites, *, x0, y0): + """Place 2D physical sites on a stable schematic grid.""" + return { + site: (x0 + float(site[1]), y0 - float(site[0])) + for site in sites + } + + +def _draw_clean_stage( + drawing, + schedule, + layer_spec, + blocks, + stage, + *, + x0, + y0, + label_sites, + label_blocks, +): + """Draw one clean input, gate, or coarse-output stage.""" + geometry = schedule.geometry + register_sites = ( + layer_spec.output_sites + if stage == "output" + else layer_spec.input_sites + ) + sites = _unique_physical_sites(geometry, register_sites) + positions = _clean_stage_positions(geometry, sites, x0=x0, y0=y0) + + # Keep the physical graph visible in every stage, like the simple wires in + # quimb's manual schematic examples. + site_set = set(sites) + for left, right in geometry.nearest_neighbor_edges(): + if left not in site_set or right not in site_set: + continue + a = positions[left] + b = positions[right] + drawing.line(a, b, preset="wire") + + for site, coo in positions.items(): + drawing.circle(coo, radius=0.12, preset="site") + if label_sites and stage in {"input", "output"}: + drawing.text( + (coo[0], coo[1] + 0.23), + str(site), + preset="site_label", + ) + + if stage not in {"disentangler", "isometry"}: + return + + for block in blocks: + block_sites = tuple(dict.fromkeys(block.sites)) + coos = [positions[site] for site in block_sites if site in positions] + if not coos: + continue + _patch_block(drawing, coos, block=block, label_blocks=False) + if label_blocks: + x = sum(coo[0] for coo in coos) / len(coos) + y = sum(coo[1] for coo in coos) / len(coos) + drawing.text( + (x, y), + f"{block.stage_label}{block.block}", + preset="block_label", + ) + + +def _draw_clean_2d( + drawing, + schedule, + layer_indices, + blocks_by_layer, + *, + label_sites, + label_blocks, +): + """Draw 2D layers as separated quimb-style input/D/W/output panels.""" + height, width = schedule.geometry.shape + panel_width = float(max(width, height)) + 2.0 + panel_gap = 1.5 + row_gap = float(max(height, 1)) + 3.2 + cursor_x = 0.0 + + for row, layer_index in enumerate(layer_indices): + layer = schedule.layers[layer_index] + y0 = -row_gap * row + stages = [("input", "input")] + disentanglers = tuple( + block + for block in blocks_by_layer.get(layer_index, ()) + if block.stage == "disentangler" + ) + isometries = tuple( + block + for block in blocks_by_layer.get(layer_index, ()) + if block.stage == "isometry" + ) + if disentanglers: + stages.append(("disentangler", "D")) + if isometries: + stages.append(("isometry", "W")) + stages.append(("output", "coarse")) + + layer_start = cursor_x + previous_right = None + for stage, label in stages: + x0 = cursor_x + blocks = ( + disentanglers + if stage == "disentangler" + else isometries + if stage == "isometry" + else () + ) + _draw_clean_stage( + drawing, + schedule, + layer, + blocks, + stage, + x0=x0, + y0=y0, + label_sites=label_sites, + label_blocks=label_blocks, + ) + drawing.text( + (x0 + 0.5 * (width - 1), y0 + 0.85), + label, + preset="stage_label", + ) + + if previous_right is not None: + arrow_y = y0 - float(height) - 0.7 + start = (previous_right, arrow_y) + end = (x0 - 0.35, arrow_y) + drawing.line(start, end, preset="flow") + drawing.arrowhead(start, end, preset="flow") + previous_right = x0 + panel_width - 0.35 + cursor_x += panel_width + panel_gap + + drawing.text( + (layer_start - 0.7, y0), + f"L{layer_index}", + preset="layer_label", + ) + cursor_x += 0.8 + + drawing.text( + (0.0, -row_gap * len(layer_indices) + 1.0), + "D = boundary disentangler W = isometry arrows = coarse-graining", + preset="legend", + ) + + def _draw_2d_layers( drawing, schedule, @@ -244,6 +411,7 @@ def draw_qmera_schedule( schedule, *, layer=None, + style="clean", figsize=None, label_sites=True, label_blocks=True, @@ -260,6 +428,10 @@ def draw_qmera_schedule( layer Optional layer index or iterable of layer indices. ``None`` draws all layers. + style : {"clean", "register"}, default="clean" + ``"clean"`` separates 2D input, disentangler, isometry, and coarse + output panels. ``"register"`` keeps the lower-level register wiring + view. figsize Optional matplotlib figure size forwarded to ``schematic.Drawing``. label_sites, label_blocks @@ -271,6 +443,8 @@ def draw_qmera_schedule( ax Optional matplotlib axes. """ + if style not in {"clean", "register"}: + raise ValueError("style must be 'clean' or 'register'.") try: from quimb import schematic except ImportError as exc: # pragma: no cover - optional plotting dep @@ -306,6 +480,13 @@ def draw_qmera_schedule( "verticalalignment": "center", }, "layer_label": {"fontsize": 10, "fontweight": "bold"}, + "stage_label": { + "fontsize": 10, + "fontweight": "bold", + "horizontalalignment": "center", + }, + "flow": {"linewidth": 1.5, "color": neutral}, + "legend": {"fontsize": 9, "color": neutral_dark}, } if presets is not None: merged = dict(default_presets) @@ -326,14 +507,24 @@ def draw_qmera_schedule( blocks_by_layer.setdefault(block.scale, []).append(block) if schedule.geometry.ndim == 2: - _draw_2d_layers( - drawing, - schedule, - layer_indices, - blocks_by_layer, - label_sites=label_sites, - label_blocks=label_blocks, - ) + if style == "clean": + _draw_clean_2d( + drawing, + schedule, + layer_indices, + blocks_by_layer, + label_sites=label_sites, + label_blocks=label_blocks, + ) + else: + _draw_2d_layers( + drawing, + schedule, + layer_indices, + blocks_by_layer, + label_sites=label_sites, + label_blocks=label_blocks, + ) if scale_figsize: drawing.scale_figsize(1.0) return drawing diff --git a/src/pepsy/tensors/symmetric.py b/src/pepsy/tensors/symmetric.py index ecdce1b..5744614 100644 --- a/src/pepsy/tensors/symmetric.py +++ b/src/pepsy/tensors/symmetric.py @@ -3719,6 +3719,7 @@ def symm_operator_from_dense( fermionic=False, sites=None, index_maps=None, + label=None, ): """Convert a dense local operator to a Symmray block-sparse array. @@ -3793,6 +3794,8 @@ def symm_operator_from_dense( kwargs = {} if array_cls.__name__ in {"AbelianArray", "FermionicArray"}: kwargs["symmetry"] = symmetry + if label is not None: + kwargs["label"] = label return array_cls.from_dense( arr, index_maps=index_maps, @@ -7311,11 +7314,13 @@ def fermion_hopping_param_gen( class Fermion: """Native spinless or spinful fermion observables, gates, and streams. - The helper keeps the symmetry convention, local fermionic operators, and - optional hopping/interaction parameters together. It is intended for direct - Symmray-backed fermionic MPS or PEPS workflows; it does not introduce a - qubit or Jordan-Wigner circuit representation. ``hamiltonian(...)`` is a - convenience only; the object is equally useful for measurements and gates. + The helper owns only the local fermionic space, symmetry convention, and + optional backend conversion. Hamiltonian couplings are deliberately not + stored here: construct them as explicit native terms, then validate and + bundle them with :meth:`hamiltonian`. This prevents a native Hamiltonian, + a gate stream, and a VMC adapter from silently using different couplings. + It is intended for direct Symmray-backed fermionic MPS or PEPS workflows; + it does not introduce a qubit or Jordan-Wigner circuit representation. ``strang_gate_stream`` uses a deterministic edge colouring and a forward/reverse half-step sequence. Consequently its hopping layers are @@ -7325,13 +7330,9 @@ class Fermion: """ symmetry: str | None = None - t: object = 1.0 - U: object = 8.0 dtype: object = "complex128" to_backend: object = None spinful: bool = True - V: object = 0.0 - mu: object = 0.0 _dense_ops: dict = field(default_factory=dict, init=False, repr=False) _observable_cache: dict = field(default_factory=dict, init=False, repr=False) _gate_cache: dict = field(default_factory=dict, init=False, repr=False) @@ -7848,6 +7849,7 @@ def operator_term( charge=None, like=None, add_hc=False, + label=None, ): """Return a native operator made from explicit fermion monomials. @@ -7990,9 +7992,185 @@ def operator_term( } for _ in range(2 * len(sites)) ), + label=label, ) return _apply_to_array_blocks(operator, self.to_backend) + @staticmethod + def _normalize_majorana_component(component): + key = str(component).strip().lower() + if component in {0, "0"} or key in {"x", "real", "gamma_x", "gamma0"}: + return 0 + if component in {1, "1"} or key in {"y", "imag", "gamma_y", "gamma1"}: + return 1 + raise ValueError("Majorana component must be 0/'x' or 1/'y'.") + + def _require_majorana(self, feature): + if self.spinful: + raise NotImplementedError( + f"{feature} currently targets one complex mode per site; " + "use Fermion(spinful=False) or provide an explicit flavor map." + ) + if self.symmetry != "Z2": + raise ValueError( + f"{feature} uses the native parity convention and requires " + "symmetry='Z2'; U1/U1U1 does not make a single Majorana " + "operator homogeneous." + ) + + def _majorana_charge(self): + self._require_majorana("Majorana operators") + return _normalize_group_charge(1, self.symmetry) + + def _majorana_mode_terms(self, site, component): + component = self._normalize_majorana_component(component) + if component == 0: + return ((1.0, ((site, "create"),)), (1.0, ((site, "annihilate"),))) + return ( + (1.0j, ((site, "create"),)), + (-1.0j, ((site, "annihilate"),)), + ) + + def majorana_operator(self, component=0, *, site=0): + """Return a native parity-odd Majorana operator. + + The convention is ``gamma_x = c + c^†`` and + ``gamma_y = -i (c - c^†)``. It is intentionally a ``Z2`` path: + individual Majoranas are not homogeneous under particle-number ``U1``. + """ + charge = self._majorana_charge() + return self.operator_term( + self._majorana_mode_terms(site, component), + sites=(site,), + charge=charge, + label=f"majorana_{site!r}", + ) + + def _majorana_bilinear_terms( + self, + left, + right, + *, + left_component=0, + right_component=0, + coefficient=1.0, + canonical=True, + ): + if left == right: + raise ValueError("Majorana bilinears require distinct mode sites.") + terms = [] + for left_coeff, left_ops in self._majorana_mode_terms(left, left_component): + for right_coeff, right_ops in self._majorana_mode_terms(right, right_component): + coefficient_term = 1.0j * coefficient * left_coeff * right_coeff + # Symmray's graded local-element builder canonicalizes the + # all-annihilator monomial in site order. Compensate that + # reversal sign so ``i * gamma_left * gamma_right`` is + # Hermitian in the native fermionic representation. + if canonical and ( + left_ops[0][1] == "annihilate" + and right_ops[0][1] == "annihilate" + ): + coefficient_term = -coefficient_term + terms.append( + ( + coefficient_term, + (*left_ops, *right_ops), + ) + ) + return tuple(terms) + + def majorana_bilinear_operator( + self, + edge, + *, + left_component=0, + right_component=0, + coefficient=1.0, + ): + """Return ``coefficient * i gamma_left gamma_right``.""" + self._require_majorana("Majorana bilinears") + try: + left, right = tuple(edge) + except (TypeError, ValueError) as exc: + raise ValueError("edge must contain exactly two mode sites.") from exc + return self.operator_term( + self._majorana_bilinear_terms( + left, + right, + left_component=left_component, + right_component=right_component, + coefficient=coefficient, + canonical=True, + ), + sites=(left, right), + charge=self.zero_charge, + ) + + def pairing_operator(self, edge, *, coefficient=1.0, phase=0.0): + """Return a Hermitian spinless pairing operator on ``edge``.""" + self._require_majorana("Pairing operators") + try: + left, right = tuple(edge) + except (TypeError, ValueError) as exc: + raise ValueError("edge must contain exactly two mode sites.") from exc + amplitude = coefficient * _fermion_complex_phase(phase, like=coefficient) + return self.operator_term( + ( + (amplitude, ((left, "create"), (right, "create"))), + ( + ar.do("conj", amplitude), + ((left, "annihilate"), (right, "annihilate")), + ), + ), + sites=(left, right), + charge=self.zero_charge, + ) + + def majorana_gate( + self, + dt, + *, + edge, + left_component=0, + right_component=0, + coefficient=1.0, + imaginary=False, + ): + """Return ``exp(-i dt * i gamma_left gamma_right)``.""" + self._require_majorana("Majorana gates") + left, right = tuple(edge) + return self.exponential( + self._majorana_bilinear_terms( + left, + right, + left_component=left_component, + right_component=right_component, + coefficient=coefficient, + canonical=False, + ), + dt, + sites=(left, right), + imaginary=imaginary, + ) + + def pairing_gate(self, dt, *, edge, coefficient=1.0, phase=0.0, imaginary=False): + """Return ``exp(-i dt H_pair)`` for a parity-preserving pairing term.""" + self._require_majorana("Pairing gates") + left, right = tuple(edge) + amplitude = coefficient * _fermion_complex_phase(phase, like=coefficient) + return self.exponential( + ( + (amplitude, ((left, "create"), (right, "create"))), + ( + -ar.do("conj", amplitude), + ((left, "annihilate"), (right, "annihilate")), + ), + ), + dt, + sites=(left, right), + imaginary=imaginary, + ) + def eta_pair_operator(self, *, coefficient=1.0): """Return ``coefficient * Delta_0^dag Delta_1 + h.c.``. @@ -8085,13 +8263,12 @@ def hopping_operator(self, *, spin=None, peierls_angle=0.0): peierls_angle=peierls_angle, ) - def hopping_term(self, edge, *, spin=None, t=None, peierls_angle=0.0): + def hopping_term(self, edge, *, spin=None, t, peierls_angle=0.0): """Return ``-t`` times the hopping operator on ``edge``.""" try: left, right = tuple(edge) except (TypeError, ValueError) as exc: raise ValueError("edge must contain exactly two site labels.") from exc - t = self.t if t is None else t t = _edge_parameter(t, left, right) return self._hopping_operator_on_sites( left, @@ -8114,9 +8291,8 @@ def interaction_operator(self): sites=(0,), ) - def interaction_term(self, site, *, U=None): + def interaction_term(self, site, *, U): """Return ``U n_up n_down`` on one physical site.""" - U = self.U if U is None else U U = _node_parameter(U, site) return self.operator_term( [(U, ((site, "double"),))], @@ -8134,9 +8310,8 @@ def chemical_potential_operator(self): terms = [(1.0, ((0, "number"),))] return self.operator_term(terms, sites=(0,)) - def chemical_potential_term(self, site, *, mu=None): + def chemical_potential_term(self, site, *, mu): """Return ``-mu n`` on one physical site.""" - mu = self.mu if mu is None else mu if self.spinful: mu = _node_parameter(mu, site) mu_up, mu_down = _as_spin_pair(mu, name="mu") @@ -8148,12 +8323,12 @@ def chemical_potential_term(self, site, *, mu=None): terms = [(-_node_parameter(mu, site), ((site, "number"),))] return self.operator_term(terms, sites=(site,)) - def onsite_term(self, site, *, U=None, mu=None): + def onsite_term(self, site, *, U=None, mu=0.0): """Return ``U n_up n_down - mu n`` on one site.""" - U = self.U if U is None else U - mu = self.mu if mu is None else mu terms = [] if self.spinful: + if U is None: + raise TypeError("onsite_term requires explicit U=... for spinful fermions.") terms.append((_node_parameter(U, site), ((site, "double"),))) mu = _node_parameter(mu, site) mu_up, mu_down = _as_spin_pair(mu, name="mu") @@ -8180,13 +8355,12 @@ def density_operator(self): ] return self.operator_term(terms, sites=(0, 1)) - def density_term(self, edge, *, V=None): + def density_term(self, edge, *, V): """Return ``V n_i n_j`` on a physical edge.""" try: left, right = tuple(edge) except (TypeError, ValueError) as exc: raise ValueError("edge must contain exactly two site labels.") from exc - V = self.V if V is None else V V = _edge_parameter(V, left, right) if self.spinful: names = ("number_up", "number_down") @@ -8334,7 +8508,7 @@ def heisenberg_gate(self, theta, *, edge=None, imaginary=False): szz_gate = spin_z_correlator_gate xy_gate = xy_exchange_gate - def interaction_gate(self, dt, *, site=None, U=None, imaginary=False): + def interaction_gate(self, dt, *, site=None, U, imaginary=False): """Return the exact onsite interaction gate. With a site-dependent ``U`` mapping or callable, pass ``site`` so the @@ -8346,7 +8520,6 @@ def interaction_gate(self, dt, *, site=None, U=None, imaginary=False): "Spinless fermions have no onsite doublon interaction; use " "density_gate(...) for the nearest-neighbor V interaction." ) - U = self.U if U is None else U U = U if site is None else _node_parameter(U, site) theta = dt * U @@ -8360,20 +8533,20 @@ def build(): return self._cached_gate(("interaction", dt, site, U, imaginary), build) - def onsite_gate(self, dt, *, site=None, U=None, mu=None, imaginary=False): + def onsite_gate(self, dt, *, site=None, U=None, mu=0.0, imaginary=False): """Return the complete one-site Hubbard gate. The generated gate represents ``U n_up n_down - mu n`` for spinful fermions and ``-mu n`` for spinless fermions. ``U`` and ``mu`` may be site-dependent mappings or callables when ``site`` is supplied. """ - U = self.U if U is None else U - mu = self.mu if mu is None else mu if site is not None: U = _node_parameter(U, site) mu = _node_parameter(mu, site) if self.spinful: + if U is None: + raise TypeError("onsite_gate requires explicit U=... for spinful fermions.") mu_up, mu_down = _as_spin_pair(mu, name="mu") U_site = U diagonal = ( @@ -8398,10 +8571,8 @@ def build(): return self._cached_gate(("onsite", dt, site, U, mu, imaginary), build) - def hopping_gate(self, dt, *, t=None, peierls_angle=0.0, imaginary=False): + def hopping_gate(self, dt, *, t, peierls_angle=0.0, imaginary=False): """Return a two-site native fermionic hopping gate with Peierls phase.""" - t = self.t if t is None else t - def build(): if not self.spinful: gate = _spinless_hopping_gate( @@ -8423,13 +8594,12 @@ def build(): return self._cached_gate(("hopping", dt, t, peierls_angle, imaginary), build) - def density_gate(self, dt, *, V=None, imaginary=False): + def density_gate(self, dt, *, V, imaginary=False): """Return the nearest-neighbor density interaction gate. For spinless fermions this is ``V n_i n_j``. For spinful fermions it is ``V (n_up + n_down)_i (n_up + n_down)_j``. """ - V = self.V if V is None else V theta = dt * V def build(): @@ -8445,9 +8615,8 @@ def build(): return self._cached_gate(("density", dt, V, imaginary), build) - def chemical_potential_gate(self, dt, *, mu=None, site=None, imaginary=False): + def chemical_potential_gate(self, dt, *, mu, site=None, imaginary=False): """Return the chemical-potential part of an onsite gate.""" - mu = self.mu if mu is None else mu mu = mu if site is None else _node_parameter(mu, site) if self.spinful: mu_up, mu_down = _as_spin_pair(mu, name="mu") @@ -8692,47 +8861,52 @@ def gate_stream( imaginary=False, t=None, U=None, - V=None, - mu=None, + V=0.0, + mu=0.0, ): - """Return a canonical first- or second-order fermion gate stream.""" + """Return a canonical fermion gate stream with explicit couplings.""" if order not in {1, 2}: raise ValueError("order must be 1 or 2.") + if t is None: + raise TypeError("gate_stream requires explicit t=... .") + if self.spinful and U is None: + raise TypeError("gate_stream requires explicit U=... for spinful fermions.") edges = _as_edges(edges) sites = _sites_from_edges(edges, sites) - target = self - if any(value is not None for value in (t, U, V, mu)): - target = type(self)( - symmetry=self.symmetry, - t=self.t if t is None else t, - U=self.U if U is None else U, - dtype=self.dtype, - to_backend=self.to_backend, - spinful=self.spinful, - V=self.V if V is None else V, - mu=self.mu if mu is None else mu, - ) if order == 2: - return target.strang_gate_stream( + return self.strang_gate_stream( edges, dt, sites=sites, peierls_angle=peierls_angle, imaginary=imaginary, + t=t, + U=U, + V=V, + mu=mu, ) entries = [] entries.extend( - (target.onsite_gate(dt, site=site, imaginary=imaginary), site) + ( + self.onsite_gate( + dt, + site=site, + U=U, + mu=mu, + imaginary=imaginary, + ), + site, + ) for site in sites ) - if target.V != 0 or isinstance(target.V, Mapping) or callable(target.V): + if V != 0 or isinstance(V, Mapping) or callable(V): entries.extend( ( - target.density_gate( + self.density_gate( dt, - V=_edge_parameter(target.V, left, right), + V=_edge_parameter(V, left, right), imaginary=imaginary, ), (left, right), @@ -8741,9 +8915,9 @@ def gate_stream( ) entries.extend( ( - target.hopping_gate( + self.hopping_gate( dt, - t=_edge_parameter(target.t, left, right), + t=_edge_parameter(t, left, right), peierls_angle=_edge_angle_parameter(peierls_angle, left, right), imaginary=imaginary, ), @@ -8753,7 +8927,7 @@ def gate_stream( ) return SymGateStream( entries, - hamiltonian=target.hamiltonian(edges), + hamiltonian=self.hamiltonian(edges, t=t, U=U, V=V, mu=mu), dt=dt, imaginary=imaginary, order=1, @@ -8767,84 +8941,66 @@ def strang_gate_stream( sites=None, peierls_angle=0.0, imaginary=False, + t=None, + U=None, + V=0.0, + mu=0.0, ): - """Return an edge-coloured second-order native fermionic gate stream.""" + """Return an edge-coloured second-order stream with explicit couplings.""" + if t is None: + raise TypeError("strang_gate_stream requires explicit t=... .") + if self.spinful and U is None: + raise TypeError( + "strang_gate_stream requires explicit U=... for spinful fermions." + ) edges = _as_edges(edges) sites = _sites_from_edges(edges, sites) half_dt = dt / 2 layers = self.edge_coloring_layers(edges) entries = [ - (self.onsite_gate(half_dt, site=site, imaginary=imaginary), site) + ( + self.onsite_gate( + half_dt, + site=site, + U=U, + mu=mu, + imaginary=imaginary, + ), + site, + ) for site in sites ] - if self.V != 0 or isinstance(self.V, Mapping) or callable(self.V): + if V != 0 or isinstance(V, Mapping) or callable(V): entries.extend( ( self.density_gate( half_dt, - V=_edge_parameter(self.V, left, right), + V=_edge_parameter(V, left, right), imaginary=imaginary, ), (left, right), ) for left, right in edges ) - if not self.spinful: - for layer in layers: - entries.extend( - ( - self.hopping_gate( - half_dt, - t=_edge_parameter(self.t, left, right), - peierls_angle=_edge_angle_parameter(peierls_angle, left, right), - imaginary=imaginary, - ), - (left, right), - ) - for left, right in layer - ) - for layer in reversed(layers): - entries.extend( - ( - self.hopping_gate( - half_dt, - t=_edge_parameter(self.t, left, right), - peierls_angle=_edge_angle_parameter(peierls_angle, left, right), - imaginary=imaginary, - ), - (left, right), - ) - for left, right in layer - ) + for layer in layers: entries.extend( ( - self.density_gate( + self.hopping_gate( half_dt, - V=_edge_parameter(self.V, left, right), + t=_edge_parameter(t, left, right), + peierls_angle=_edge_angle_parameter(peierls_angle, left, right), imaginary=imaginary, ), (left, right), ) - for left, right in edges - ) - entries.extend( - (self.chemical_potential_gate(half_dt, site=site, imaginary=imaginary), site) - for site in sites - ) - stream = SymGateStream( - entries, - hamiltonian=self.hamiltonian(edges), - dt=dt, - imaginary=imaginary, - order=2, + for left, right in layer ) - return stream - for layer in layers: + for layer in reversed(layers): entries.extend( ( self.hopping_gate( half_dt, - t=_edge_parameter(self.t, left, right), + t=_edge_parameter(t, left, right), peierls_angle=_edge_angle_parameter(peierls_angle, left, right), imaginary=imaginary, ), @@ -8852,69 +9008,147 @@ def strang_gate_stream( ) for left, right in layer ) - for layer in reversed(layers): + if V != 0 or isinstance(V, Mapping) or callable(V): entries.extend( ( - self.hopping_gate( + self.density_gate( half_dt, - t=_edge_parameter(self.t, left, right), - peierls_angle=_edge_angle_parameter(peierls_angle, left, right), + V=_edge_parameter(V, left, right), imaginary=imaginary, ), (left, right), ) - for left, right in layer + for left, right in edges ) entries.extend( - (self.onsite_gate(half_dt, site=site, imaginary=imaginary), site) + ( + self.onsite_gate( + half_dt, + site=site, + U=U, + mu=mu, + imaginary=imaginary, + ), + site, + ) for site in sites ) return SymGateStream( entries, - hamiltonian=self.hamiltonian(edges), + hamiltonian=self.hamiltonian(edges, t=t, U=U, V=V, mu=mu), dt=dt, imaginary=imaginary, order=2, ) - def hamiltonian(self, edges, **params): - """Build a native Hamiltonian from edges or explicit local terms. + def _validate_hamiltonian_terms(self, terms): + """Validate native local terms against this Fermion's local space.""" + terms = dict(terms) + coordinate_sites = _term_mapping_uses_coordinate_sites(terms) + expected_physical = { + charge: int(size) + for charge, size in self.physical_sectors.items() + } + backends = set() - Passing an edge iterable builds the configured Fermi-Hubbard model. - Passing a mapping ``{site_or_edge: operator}`` preserves the supplied - coefficient-free operators and their locations in a - :class:`SymHamiltonian`; that container can then be passed to - ``to_mpo`` without losing fermionic ordering metadata. + for where, term in terms.items(): + support = _as_term_where( + where, + coordinate_sites=coordinate_sites, + ) + if not _is_fermionic_symmray_array(term): + raise TypeError( + "Fermion.hamiltonian requires native fermionic Symmray " + f"arrays; term at {where!r} is {type(term).__name__}." + ) + if str(getattr(term, "symmetry", None)) != self.symmetry: + raise ValueError( + f"Term at {where!r} has symmetry " + f"{getattr(term, 'symmetry', None)!r}, expected " + f"{self.symmetry!r}." + ) + indices = tuple(getattr(term, "indices", ())) + expected_rank = 2 * len(support) + if len(indices) != expected_rank: + raise ValueError( + f"Term at {where!r} has rank {len(indices)}, but its " + f"{len(support)}-site key requires rank {expected_rank}." + ) + for axis, index in enumerate(indices): + actual_physical = { + charge: int(size) + for charge, size in dict(getattr(index, "chargemap", {})).items() + } + if actual_physical != expected_physical: + raise ValueError( + f"Term at {where!r} axis {axis} has physical sectors " + f"{actual_physical!r}, expected {expected_physical!r}." + ) + for block in getattr(term, "blocks", {}).values(): + backends.add(ar.infer_backend(block)) + + if len(backends) > 1: + raise TypeError( + "Fermion.hamiltonian terms use mixed array backends " + f"{sorted(backends)!r}. Supply to_backend=... so every native " + "block is converted consistently." + ) + + def hamiltonian( + self, + terms_or_edges, + *, + t=None, + U=None, + V=0.0, + mu=0.0, + flat=False, + to_backend=None, + ): + """Validate explicit terms or build a model only from explicit couplings. + + The canonical form is a mapping from one-site or two-site locations to + native fermionic Symmray arrays. It is checked for symmetry, physical + sectors, support rank, and backend consistency before being bundled in + a :class:`SymHamiltonian`. Passing lattice edges remains a compact + convenience, but requires its couplings explicitly; no coupling is + stored on :class:`Fermion`. """ - to_backend = params.pop("to_backend", self.to_backend) - if isinstance(edges, Mapping): + to_backend = self.to_backend if to_backend is None else to_backend + if isinstance(terms_or_edges, Mapping): + if any(value is not None for value in (t, U)) or V != 0 or mu != 0: + raise TypeError( + "When passing explicit terms, put every coupling in the " + "native arrays rather than passing t/U/V/mu again." + ) + terms = _apply_to_hamiltonian_terms(terms_or_edges, to_backend) + self._validate_hamiltonian_terms(terms) return SymHamiltonian.from_terms( self.model, self.symmetry, - edges, - to_backend=to_backend, - parameters=params, + terms, + parameters={}, ) - flat = params.pop("flat", False) + + if t is None: + raise TypeError("hamiltonian(edges, ...) requires explicit t=... .") + if self.spinful and U is None: + raise TypeError( + "hamiltonian(edges, ...) requires explicit U=... for spinful fermions." + ) + params = {"t": t, "V": V, "mu": mu} if self.spinful: - params = { - "t": self.t, - "U": self.U, - "mu": self.mu, - **params, - } - if self.V != 0 or "V" in params: - params.setdefault("V", self.V) - else: - params = {"t": self.t, "V": self.V, "mu": self.mu, **params} - return SymHamiltonian.from_edges( + params["U"] = U + hamiltonian = SymHamiltonian.from_edges( self.model, self.symmetry, - edges, + terms_or_edges, flat=flat, to_backend=to_backend, **params, ) + self._validate_hamiltonian_terms(hamiltonian.terms) + return hamiltonian def local_terms(self, edges, *, layout="site", **params): """Return native local terms for site or qMERA energy workflows. @@ -8931,8 +9165,6 @@ def local_terms(self, edges, *, layout="site", **params): if layout in {"site", "sites", "native"}: return self.hamiltonian(edges, **params).terms if layout in {"qmera", "qmera_modes", "modes"}: - if not self.spinful: - raise ValueError("qMERA Hubbard terms require spinful fermions.") from ..optimizers.mera import ( # pylint: disable=import-outside-toplevel qmera_symmray_fermi_hubbard_terms, ) @@ -8942,15 +9174,29 @@ def local_terms(self, edges, *, layout="site", **params): fermion=self, **params, ) + if layout in {"majorana", "qmera_majorana"}: + from ..optimizers.mera import ( # pylint: disable=import-outside-toplevel + qmera_symmray_majorana_terms, + ) + + return qmera_symmray_majorana_terms( + edges, + fermion=self, + **params, + ) raise ValueError( "layout must be 'site' for native site terms or 'qmera' for " - "two-state qMERA mode terms." + "two-state qMERA mode terms, or 'majorana'." ) def qmera_terms(self, geometry, **params): """Return the explicit two-state qMERA terms for ``geometry``.""" return self.local_terms(geometry, layout="qmera", **params) + def majorana_terms(self, geometry, **params): + """Return parity-preserving Majorana terms for a qMERA geometry.""" + return self.local_terms(geometry, layout="majorana", **params) + # Kept for callers that adopted the initial public names before the helper was # generalized to the operator/gate-focused ``Fermion`` API. diff --git a/tests/test_optimize_mera.py b/tests/test_optimize_mera.py index 7152248..8be752c 100644 --- a/tests/test_optimize_mera.py +++ b/tests/test_optimize_mera.py @@ -12,12 +12,18 @@ QMeraBlockSpec, QMeraBuilder, QMeraCompiledLightconeChunk, + QMeraContractionPathCache, + QMeraDisentanglerSpec, QMeraGeometry, + QMeraIsometrySpec, + QMeraLightconeGroup, QMeraLightconeTN, QMeraSchematicBlock, QMeraParametricEnergyOptimizer, QMeraParametricLightconeChunk, QMeraSymmrayFermionBackend, + QMeraScaleSpec, + QMeraUnitarySpec, UserGateFamily, build_lightcone_chunks, build_qmera_contraction_optimizer, @@ -25,6 +31,8 @@ build_qmera_parametric_lightcone_chunks, compile_qmera_parametric_lightcones, contract_qmera_lightcone_tn, + group_qmera_parametric_lightcone_chunks, + lightcone_energy, default_gate_registry, draw_qmera_schedule, local_qmera_compiled_lightcone_expectation, @@ -32,12 +40,15 @@ local_lightcone_expectation, normalize_local_terms, qmera_compiled_parametric_energy, + qmera_direct_parametric_energy, qmera_parametric_energy, qmera_parametric_lightcone_state, qmera_parametric_lightcone_tn, qmera_schematic_blocks, qmera_symmray_fermi_hubbard_terms, + qmera_symmray_majorana_terms, symmray_fermion_gate_registry, + symmray_majorana_gate_registry, ) from pepsy.optimizers.mera.optimizer import ( MeraEnergyOptimizer as ModuleMeraEnergyOptimizer, @@ -94,6 +105,40 @@ def fake_build_qmera_contraction_optimizer(**kwargs): assert calls == [{"directory": "/tmp/qmera-cache"}] +def test_qmera_builder_infers_fermion_register_convention(): + """A stored Fermion model should supply qMERA mode metadata.""" + fermion = Fermion(spinful=True, symmetry="U1U1") + builder = QMeraBuilder(shape=(2, 2), fermion=fermion) + + assert builder.fermion is fermion + assert builder.geometry.site_modes == ("up", "down") + assert builder.geometry.mode_order == "mode-major" + assert builder.geometry.num_modes == 8 + + backend = QMeraSymmrayFermionBackend.from_fermion(fermion) + assert backend.symmetry == "U1U1" + assert backend.site_modes == ("up", "down") + assert backend.mode_order == "mode-major" + + +def test_qmera_builder_preserves_explicit_geometry_override(): + """Advanced callers may select a different register order explicitly.""" + fermion = Fermion(spinful=True, symmetry="U1U1") + geometry = QMeraGeometry( + shape=(2, 2), + site_modes=("up", "down"), + mode_order="site-major", + ) + builder = QMeraBuilder(geometry=geometry, fermion=fermion) + + assert builder.geometry is geometry + assert builder.geometry.mode_order == "site-major" + + bad_geometry = QMeraGeometry(shape=(2, 2), site_modes=("mode",)) + with pytest.raises(ValueError, match="site_modes must match"): + QMeraBuilder(geometry=bad_geometry, fermion=fermion) + + def test_normalize_local_terms_accepts_mapping_iterable_and_local_term(): """Hamiltonian input should normalize to explicit LocalTerm objects.""" op = _zz_term() @@ -169,6 +214,59 @@ def test_mera_energy_loss_matches_quimb_exact_sum(): assert complex(opt.loss(real=False)) == pytest.approx(complex(direct)) +def test_generic_lightcone_energy_groups_select_gate_and_contract(): + """The public fixed-state helper should match the full MERA oracle.""" + mera = _small_mera(seed=241) + terms = {(0, 1): _zz_term(), (2, 3): _zz_term()} + direct = mera.compute_local_expectation_exact( + terms, + optimize="auto-hq", + normalized=True, + ) + + value = lightcone_energy( + mera, + terms, + energy_per_site=False, + normalized=True, + real=False, + group_terms=True, + ) + + assert complex(value) == pytest.approx(complex(direct)) + + +def test_qmera_parameter_sharing_per_block_reuses_round_parameters(): + """One block can share parameters across its brickwall rounds.""" + unitary = QMeraUnitarySpec( + gate_family="rxx", + parameter_sharing="per-block", + ) + builder = QMeraBuilder( + shape=(4, 4), + disentangler=QMeraDisentanglerSpec( + block_shape=(2, 2), + unitary=unitary, + circuit_depth=2, + ), + isometry=QMeraIsometrySpec( + block_shape=(2, 2), + unitary=unitary, + circuit_depth=2, + ), + max_layers=1, + ) + layer = builder.build_schedule().layers[0] + + for placements in (layer.disentanglers, layer.isometries): + by_block = {} + for placement in placements: + by_block.setdefault(placement.block, set()).add(placement.param_key) + assert by_block + assert all(len(keys) == 1 for keys in by_block.values()) + assert len({next(iter(keys)) for keys in by_block.values()}) == len(by_block) + + def test_mera_energy_estimate_reports_lightcone_metadata(): """energy() should return the shared EnergyEstimate dataclass.""" mera = _small_mera(seed=25) @@ -361,6 +459,131 @@ def test_qmera_2d_schedule_uses_rg_blocks_and_face_disentanglers(): assert any(placement.axis == "y" for placement in first.isometries) +def test_qmera_explicit_specs_build_4x4_periodic_hubbard_schedule(): + """Explicit square layers should include all 4x4 PBC interfaces.""" + backend = QMeraSymmrayFermionBackend( + symmetry="U1U1", + site_modes=("up", "down"), + mode_order="mode-major", + ) + unitary = QMeraUnitarySpec( + gate_family="symmray-hubbard", + family="fermion", + arity_kind="mode", + symmetry="U1U1", + preserves_parity=True, + parameter_sharing="per-axis", + metadata={"model": "fermi-hubbard", "term": "hopping"}, + ) + builder = QMeraBuilder( + shape=(4, 4), + boundary="periodic", + site_modes=backend.site_modes, + mode_order="mode-major", + gate_registry=symmray_fermion_gate_registry(backend=backend), + disentangler=QMeraDisentanglerSpec( + block_shape=(2, 2), + unitary=unitary, + placement="boundary-square", + circuit_depth=2, + ), + isometry=QMeraIsometrySpec( + block_shape=(2, 2), + unitary=unitary, + circuit_depth=2, + ), + max_layers=2, + ) + + schedule = builder.build_schedule() + first, second = schedule.layers + + assert schedule.num_scales == 2 + assert [len(layer.input_sites) for layer in schedule.layers] == [32, 8] + assert [len(layer.output_sites) for layer in schedule.layers] == [8, 2] + assert [len(layer.isometry_blocks) for layer in schedule.layers] == [4, 1] + assert len(first.disentangler_blocks) == 8 + assert not second.disentanglers + assert schedule.disentangler.placement == "boundary-square" + assert schedule.disentangler.unitary_spec is unitary + assert schedule.isometry.implementation == "unitary-completion" + assert {placement.axis for placement in first.disentanglers} == {"x", "y"} + assert len(builder.initialize_parameters(schedule)) == len( + set(schedule.param_keys) + ) + + physical_supports = [ + {geometry_site for geometry_site in map(schedule.geometry.to_site, block)} + for block in first.disentangler_blocks + ] + assert all( + len({site[0] for site in support}) == 2 + and len({site[1] for site in support}) == 2 + for support in physical_supports + ) + assert any( + {site[0] for site in support} == {0, 3} + for support in physical_supports + ) + assert any( + {site[1] for site in support} == {0, 3} + for support in physical_supports + ) + + +def test_qmera_explicit_layer_spec_rejects_true_isometry_until_supported(): + """The public API should not silently call a unitary a rectangular isometry.""" + with pytest.raises(NotImplementedError, match="true-isometry"): + QMeraIsometrySpec(implementation="true-isometry").to_block_spec() + + +def test_qmera_scale_plan_supports_heterogeneous_6x6_periodic_layers(): + """A scale plan should express 2x2 then 3x3 RG blocks and vertical strips.""" + scale_plan = ( + QMeraScaleSpec( + name="6x6-to-3x3", + disentangler=QMeraDisentanglerSpec( + block_shape=(2, 2), + placement="boundary-square", + ), + isometry=QMeraIsometrySpec(block_shape=(2, 2)), + ), + QMeraScaleSpec( + name="3x3-to-1", + disentangler=QMeraDisentanglerSpec( + block_shape=3, + orientation="vertical", + placement="within-block", + circuit_depth=3, + ), + isometry=QMeraIsometrySpec(block_shape=(3, 3)), + ), + ) + schedule = QMeraBuilder( + shape=(6, 6), + boundary="periodic", + scales=scale_plan, + ).build_schedule() + + first, second = schedule.layers + assert [len(layer.input_sites) for layer in schedule.layers] == [36, 9] + assert [len(layer.output_sites) for layer in schedule.layers] == [9, 1] + assert [len(layer.isometry_blocks) for layer in schedule.layers] == [9, 1] + assert len(first.disentangler_blocks) == 18 + assert len(second.disentangler_blocks) == 3 + assert second.disentangler_spec.orientation == "y" + assert second.disentangler_spec.placement == "within-block" + assert {placement.axis for placement in second.disentanglers} == {"y"} + assert any( + {schedule.geometry.to_site(site)[1] for site in placement.where} == {0, 4} + for placement in second.disentanglers + ) + assert [scale.name for scale in schedule.scale_specs] == [ + "6x6-to-3x3", + "3x3-to-1", + ] + + def test_qmera_1d_mode_geometry_schedules_register_modes(): """1D qMERA schedules should operate on mode/register positions.""" builder = QMeraBuilder( @@ -393,6 +616,26 @@ def test_qmera_1d_mode_geometry_schedules_register_modes(): assert "I3" in ansatz.state.tags +def test_qmera_1d_mode_schedule_never_pairs_different_fermion_modes(): + """1D brickwall layers should preserve each explicit mode flavor.""" + builder = QMeraBuilder( + shape=4, + site_modes=("up", "down"), + mode_order="mode-major", + gate_family="fsim", + isometry_gate_family="fsim", + max_layers=2, + ) + schedule = builder.build_schedule() + + for placement in schedule.placements: + modes = { + schedule.geometry.to_mode(register_site)[1] + for register_site in placement.where + } + assert len(modes) == 1 + + def test_qmera_symmray_fermion_backend_builds_native_hubbard_terms(): """qMERA Hubbard terms should be native Symmray fermionic mode arrays.""" pytest.importorskip("symmray") @@ -457,6 +700,13 @@ def test_unified_fermion_qmera_optimizer_runs_torch_autodiff(): array_backend = backend_torch(dtype=torch.complex128) backend = QMeraSymmrayFermionBackend(to_backend=array_backend) registry = symmray_fermion_gate_registry(backend=backend) + fermion = Fermion( + spinful=True, + symmetry="U1U1", + t=0.2, + U=0.5, + mu=0.1, + ) def product_state_factory(schedule, sites, **kwargs): return backend.product_state( @@ -468,8 +718,7 @@ def product_state_factory(schedule, sites, **kwargs): builder = QMeraBuilder( shape=2, - site_modes=backend.site_modes, - mode_order="mode-major", + fermion=fermion, gate_registry=registry, array_backend=array_backend, disentangler={"block_size": 2, "circuit_depth": 0}, @@ -483,17 +732,8 @@ def product_state_factory(schedule, sites, **kwargs): param_scale=0.01, product_state_factory=product_state_factory, ) - fermion = Fermion( - spinful=True, - symmetry="U1U1", - t=0.2, - U=0.5, - mu=0.1, - ) - - terms = builder.fermion_terms(fermion) + terms = builder.fermion_terms() optimizer = builder.fermion_parametric_optimizer( - fermion, energy_per_site=False, ) initial = optimizer.loss(energy_per_site=False) @@ -583,7 +823,7 @@ def test_qmera_symmray_fsim_runs_full_fermionic_lightcone(): ) assert complex(value) == pytest.approx(0.0) - bad_builder = QMeraBuilder( + single_site_builder = QMeraBuilder( shape=1, site_modes=backend.site_modes, gate_registry=registry, @@ -594,11 +834,13 @@ def test_qmera_symmray_fsim_runs_full_fermionic_lightcone(): }, max_layers=1, ) - with pytest.raises(ValueError, match="spin-changing"): - bad_builder.gate_tensors( - bad_builder.initialize_parameters(), - bad_builder.build_schedule(), - ) + single_site_schedule = single_site_builder.build_schedule() + assert not single_site_schedule.layers + assert not single_site_schedule.placements + assert single_site_builder.gate_tensors( + single_site_builder.initialize_parameters(single_site_schedule), + single_site_schedule, + ) == {} def test_qmera_symmray_fermion_lightcone_contracts_native_term(): @@ -638,6 +880,14 @@ def test_qmera_symmray_fermion_lightcone_contracts_native_term(): energy_per_site=False, real=False, ) + fixed_state_value = lightcone_energy( + builder.build(params), + terms[:1], + schedule=schedule, + convert_terms=False, + energy_per_site=False, + real=False, + ) assert isinstance(lightcone, QMeraLightconeTN) assert lightcone.ket.num_tensors == 2 @@ -647,14 +897,258 @@ def test_qmera_symmray_fermion_lightcone_contracts_native_term(): ) assert complex(value) == pytest.approx(0.0) assert complex(builder_value) == pytest.approx(complex(value)) + assert complex(fixed_state_value) == pytest.approx(complex(value)) -def test_qmera_2d_multi_mode_schedule_requires_explicit_design(): - """2D multi-mode qMERA should fail before silently dropping modes.""" - builder = QMeraBuilder(shape=(2, 2), site_modes=("up", "down")) +def test_qmera_2d_multi_mode_schedule_retains_modes_and_pairs_like_modes(): + """2D RG blocks should retain modes and never pair different flavors.""" + builder = QMeraBuilder( + shape=(2, 2), + site_modes=("up", "down"), + mode_order="mode-major", + disentangler={"block_size": 2, "circuit_depth": 1}, + isometry={"block_size": (2, 2), "circuit_depth": 1}, + max_layers=1, + ) - with pytest.raises(NotImplementedError, match="mode-blocking"): - builder.build_schedule() + schedule = builder.build_schedule() + first = schedule.layers[0] + + assert first.isometry_blocks[0] == (0, 4, 1, 5, 2, 6, 3, 7) + assert first.output_sites == (0, 4) + for placement in (*first.disentanglers, *first.isometries): + modes = [schedule.geometry.to_mode(site)[1] for site in placement.where] + assert len(set(modes)) == 1 + + +def test_qmera_2d_multimode_rg_keeps_populated_axis_after_coarse_graining(): + """An anisotropic coarse grid should still receive its final isometry.""" + builder = QMeraBuilder( + shape=(2, 4), + site_modes=("up", "down"), + mode_order="mode-major", + isometry={"block_size": (2, 2), "circuit_depth": 1}, + max_layers=2, + ) + + schedule = builder.build_schedule() + + assert len(schedule.layers) == 2 + assert schedule.layers[-1].isometries + assert schedule.layers[-1].isometries[0].axis == "y" + + +def test_fermion_majorana_convention_is_z2_and_parity_preserving(): + """Majoranas are odd Z2 operators; bilinears and gates are neutral.""" + pytest.importorskip("symmray") + majorana = Fermion(spinful=False, symmetry="Z2") + + gamma_x = majorana.majorana_operator("x", site=0) + gamma_y = majorana.majorana_operator("y", site=0) + bilinear = majorana.majorana_bilinear_operator( + (0, 1), + left_component="y", + right_component="x", + ) + pairing = majorana.pairing_operator((0, 1), phase=0.25) + gates = ( + majorana.majorana_gate(0.1, edge=(0, 1)), + majorana.pairing_gate(0.1, edge=(0, 1), phase=0.25), + ) + + assert gamma_x.charge == 1 + assert gamma_y.charge == 1 + assert bilinear.charge == 0 + assert pairing.charge == 0 + gamma_x_dense = np.asarray(gamma_x.to_dense()) + gamma_y_dense = np.asarray(gamma_y.to_dense()) + np.testing.assert_allclose(gamma_x_dense @ gamma_x_dense, np.eye(2)) + np.testing.assert_allclose(gamma_y_dense @ gamma_y_dense, np.eye(2)) + np.testing.assert_allclose( + gamma_x_dense @ gamma_y_dense + gamma_y_dense @ gamma_x_dense, + np.zeros((2, 2)), + ) + for operator in (bilinear, pairing): + dense = np.asarray(operator.to_dense()) + matrix = dense.reshape((4, 4)) + np.testing.assert_allclose(matrix, matrix.conj().T) + for gate in gates: + dense = np.asarray(gate.to_dense()).reshape((4, 4)) + np.testing.assert_allclose(dense.conj().T @ dense, np.eye(4), atol=1.0e-12) + assert all("Z2FermionicArray" in type(value).__name__ for value in ( + gamma_x, + gamma_y, + bilinear, + pairing, + *gates, + )) + with pytest.raises(ValueError, match="requires symmetry='Z2'"): + Fermion(spinful=False, symmetry="U1").majorana_operator() + + +def test_qmera_2d_fermion_and_majorana_direct_oracles_match_lightcones(): + """Native graded 2D Hubbard and Majorana paths agree with direct TNs.""" + pytest.importorskip("symmray") + + def run_case(geometry, backend, registry, fermion, terms, gate_family): + def product_state_factory(schedule, sites, **kwargs): + occupations = { + site: int( + (sum(schedule.geometry.to_site(site)) % 2 == 0) + == (schedule.geometry.to_mode(site)[-1] == "up") + ) + for site in sites + } + return backend.product_state( + schedule, + sites, + occupations=occupations, + **kwargs, + ) + + builder = QMeraBuilder( + geometry=geometry, + gate_registry=registry, + gate_family=gate_family, + disentangler={ + "block_size": 2, + "circuit_depth": 1, + "gate_family": gate_family, + }, + isometry={ + "block_size": (2, 2), + "circuit_depth": 1, + "gate_family": gate_family, + }, + max_layers=1, + seed=91, + param_scale=0.01, + product_state_factory=product_state_factory, + ) + schedule = builder.build_schedule() + parameters = builder.initialize_parameters(schedule) + lightcone = builder.parametric_loss( + parameters, + terms, + schedule=schedule, + convert_terms=False, + energy_per_site=False, + real=False, + ) + direct = builder.direct_parametric_loss( + parameters, + terms, + schedule=schedule, + convert_terms=False, + energy_per_site=False, + real=False, + ) + return lightcone, direct + + hubbard_geometry = QMeraGeometry( + shape=(2, 2), + site_modes=("up", "down"), + mode_order="mode-major", + ) + hubbard_backend = QMeraSymmrayFermionBackend() + hubbard_registry = symmray_fermion_gate_registry(backend=hubbard_backend) + hubbard = Fermion( + spinful=True, + symmetry="U1U1", + t=0.2, + U=0.5, + mu=0.1, + ) + hubbard_terms = qmera_symmray_fermi_hubbard_terms( + hubbard_geometry, + fermion=hubbard, + ) + hubbard_lightcone, hubbard_direct = run_case( + hubbard_geometry, + hubbard_backend, + hubbard_registry, + hubbard, + hubbard_terms, + "symmray-fsim", + ) + + majorana_geometry = QMeraGeometry( + shape=(2, 2), + site_modes=("mode",), + mode_order="mode-major", + ) + majorana_backend = QMeraSymmrayFermionBackend( + symmetry="Z2", + site_modes=("mode",), + ) + majorana_registry = symmray_majorana_gate_registry(backend=majorana_backend) + majorana = Fermion(spinful=False, symmetry="Z2") + majorana_terms = qmera_symmray_majorana_terms( + majorana_geometry, + fermion=majorana, + coupling=0.4, + pairing=0.2, + ) + majorana_lightcone, majorana_direct = run_case( + majorana_geometry, + majorana_backend, + majorana_registry, + majorana, + majorana_terms, + "symmray-majorana", + ) + + assert complex(hubbard_lightcone) == pytest.approx(complex(hubbard_direct)) + assert complex(majorana_lightcone) == pytest.approx(complex(majorana_direct)) + + +def test_qmera_grouped_and_direct_energy_match_schedule_lightcones(): + """Grouping and the full direct-gate oracle must preserve local energy.""" + builder = QMeraBuilder(shape=8, seed=12, param_scale=0.02) + schedule = builder.build_schedule() + parameters = builder.initialize_parameters(schedule) + terms = {(0, 1): _zz_term(), (2, 3): _zz_term(), (4, 5): _zz_term()} + chunks = builder.parametric_lightcone_chunks(terms, schedule) + groups = group_qmera_parametric_lightcone_chunks(chunks) + + assert all(isinstance(group, QMeraLightconeGroup) for group in groups) + assert sum(group.num_terms for group in groups) == len(chunks) + local = builder.parametric_loss( + parameters, + terms, + schedule=schedule, + energy_per_site=False, + ) + direct = qmera_direct_parametric_energy( + schedule, + parameters, + terms, + energy_per_site=False, + ) + assert complex(local) == pytest.approx(complex(direct)) + + +def test_qmera_contraction_path_cache_reuses_topology_optimizer(monkeypatch): + """Path caches should lazily create one reusable optimizer per topology.""" + calls = [] + + def fake_builder(**kwargs): + calls.append(kwargs) + return object() + + monkeypatch.setattr( + "pepsy.optimizers.mera.cache.build_qmera_contraction_optimizer", + fake_builder, + ) + cache = QMeraContractionPathCache({"directory": False}) + first = cache.optimizer_for(("cone",)) + second = cache.optimizer_for(("cone",)) + other = cache.resolve("auto-hq", key=("other",)) + + assert first is second + assert other is not first + assert cache.num_cached_paths == 2 + assert calls == [{"directory": False}, {"directory": False}] def test_qmera_gate_registry_generates_parametrized_two_qubit_gates(): @@ -1230,3 +1724,19 @@ def test_qmera_2d_draw_schematic_builds_quimb_drawing(): ) assert isinstance(drawing, schematic.Drawing) + clean_drawing = builder.draw_schematic( + layer=0, + style="clean", + label_sites=False, + scale_figsize=False, + ) + register_drawing = builder.draw_schematic( + layer=0, + style="register", + label_sites=False, + scale_figsize=False, + ) + assert isinstance(clean_drawing, schematic.Drawing) + assert isinstance(register_drawing, schematic.Drawing) + with pytest.raises(ValueError, match="style"): + builder.draw_schematic(layer=0, style="unknown") From 47683f11b246bf7f0dc25d2a2ddd357347b1c499 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Tue, 28 Jul 2026 16:42:58 -0700 Subject: [PATCH 14/70] tree: preserve canonical isometry metadata --- .github/skills/tree-optimizer/SKILL.md | 14 +- .../references/performance-layout.md | 10 +- .../skills/tree-stabilizer-optimizer/SKILL.md | 9 + docs/api/optimizers/tree.md | 26 ++- docs/api/optimizers/tree_stabilizer.md | 11 ++ src/pepsy/optimizers/tree/optimizer.py | 68 +++++-- src/pepsy/optimizers/tree/ttn.py | 127 +++++++++++- .../optimizers/tree_stabilizer/optimizer.py | 37 +++- src/pepsy/tensors/constructors.py | 7 + tests/test_optimize_tree.py | 186 ++++++++++++++++-- tests/test_optimize_tree_stabilizer.py | 111 +++++++++++ 11 files changed, 562 insertions(+), 44 deletions(-) diff --git a/.github/skills/tree-optimizer/SKILL.md b/.github/skills/tree-optimizer/SKILL.md index fb6f684..f97e24d 100644 --- a/.github/skills/tree-optimizer/SKILL.md +++ b/.github/skills/tree-optimizer/SKILL.md @@ -155,6 +155,13 @@ telescopes to identity between bra and ket. `shift_orthogonality_center` first peels that region with lossless QR and then walks only the remaining path. Do not regress this regional recovery to an unconditional O(N) recanonicalisation. +- Local isometry proofs live only on each tensor's ``left_inds``. + `TreeTensorNetwork.isometry_direction` / `isometry_map` derive read-only + orientations, `can_skip_canonize` recognizes an already-proven dense edge, + and `validate_isometry_metadata` checks alignment with the canonical region. + `TreeOptimizer` delegates these methods; never add a second mutable + optimizer-owned orientation map. Native fermionic edges always retain their + explicit graded QR and are never skipped through this dense metadata path. - `ttn.is_canonical_form(center)` verifies the invariant directly (every non-centre tensor is an isometry toward the centre) — use it in tests/diagnostics. - A freshly built product state is **already canonical at the root** (all @@ -272,8 +279,11 @@ covering range then compressed (quimb's `gate_with_submpo` is `MatrixProductStat state machine without repeating those QRs; native fermionic trees retain their explicit graded QR recovery. Finally make one depth-first canonical SVD sweep: every affected tree edge is truncated once, after the complete - operator has arrived. `renormalize=True` renormalises afterwards (for - Kraus/projection). + operator has arrived. Dense path and subtree sweeps select one-sided + ``reduced="left"`` compression only when the destination tensor's live + ``left_inds`` proves the required isometry; missing proofs and native + graded tensors use the full reduction. `renormalize=True` renormalises + afterwards (for Kraus/projection). State bonds are always read from the live tensors because gate application can rename them. New state message bonds are fresh per-update names, while operator diff --git a/.github/skills/tree-optimizer/references/performance-layout.md b/.github/skills/tree-optimizer/references/performance-layout.md index f89064d..5148826 100644 --- a/.github/skills/tree-optimizer/references/performance-layout.md +++ b/.github/skills/tree-optimizer/references/performance-layout.md @@ -15,8 +15,14 @@ Tree Optimizer skill so the upload-facing `SKILL.md` stays concise. ids against `self.tn.tensor_map`; a stale entry is recomputed safely. - Dense path and subtree routing preserve each QR-produced Q tensor's `left_inds`. Canonical recovery therefore recognizes an already-isometric - routed branch without repeating its decomposition. Native fermionic routing - deliberately retains explicit graded QR recovery. + routed branch without repeating its decomposition or entering Quimb's dense + canonicalization kernel. Path and subtree compression also reads that proof + before selecting one-sided `reduced="left"` compression, avoiding the + redundant reduction QR only when the destination tensor is proven + isometric. Missing proofs fall back to two-sided reduction. The network + derives orientation views directly from live tensors; do not cache a + duplicate map in the optimizer. Native fermionic routing deliberately + retains explicit graded QR/SVD recovery. - `copy()` shares the immutable `TreePlan`, owns `self.tn.copy()`, resets the tid cache, and derives a deterministic child seed for an independent RNG. diff --git a/.github/skills/tree-stabilizer-optimizer/SKILL.md b/.github/skills/tree-stabilizer-optimizer/SKILL.md index 4191840..594dc10 100644 --- a/.github/skills/tree-stabilizer-optimizer/SKILL.md +++ b/.github/skills/tree-stabilizer-optimizer/SKILL.md @@ -37,6 +37,15 @@ swap/split updates, and linear canonical metadata are chain-specific. Keep `to_statevector()` equal to `C @ p_dense` in logical big-endian order, and delegate `norm()` to the coefficient tree. +Keep coefficient canonicality single-owned as well. Local isometry proofs live +only on ``TreeTensorNetwork`` tensors' ``left_inds``; TreeStab exposes +read-only delegates to the tree API and must not cache another map. Direct, +MPO, and coefficient-frame sub-MPO paths inherit TreeOptimizer's +metadata-gated one-sided compression. Use ``apply_to_arrays`` for backend-only +conversion, and when a TreeStab constructor independently proves a canonical +tree (such as dense cap factorization), install the proven metadata without a +redundant numerical QR sweep. + Clifford events update `C`. Physical Pauli rotations, measurements, resets, and magic gadgets map through `C† P C` and update `|p>`. Coefficient-frame sub-MPO events go directly to `TreeOptimizer.apply_submpo`; they are not diff --git a/docs/api/optimizers/tree.md b/docs/api/optimizers/tree.md index a5608cd..2a2a3e3 100644 --- a/docs/api/optimizers/tree.md +++ b/docs/api/optimizers/tree.md @@ -88,9 +88,20 @@ the `orthogonality_center` name-parity alias), `shift_orthogonality_center(node) and `is_canonical_form(center)` delegate to the state, so the optimizer and its `TreeTensorNetwork` speak the same canonicalisation vocabulary. +Local isometry orientation also has one owner: each live Quimb tensor carries +its proven `left_inds`, while `TreeTensorNetwork.isometry_direction(node)` and +`isometry_map()` derive read-only node-to-neighbour views from those tensors. +`can_skip_canonize(a, b)` exposes the exact dense-edge condition used to avoid +an already-proven QR, and `validate_isometry_metadata()` checks the local +orientations against the tracked canonical region. `TreeOptimizer` delegates +the same four methods without maintaining another mutable map. Native +fermionic trees retain explicit graded QR and therefore never report a +skippable edge through this API. + `TreeTensorNetwork.validate()` checks the live tensor set, physical legs, tree edges, and bond ownership against the `TreePlan`; pass -`check_canonical=True` when the more expensive isometry check is also desired. +`check_canonical=True` when the metadata alignment and more expensive numerical +isometry check are also desired. Direct Quimb mutations such as `gate_inds_`, `canonize_between`, `compress_between`, and `canonize_around_` invalidate the tracked canonical region. Call `invalidate_canonical_form()` after mutating tensor data directly; @@ -600,8 +611,7 @@ finder = py.TreeLayoutFinder(gates, n=L, weight_mode="operator_schmidt") plan = finder.layered(block_size=4) state = py.TreeTensorNetwork.from_plan(plan) -for tensor in state.tensor_map.values(): - tensor.modify(data=to_backend(tensor.data)) +state.apply_to_arrays(to_backend) # backend-only conversion preserves left_inds # Convert user-provided gate arrays once, at their source. native_gates = [(to_backend(gate), where) for gate, where in gates] @@ -705,8 +715,14 @@ available through `truncation_report()`, `get_infidelities()`, and additionally normalized by their exact graded norm readout. - **Routed isometry reuse.** Dense geodesic and subtree QR routing retains each Q tensor's `left_inds`, allowing later canonical recovery to reuse the proven - isometry without repeating the decomposition. Native fermionic trees keep - their separate explicit graded QR path. + isometry without repeating the decomposition or entering Quimb's dense + canonicalization kernel. Final path and subtree compression also consults + that live proof: when the destination-side tensor is already isometric, + Quimb uses one-sided `reduced="left"` compression and avoids its redundant + reduction QR; otherwise it falls back to the full two-sided reduction. The + network derives orientation diagnostics from those tensors; the optimizer + does not keep a duplicate map. Native fermionic trees keep their separate + explicit graded QR/SVD path. - **State-owned centre.** The orthogonality centre lives on the `TreeTensorNetwork` (`orthogonality_center`, an `_EXTRA_PROPS` field), so the optimizer and the state cannot disagree and the centre is carried by diff --git a/docs/api/optimizers/tree_stabilizer.md b/docs/api/optimizers/tree_stabilizer.md index 71fb159..c0cd5fb 100644 --- a/docs/api/optimizers/tree_stabilizer.md +++ b/docs/api/optimizers/tree_stabilizer.md @@ -10,6 +10,17 @@ milestone. It represents the state as where `C` is a Stim tableau Clifford and `|p>` is a dense two-level `TreeTensorNetwork` evolved by `TreeOptimizer`. +Canonical and compression state has the same single owner as ordinary tree +simulation: local isometry proofs live on the coefficient tensors' +``left_inds`` and are interpreted by ``TreeTensorNetwork``. TreeStab delegates +``isometry_direction()``, ``isometry_map()``, ``can_skip_canonize()``, and +``validate_isometry_metadata()`` to its coefficient ``TreeOptimizer``; it does +not keep another map. Direct, MPO, and coefficient-frame sub-MPO routes +therefore reuse proven path/subtree Q tensors and select one-sided SVD +compression only when the live proof is valid. Backend conversion and dense +cap reconstruction preserve or install those proofs rather than forcing a +second canonicalization sweep. + The first milestone supports: - named and matrix-valued Clifford gates, which update only the tableau; diff --git a/src/pepsy/optimizers/tree/optimizer.py b/src/pepsy/optimizers/tree/optimizer.py index 1475cff..be57606 100644 --- a/src/pepsy/optimizers/tree/optimizer.py +++ b/src/pepsy/optimizers/tree/optimizer.py @@ -1159,7 +1159,8 @@ def _remount_product_state(self, state): root_tensor = target.node_tensor(target.plan.root) root_tensor.modify(data=root_tensor.data * factor) - target._with_center(self.plan.root).validate() + target._with_center(self.plan.root) + target._set_isometry_metadata_from_region({self.plan.root}).validate() self._state_backend_info(target) return target @@ -1385,6 +1386,23 @@ def is_canonical_form(self, center=None, *, tol=1e-9): """ return self.tn.is_canonical_form(center, tol=tol) + def isometry_direction(self, node): + """Neighbour toward which ``node`` has a proven local isometry.""" + return self.tn.isometry_direction(node) + + def isometry_map(self): + """Return the live network-owned node-isometry orientation map.""" + return self.tn.isometry_map() + + def can_skip_canonize(self, a, b, *, absorb="right"): + """Whether local metadata proves this dense edge QR is redundant.""" + return self.tn.can_skip_canonize(a, b, absorb=absorb) + + def validate_isometry_metadata(self, region=None): + """Validate live tensor ``left_inds`` against the canonical region.""" + self.tn.validate_isometry_metadata(region) + return self + @property def canonical_region(self): """Frozenset of node ids forming the canonicalised subtree (``None`` if unknown). @@ -1946,18 +1964,21 @@ def _apply_1q_impl(self, gate, q, *, renormalize=False): gate_np.conj().T @ gate_np, np.eye(d, dtype=gate_np.dtype), rtol=1e-10, atol=1e-12, ) + site_node = self.plan.node_of_qubit[q] if not unitary: - self._move_center(self.plan.node_of_qubit[q]) + self._move_center(site_node) region = self.tn.canonical_region + left_inds = self.tn.node_tensor(site_node).left_inds self.tn.gate_inds_(gate, [self._phys(q)], contract=True) if unitary: # A physical unitary preserves the isometric exterior, but the # state-owned gate mutator deliberately invalidates metadata for - # direct callers. Restore the known region only for this proven - # canonical-preserving operation. + # direct callers. Restore both the local proof and known region + # only for this proven canonical-preserving operation. + self.tn.node_tensor(site_node).modify(left_inds=left_inds) self.tn.canonical_region = region if not unitary: - self.center = self.plan.node_of_qubit[q] + self.center = site_node if renormalize: self.normalize() return self @@ -2269,9 +2290,17 @@ def _apply_2q_sibling_factors( max_bond=self.chi if max_bond is None else max_bond, cutoff=self.cutoff if cutoff is None else cutoff, ) - tla.modify(data=la_t.data, inds=la_t.inds) - tlb.modify(data=lb_t.data, inds=lb_t.inds) - tp.modify(data=p_t.data, inds=p_t.inds) + tla.modify( + data=la_t.data, + inds=la_t.inds, + left_inds=la_t.left_inds, + ) + tlb.modify( + data=lb_t.data, + inds=lb_t.inds, + left_inds=lb_t.left_inds, + ) + tp.modify(data=p_t.data, inds=p_t.inds, left_inds=None) self.center = parent return self @@ -2583,6 +2612,19 @@ def _compress_edge_with_diagnostics( full_spectrum=full_spectrum, max_bond=max_bond, cutoff=cutoff, ) + def _metadata_aware_reduction(self, u, v): + """Choose one-sided compression when ``v`` is proven isometric. + + Every caller compresses ``u -> v`` with ``absorb="right"``. If the + live ``left_inds`` on ``v`` prove that it is already isometric toward + ``u``, Quimb can SVD only ``u`` and reuse ``v`` directly. Missing or + native graded metadata conservatively falls back to the usual + two-sided QR reduction. + """ + if self.tn.can_skip_canonize(u, v, absorb="left"): + return "left" + return True + def _compress_path(self, path, *, max_bond=None, cutoff=None): """Canonically compress every bond along ``path`` down to ``chi``. @@ -2595,6 +2637,7 @@ def _compress_path(self, path, *, max_bond=None, cutoff=None): for v, u in zip(path[::-1], path[-2::-1]): self._compress_edge_with_diagnostics( v, u, max_bond=max_bond, cutoff=cutoff, + reduced=self._metadata_aware_reduction(v, u), ) self.center = path[0] @@ -2608,7 +2651,6 @@ def _compress_subtree(self, snodes, hub, *, max_bond=None, cutoff=None): """ snodes = frozenset(snodes) self._move_center(hub) - forward_reduced = True if self.tn.fermionic else "left" def descend(node, parent): children = sorted( @@ -2619,7 +2661,7 @@ def descend(node, parent): for child in children: self._compress_edge_with_diagnostics( node, child, max_bond=max_bond, cutoff=cutoff, - reduced=forward_reduced, + reduced=self._metadata_aware_reduction(node, child), ) descend(child, node) self.tn.canonize_edge_(child, node, absorb="right") @@ -3315,7 +3357,11 @@ def _apply_product_pauli_projector_impl( ) for nid in snodes: node_t = self.tn.tensor_map[self._tid(nid)] - node_t.modify(data=local[nid].data, inds=local[nid].inds) + node_t.modify( + data=local[nid].data, + inds=local[nid].inds, + left_inds=None if nid == hub else local[nid].left_inds, + ) self.center = hub def _apply_product_pauli_projector( diff --git a/src/pepsy/optimizers/tree/ttn.py b/src/pepsy/optimizers/tree/ttn.py index e08d484..ee964a4 100644 --- a/src/pepsy/optimizers/tree/ttn.py +++ b/src/pepsy/optimizers/tree/ttn.py @@ -683,6 +683,115 @@ def bond(self, a, b): ) return next(iter(shared)) + def isometry_direction(self, nid): + """Return the neighbour proven by ``left_inds`` to receive node ``nid``. + + A dense tree tensor is an isometry toward exactly one adjacent node + when its ``left_inds`` contain every leg except that shared tree bond. + ``None`` means no usable local proof is currently recorded. This is a + derived view of the live tensor metadata, not separately tracked state. + """ + if nid not in self._plan.children: + raise ValueError(f"{nid!r} is not a node of the tree.") + tensor = self.node_tensor(nid) + if tensor.left_inds is None: + return None + left_inds = set(tensor.left_inds) + right_inds = [ + index for index in tensor.inds if index not in left_inds + ] + if len(right_inds) != 1: + return None + toward_bond = right_inds[0] + for neighbour in self.neighbors(nid): + if toward_bond == self.bond(nid, neighbour): + return neighbour + return None + + def isometry_map(self): + """Return ``{node: toward_node_or_None}`` from live ``left_inds``.""" + return { + nid: self.isometry_direction(nid) + for nid in self._plan.nodes() + } + + def _set_isometry_metadata_from_region(self, region): + """Record orientations for a state already proven canonical. + + This changes metadata only and therefore must be called solely by + constructors or kernels that independently establish the stated + canonical region. + """ + region = self._validated_region(region) + for nid in self._plan.nodes(): + tensor = self.node_tensor(nid) + if nid in region: + tensor.modify(left_inds=None) + continue + toward = self._toward_region(nid, region) + bond = self.bond(nid, toward) + tensor.modify( + left_inds=tuple( + index for index in tensor.inds if index != bond + ) + ) + return self + + def can_skip_canonize(self, a, b, *, absorb="right"): + """Whether local metadata proves edge canonicalisation is redundant. + + With ``absorb="right"`` node ``a`` must already be isometric toward + ``b``; the ``"left"`` orientation is symmetric. Native fermionic trees + always return ``False`` because their graded QR path remains explicit. + """ + if absorb not in {"right", "left"}: + raise ValueError("absorb must be 'right' or 'left'.") + bond = self.bond(a, b) # validate the requested tree edge + if self.fermionic: + return False + if absorb == "right": + node = a + else: + node = b + tensor = self.node_tensor(node) + if tensor.left_inds is None: + return False + return set(tensor.left_inds) == set(tensor.inds) - {bond} + + def validate_isometry_metadata(self, region=None): + """Validate local ``left_inds`` against a canonical region. + + Every tensor outside ``region`` must point along its unique next edge + toward that region. Tensors inside the region need not be isometric. + When no explicit or tracked region exists, only malformed non-``None`` + metadata is rejected. Returns ``self`` when valid. + """ + if region is None: + region = self.canonical_region + else: + region = self._validated_region(region) + + for nid in self._plan.nodes(): + tensor = self.node_tensor(nid) + direction = self.isometry_direction(nid) + if tensor.left_inds is not None and direction is None: + raise ValueError( + f"tree node {nid} has left_inds that do not identify " + "exactly one adjacent isometry direction." + ) + if ( + not self.fermionic + and region is not None + and nid not in region + ): + expected = self._toward_region(nid, region) + if direction != expected: + raise ValueError( + f"tree node {nid} must be isometric toward node " + f"{expected}, but left_inds point toward {direction}." + ) + return self + def validate(self, *, check_canonical=False, tol=1e-9): """Validate the live network against its :class:`TreePlan`. @@ -823,10 +932,12 @@ def validate(self, *, check_canonical=False, tol=1e-9): ) if self._validated_region(region) != region: raise ValueError("canonical region is not a connected subtree.") - if check_canonical and not self.is_subtree_canonical_form( - region, tol=tol - ): - raise ValueError("tracked canonical region failed the isometry check.") + if check_canonical: + self.validate_isometry_metadata(region) + if not self.is_subtree_canonical_form(region, tol=tol): + raise ValueError( + "tracked canonical region failed the isometry check." + ) return self # -- plan delegators ------------------------------------------------------ @@ -1190,7 +1301,8 @@ def _recover_center_from_region(self, region, target, *, absorb="right"): a, b = node, neighbour else: a, b = neighbour, node - self.canonize_edge_(a, b, absorb=absorb) + if not self.can_skip_canonize(a, b, absorb=absorb): + self.canonize_edge_(a, b, absorb=absorb) remaining.remove(node) self._canonical_region = frozenset({target}) @@ -1646,13 +1758,14 @@ def from_plan(cls, plan, *, dtype=complex, phys_dim=2, site_tag_id="I{}", else: data = np.ones(shape, dtype=dtype) tensors.append(qtn.Tensor(data, inds=inds, tags=tags)) - return cls( + ttn = cls( tensors, plan=plan, site_tag_id=site_tag_id, site_ind_id=site_ind_id, node_tag_id=node_tag_id, - )._with_center(plan.root).validate() + )._with_center(plan.root) + return ttn._set_isometry_metadata_from_region({plan.root}).validate() @classmethod def from_symmray_plan( diff --git a/src/pepsy/optimizers/tree_stabilizer/optimizer.py b/src/pepsy/optimizers/tree_stabilizer/optimizer.py index 52baecd..f56c2d9 100644 --- a/src/pepsy/optimizers/tree_stabilizer/optimizer.py +++ b/src/pepsy/optimizers/tree_stabilizer/optimizer.py @@ -421,7 +421,13 @@ def _dense_to_tree_state(state, plan, *, max_bond=None, cutoff=0.0, dtype=comple tensors.append(qtn.Tensor(np.asarray(data, dtype=dtype), inds=inds, tags=tags)) - return TreeTensorNetwork(tensors, plan=plan)._with_center(plan.root).validate() + tree = TreeTensorNetwork(tensors, plan=plan)._with_center(plan.root) + # The hierarchical bases above are orthonormal by construction, so every + # non-root tensor is already an isometry toward the root. Record that + # proven local orientation without repeating the dense decomposition with + # a numerical canonicalization sweep. + tree._set_isometry_metadata_from_region({plan.root}) + return tree.validate() def _tree_bond_index(left, right): @@ -830,8 +836,9 @@ def __init__( ) self.to_backend = to_backend if to_backend is not None: - for tensor in self._tree.tn.tensor_map.values(): - tensor.modify(data=to_backend(tensor.data)) + # Backend-only conversion must not erase the QR/SVD isometry + # proofs stored in each tensor's ``left_inds``. + self._tree.tn.apply_to_arrays(to_backend) self._tree.tn.validate() self.max_operator_qubits = max_operator_qubits self.max_pauli_decomposition_qubits = max_operator_qubits @@ -1069,6 +1076,23 @@ def plan(self): def center(self): return self._tree.center + def isometry_direction(self, node): + """Return the coefficient-tree neighbour proven by ``left_inds``.""" + return self._tree.isometry_direction(node) + + def isometry_map(self): + """Return the live coefficient-tree isometry orientation map.""" + return self._tree.isometry_map() + + def can_skip_canonize(self, a, b, *, absorb="right"): + """Whether coefficient metadata proves an edge QR is redundant.""" + return self._tree.can_skip_canonize(a, b, absorb=absorb) + + def validate_isometry_metadata(self, region=None): + """Validate coefficient ``left_inds`` against its canonical region.""" + self._tree.validate_isometry_metadata(region) + return self + @property def tree_optimizer(self): """Return the coefficient-side ``TreeOptimizer``.""" @@ -3478,10 +3502,11 @@ def cap(self, where, vec, *, absorb="left") -> "TreeStabOptimizer": dtype=old_tree.dtype, ) if self.to_backend is not None: - for tensor in coefficient.tensor_map.values(): - tensor.modify(data=self.to_backend(tensor.data)) + coefficient.apply_to_arrays(self.to_backend) elif self._tree.backend_info()["backend"] != "numpy": - self._tree._coerce_tensor_network_backend(coefficient, warn=False) + coefficient.apply_to_arrays( + self._tree._backend_converter(self._tree._state_like()) + ) new_tree = TreeOptimizer( None, n=reduced_n, diff --git a/src/pepsy/tensors/constructors.py b/src/pepsy/tensors/constructors.py index 0aed728..64d6b14 100644 --- a/src/pepsy/tensors/constructors.py +++ b/src/pepsy/tensors/constructors.py @@ -828,6 +828,11 @@ def ps_to_ttn( seed_rand(seed) ttn.expand_bond_dimension_(chi, rand_strength=rand_strength) ttn.canonize_around_node_(plan.root) + else: + # Replacing each product vector clears Quimb's local ``left_inds``. + # Bond-one normalized product tensors remain trivially canonical, so + # restore the network-owned orientation metadata without new QR work. + ttn._set_isometry_metadata_from_region({plan.root}) return ttn.validate() @@ -990,6 +995,8 @@ def hrs_to_ttn( if chi > 1: ttn.expand_bond_dimension_(chi, rand_strength=rand_strength) ttn.canonize_around_node_(plan.root) + else: + ttn._set_isometry_metadata_from_region({plan.root}) return ttn.validate() diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index 5e578de..f0e844e 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -167,6 +167,112 @@ def check_then_compress(path, **kwargs): assert opt.tn.validate(check_canonical=True) is opt.tn +@pytest.mark.parametrize("mode", ("direct", "mpo", "submpo")) +def test_two_site_modes_reuse_path_isometries_for_compression( + mode, monkeypatch, +): + """All two-site routes skip the QR already proven by ``left_inds``.""" + rng = np.random.default_rng(920) + n = 8 + where = (0, 7) + gate = _rand_unitary(2, rng) + opt = TreeOptimizer( + None, n=n, chi=16, cutoff=1e-12, mode=mode, run=False, + ) + + reductions = [] + compress_edge = opt._compress_edge_with_diagnostics + + def traced_compress_edge( + u, v, *, max_bond=None, cutoff=None, reduced=True, + ): + assert opt.tn.can_skip_canonize(u, v, absorb="left") + reductions.append(reduced) + return compress_edge( + u, + v, + max_bond=max_bond, + cutoff=cutoff, + reduced=reduced, + ) + + monkeypatch.setattr( + opt, "_compress_edge_with_diagnostics", traced_compress_edge, + ) + if mode == "submpo": + submpo = qtn.MatrixProductOperator.from_dense( + gate.reshape((2,) * 4), + dims=(2, 2), + sites=where, + L=n, + max_bond=None, + cutoff=0.0, + ) + opt.apply_submpo(submpo, where) + else: + opt.apply_2q(gate, *where) + + expected = np.zeros(2**n, dtype=complex) + expected[0] = 1.0 + expected = _sv_apply_kq(expected, gate, where, n) + assert reductions + assert all(reduced == "left" for reduced in reductions) + assert _fidelity(expected, opt.to_dense()) > 1 - 1e-10 + assert opt.tn.validate(check_canonical=True) is opt.tn + + +def test_dense_path_one_sided_compression_matches_full_reduction(monkeypatch): + """Reusing routed Q tensors is exact even when the path truncates.""" + rng = np.random.default_rng(921) + n = 8 + plan = TreePlan.from_order(range(n), structure="balanced") + seed = TreeTensorNetwork.rand(plan, D=2, seed=921) + optimized = TreeOptimizer( + None, + tree=plan, + state=seed.copy(), + chi=2, + cutoff=1e-12, + mode="direct", + run=False, + ) + reference = TreeOptimizer( + None, + tree=plan, + state=seed.copy(), + chi=2, + cutoff=1e-12, + mode="direct", + run=False, + ) + monkeypatch.setattr( + reference, "_metadata_aware_reduction", lambda _u, _v: True, + ) + + for where in ((0, 7), (1, 6), (2, 5), (0, 4)): + gate = _rand_unitary(2, rng) + optimized.apply_2q(gate, *where) + reference.apply_2q(gate, *where) + + assert optimized.max_bond() <= 2 + assert reference.max_bond() <= 2 + assert _fidelity(optimized.to_dense(), reference.to_dense()) > 1 - 1e-10 + assert optimized.tn.validate(check_canonical=True) is optimized.tn + assert reference.tn.validate(check_canonical=True) is reference.tn + + +def test_compression_reduction_falls_back_without_local_isometry_proof(): + """One-sided compression is selected only from live ``left_inds``.""" + plan = TreePlan.from_order(range(4), structure="balanced") + opt = TreeOptimizer(None, tree=plan, run=False) + child = plan.leaf_of_qubit[0] + parent = plan.parent[child] + + assert opt._metadata_aware_reduction(parent, child) == "left" + opt.tn.node_tensor(child).modify(left_inds=None) + assert opt._metadata_aware_reduction(parent, child) is True + + def test_tree_mpo_mode_keeps_small_operator_schmidt_components(): """MPO lowering must not apply Quimb's default gate-SVD cutoff.""" x = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex) @@ -252,12 +358,18 @@ def test_dense_subtree_hub_recovery_reuses_routed_q_metadata(monkeypatch): qr_calls = [] tensor_split = qtc.tensor_split + canonize_calls = [] + canonize_between = opt.tn.canonize_between def traced_tensor_split(*args, **kwargs): if kwargs.get("method") == "qr": qr_calls.append(args[0]) return tensor_split(*args, **kwargs) + def traced_canonize_between(*args, **kwargs): + canonize_calls.append(args) + return canonize_between(*args, **kwargs) + recoveries = [] move_center = opt._move_center @@ -273,17 +385,23 @@ def traced_move_center(target): bond = opt.tn.bond(nid, toward_hub) assert tensor.left_inds is not None assert set(tensor.left_inds) == set(tensor.inds) - {bond} - before = len(qr_calls) + before = (len(qr_calls), len(canonize_calls)) result = move_center(target) - recoveries.append(len(qr_calls) - before) + recoveries.append( + ( + len(qr_calls) - before[0], + len(canonize_calls) - before[1], + ) + ) return result return move_center(target) monkeypatch.setattr(qtc, "tensor_split", traced_tensor_split) + monkeypatch.setattr(opt.tn, "canonize_between", traced_canonize_between) monkeypatch.setattr(opt, "_move_center", traced_move_center) opt.apply_subtree_operator(gate, where) - assert recoveries == [0] + assert recoveries == [(0, 0)] assert _fidelity(expected, opt.to_dense()) > 1 - 1e-10 assert opt.tn.validate(check_canonical=True) is opt.tn @@ -2306,8 +2424,8 @@ def test_tree_torch_state_stays_native_across_public_operations(): to_backend = pepsy.backend_torch(device="cpu", dtype=torch.complex128) plan = TreePlan.from_order(range(3), structure="balanced") state = TreeTensorNetwork.from_plan(plan) - for tensor in state.tensor_map.values(): - tensor.modify(data=to_backend(tensor.data)) + state.apply_to_arrays(to_backend) + assert state.validate_isometry_metadata() is state h = to_backend( np.array([[1.0, 1.0], [1.0, -1.0]], dtype=complex) / np.sqrt(2.0) ) @@ -2339,8 +2457,7 @@ def test_tree_warns_once_when_a_gate_does_not_match_the_state_backend(): to_backend = pepsy.backend_torch(device="cpu", dtype=torch.complex128) plan = TreePlan.from_order(range(2), structure="balanced") state = TreeTensorNetwork.from_plan(plan) - for tensor in state.tensor_map.values(): - tensor.modify(data=to_backend(tensor.data)) + state.apply_to_arrays(to_backend) opt = TreeOptimizer(None, state=state, run=False) with pytest.warns(UserWarning, match="backend-compatible gate"): @@ -2355,8 +2472,7 @@ def test_tree_gate_stream_backend_preparation_is_stream_level(): to_backend = pepsy.backend_torch(device="cpu", dtype=torch.complex128) plan = TreePlan.from_order(range(2), structure="balanced") state = TreeTensorNetwork.from_plan(plan) - for tensor in state.tensor_map.values(): - tensor.modify(data=to_backend(tensor.data)) + state.apply_to_arrays(to_backend) gates = [ np.eye(2, dtype=complex), np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex), @@ -2381,8 +2497,7 @@ def test_tree_submpo_stream_backend_preparation_preserves_input(): to_backend = pepsy.backend_torch(device="cpu", dtype=torch.complex128) plan = TreePlan.from_order(range(2), structure="balanced") state = TreeTensorNetwork.from_plan(plan) - for tensor in state.tensor_map.values(): - tensor.modify(data=to_backend(tensor.data)) + state.apply_to_arrays(to_backend) submpo = _two_branch_flip_submpo(L=2, sites=(0, 1), targets=(0, 1)) opt = TreeOptimizer(None, state=state, tree=plan, run=False) @@ -2689,6 +2804,51 @@ def _entangled_ttn(seed=0, n=6, D=3, structure="balanced"): return TreeTensorNetwork.rand(plan, D=D, seed=seed) +def test_isometry_metadata_api_has_one_network_owned_orientation_map(): + """Product construction and optimizer delegates expose one live map.""" + plan = TreePlan.from_order(range(6), structure="balanced") + ttn = TreeTensorNetwork.from_plan(plan) + directions = ttn.isometry_map() + + assert directions[plan.root] is None + for nid in plan.nodes(): + if nid == plan.root: + continue + assert directions[nid] == plan.parent[nid] + assert ttn.can_skip_canonize(nid, plan.parent[nid]) + assert ttn.can_skip_canonize( + plan.parent[nid], nid, absorb="left", + ) + assert ttn.validate_isometry_metadata() is ttn + assert ttn.validate(check_canonical=True) is ttn + + opt = TreeOptimizer(None, tree=plan, state=ttn, run=False) + assert opt.isometry_map() == directions + leaf = plan.leaf_of_qubit[0] + assert opt.isometry_direction(leaf) == plan.parent[leaf] + assert opt.can_skip_canonize(leaf, plan.parent[leaf]) + assert opt.validate_isometry_metadata() is opt + + +def test_isometry_metadata_validation_detects_cleared_local_proof(): + """A live canonical-region claim cannot outlast cleared ``left_inds``.""" + ttn = _entangled_ttn(seed=43) + leaf = ttn.leaf_of_qubit(0) + tensor = ttn.node_tensor(leaf) + assert ttn.isometry_direction(leaf) == ttn.parent(leaf) + + # Quimb correctly clears ``left_inds`` whenever tensor data changes. + tensor.modify(data=np.array(tensor.data)) + assert ttn.isometry_direction(leaf) is None + with pytest.raises(ValueError, match="must be isometric"): + ttn.validate_isometry_metadata() + with pytest.raises(ValueError, match="must be isometric"): + ttn.validate(check_canonical=True) + + ttn.invalidate_canonical_form() + assert ttn.validate_isometry_metadata() is ttn + + def test_shift_center_lossless_and_recanonical(): """Shifting the centre preserves the state exactly and re-canonicalises.""" ttn = _entangled_ttn(seed=1) @@ -3285,6 +3445,10 @@ def dense_vector(opt): _fidelity(dense_vector(candidate), dense_vector(reference)) > 1 - 1e-10 ) + assert candidate.validate_isometry_metadata() is candidate + for nid, toward in candidate.isometry_map().items(): + if toward is not None: + assert not candidate.can_skip_canonize(nid, toward) assert candidate.tn.validate(check_canonical=True) is candidate.tn diff --git a/tests/test_optimize_tree_stabilizer.py b/tests/test_optimize_tree_stabilizer.py index bc21b9c..a05e26b 100644 --- a/tests/test_optimize_tree_stabilizer.py +++ b/tests/test_optimize_tree_stabilizer.py @@ -85,6 +85,97 @@ def test_tree_stab_is_public_and_cliffords_are_tableau_only(): _assert_same_state(opt.to_statevector(), expected) +def test_tree_stab_isometry_api_and_backend_conversion_preserve_proofs(): + """TreeStab delegates one live map and backend conversion keeps it valid.""" + def converter(array): + return np.array(array, copy=True) + + opt = pepsy.TreeStabOptimizer(6, to_backend=converter) + directions = opt.isometry_map() + + assert directions == opt.tree_optimizer.isometry_map() + assert directions[opt.plan.root] is None + for node in opt.plan.nodes(): + if node == opt.plan.root: + continue + parent = opt.plan.parent[node] + assert opt.isometry_direction(node) == parent + assert opt.can_skip_canonize(node, parent) + assert opt.validate_isometry_metadata() is opt + assert opt.p.validate(check_canonical=True) is opt.p + + # Dense cap reconstruction independently proves a new root-canonical tree + # and then crosses the same backend-only conversion boundary. + opt.cap(1, [1.0, 0.0]) + assert opt.validate_isometry_metadata() is opt + assert opt.p.validate(check_canonical=True) is opt.p + + +@pytest.mark.parametrize("route", ("direct", "mpo", "submpo")) +def test_tree_stab_routes_reuse_proven_path_isometries( + route, monkeypatch, +): + """Direct, MPO, and sub-MPO coefficient paths avoid redundant QRs.""" + from pepsy.optimizers.tree import TreePlan + + n = 8 + where = (0, 7) + plan = TreePlan.from_order(range(n), structure="balanced") + opt = pepsy.TreeStabOptimizer(n, tree=plan, mode=route) + reductions = [] + compress_edge = opt.tree_optimizer._compress_edge_with_diagnostics + + def traced_compress_edge( + u, v, *, max_bond=None, cutoff=None, reduced=True, + ): + reductions.append(( + reduced, + opt.can_skip_canonize(u, v, absorb="left"), + )) + return compress_edge( + u, + v, + max_bond=max_bond, + cutoff=cutoff, + reduced=reduced, + ) + + monkeypatch.setattr( + opt.tree_optimizer, + "_compress_edge_with_diagnostics", + traced_compress_edge, + ) + if route == "submpo": + operator = np.kron(X, X).reshape((2,) * 4) + submpo = qtn.MatrixProductOperator.from_dense( + operator, + dims=(2, 2), + sites=where, + L=n, + max_bond=None, + cutoff=0.0, + ) + opt.apply([("submpo", submpo, where)]) + expected = np.zeros(2**n, dtype=complex) + expected[(1 << (n - 1)) | 1] = 1.0 + _assert_same_state(opt.to_statevector(), expected) + else: + # The physical Cliffords spread Z_7 to Z_0 Z_7. Basis-updating + # measurement then runs a long-range coefficient CNOT localizer through + # the selected direct/MPO path. + opt.apply([("h", 0), ("cnot", 0, 7)]) + opt.measure_pauli("Z", 7, outcome=+1, absorb_basis=True) + assert opt.expectation("Z", 7) == pytest.approx(1.0) + + assert reductions + assert all( + reduced == "left" and proven + for reduced, proven in reductions + ) + assert opt.validate_isometry_metadata() is opt + assert opt.p.validate(check_canonical=True) is opt.p + + def test_tree_stab_chi_none_is_uncapped_and_sampling_is_conditional(): opt = pepsy.TreeStabOptimizer( 2, @@ -243,6 +334,8 @@ def test_tree_stab_norm_diagnostics_and_sampling_copy_contract(): def test_tree_stab_torch_backend_matches_numpy(): + from pepsy.optimizers.tree import TreePlan, TreeTensorNetwork + torch = pytest.importorskip("torch") backend = pepsy.backend_torch(dtype=torch.complex128, device="cpu") stream = [ @@ -257,6 +350,22 @@ def test_tree_stab_torch_backend_matches_numpy(): assert gpu.backend_info()["backend"] == "torch" assert "torch" in type(gpu.p[0].data).__module__ _assert_same_state(gpu.to_statevector(), cpu.to_statevector()) + assert gpu.validate_isometry_metadata() is gpu + assert gpu.p.validate(check_canonical=True) is gpu.p + + # A caller-supplied native TTN has no retained ``to_backend`` callback, so + # cap must derive the converter from the live tree without clearing proofs. + native_state = TreeTensorNetwork.from_plan( + TreePlan.from_order(range(3), structure="balanced") + ) + native_state.apply_to_arrays(backend) + inherited = pepsy.TreeStabOptimizer( + native_state, max_dense_cap_qubits=4 + ) + inherited.cap(1, [1.0, 0.0]) + assert inherited.backend_info()["backend"] == "torch" + assert inherited.validate_isometry_metadata() is inherited + assert inherited.p.validate(check_canonical=True) is inherited.p def test_tree_stab_cap_matches_mps_and_rebuilds_identity_frame(): @@ -276,6 +385,8 @@ def test_tree_stab_cap_matches_mps_and_rebuilds_identity_frame(): mps.probability("00"), abs=1e-7 ) assert tree.p.max_bond() > 1 + assert tree.validate_isometry_metadata() is tree + assert tree.p.validate(check_canonical=True) is tree.p def test_tree_stab_cap_stream_remaps_later_compact_labels(): From 9f5c0da2405a7da25c826b9b25b06ccdb61b3ed0 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Tue, 28 Jul 2026 17:51:26 -0600 Subject: [PATCH 15/70] Restore Fermion model parameters for qMERA --- src/pepsy/tensors/symmetric.py | 105 ++++++++++++++++++--------------- 1 file changed, 56 insertions(+), 49 deletions(-) diff --git a/src/pepsy/tensors/symmetric.py b/src/pepsy/tensors/symmetric.py index 5744614..2132423 100644 --- a/src/pepsy/tensors/symmetric.py +++ b/src/pepsy/tensors/symmetric.py @@ -7314,13 +7314,12 @@ def fermion_hopping_param_gen( class Fermion: """Native spinless or spinful fermion observables, gates, and streams. - The helper owns only the local fermionic space, symmetry convention, and - optional backend conversion. Hamiltonian couplings are deliberately not - stored here: construct them as explicit native terms, then validate and - bundle them with :meth:`hamiltonian`. This prevents a native Hamiltonian, - a gate stream, and a VMC adapter from silently using different couplings. - It is intended for direct Symmray-backed fermionic MPS or PEPS workflows; - it does not introduce a qubit or Jordan-Wigner circuit representation. + The helper owns the local fermionic space, symmetry convention, optional + model parameters, and backend conversion. Explicit native terms can still + be passed to :meth:`hamiltonian`; model parameters are used for compact + lattice and qMERA workflows. It is intended for direct Symmray-backed + fermionic MPS or PEPS workflows and does not introduce a qubit or + Jordan-Wigner circuit representation. ``strang_gate_stream`` uses a deterministic edge colouring and a forward/reverse half-step sequence. Consequently its hopping layers are @@ -7330,9 +7329,13 @@ class Fermion: """ symmetry: str | None = None + t: object = 1.0 + U: object = 8.0 dtype: object = "complex128" to_backend: object = None spinful: bool = True + V: object = 0.0 + mu: object = 0.0 _dense_ops: dict = field(default_factory=dict, init=False, repr=False) _observable_cache: dict = field(default_factory=dict, init=False, repr=False) _gate_cache: dict = field(default_factory=dict, init=False, repr=False) @@ -8263,12 +8266,13 @@ def hopping_operator(self, *, spin=None, peierls_angle=0.0): peierls_angle=peierls_angle, ) - def hopping_term(self, edge, *, spin=None, t, peierls_angle=0.0): + def hopping_term(self, edge, *, spin=None, t=None, peierls_angle=0.0): """Return ``-t`` times the hopping operator on ``edge``.""" try: left, right = tuple(edge) except (TypeError, ValueError) as exc: raise ValueError("edge must contain exactly two site labels.") from exc + t = self.t if t is None else t t = _edge_parameter(t, left, right) return self._hopping_operator_on_sites( left, @@ -8291,8 +8295,9 @@ def interaction_operator(self): sites=(0,), ) - def interaction_term(self, site, *, U): + def interaction_term(self, site, *, U=None): """Return ``U n_up n_down`` on one physical site.""" + U = self.U if U is None else U U = _node_parameter(U, site) return self.operator_term( [(U, ((site, "double"),))], @@ -8310,8 +8315,9 @@ def chemical_potential_operator(self): terms = [(1.0, ((0, "number"),))] return self.operator_term(terms, sites=(0,)) - def chemical_potential_term(self, site, *, mu): + def chemical_potential_term(self, site, *, mu=None): """Return ``-mu n`` on one physical site.""" + mu = self.mu if mu is None else mu if self.spinful: mu = _node_parameter(mu, site) mu_up, mu_down = _as_spin_pair(mu, name="mu") @@ -8323,12 +8329,12 @@ def chemical_potential_term(self, site, *, mu): terms = [(-_node_parameter(mu, site), ((site, "number"),))] return self.operator_term(terms, sites=(site,)) - def onsite_term(self, site, *, U=None, mu=0.0): + def onsite_term(self, site, *, U=None, mu=None): """Return ``U n_up n_down - mu n`` on one site.""" + U = self.U if U is None else U + mu = self.mu if mu is None else mu terms = [] if self.spinful: - if U is None: - raise TypeError("onsite_term requires explicit U=... for spinful fermions.") terms.append((_node_parameter(U, site), ((site, "double"),))) mu = _node_parameter(mu, site) mu_up, mu_down = _as_spin_pair(mu, name="mu") @@ -8355,12 +8361,13 @@ def density_operator(self): ] return self.operator_term(terms, sites=(0, 1)) - def density_term(self, edge, *, V): + def density_term(self, edge, *, V=None): """Return ``V n_i n_j`` on a physical edge.""" try: left, right = tuple(edge) except (TypeError, ValueError) as exc: raise ValueError("edge must contain exactly two site labels.") from exc + V = self.V if V is None else V V = _edge_parameter(V, left, right) if self.spinful: names = ("number_up", "number_down") @@ -8508,7 +8515,7 @@ def heisenberg_gate(self, theta, *, edge=None, imaginary=False): szz_gate = spin_z_correlator_gate xy_gate = xy_exchange_gate - def interaction_gate(self, dt, *, site=None, U, imaginary=False): + def interaction_gate(self, dt, *, site=None, U=None, imaginary=False): """Return the exact onsite interaction gate. With a site-dependent ``U`` mapping or callable, pass ``site`` so the @@ -8520,6 +8527,7 @@ def interaction_gate(self, dt, *, site=None, U, imaginary=False): "Spinless fermions have no onsite doublon interaction; use " "density_gate(...) for the nearest-neighbor V interaction." ) + U = self.U if U is None else U U = U if site is None else _node_parameter(U, site) theta = dt * U @@ -8533,20 +8541,20 @@ def build(): return self._cached_gate(("interaction", dt, site, U, imaginary), build) - def onsite_gate(self, dt, *, site=None, U=None, mu=0.0, imaginary=False): + def onsite_gate(self, dt, *, site=None, U=None, mu=None, imaginary=False): """Return the complete one-site Hubbard gate. The generated gate represents ``U n_up n_down - mu n`` for spinful fermions and ``-mu n`` for spinless fermions. ``U`` and ``mu`` may be site-dependent mappings or callables when ``site`` is supplied. """ + U = self.U if U is None else U + mu = self.mu if mu is None else mu if site is not None: U = _node_parameter(U, site) mu = _node_parameter(mu, site) if self.spinful: - if U is None: - raise TypeError("onsite_gate requires explicit U=... for spinful fermions.") mu_up, mu_down = _as_spin_pair(mu, name="mu") U_site = U diagonal = ( @@ -8571,8 +8579,9 @@ def build(): return self._cached_gate(("onsite", dt, site, U, mu, imaginary), build) - def hopping_gate(self, dt, *, t, peierls_angle=0.0, imaginary=False): + def hopping_gate(self, dt, *, t=None, peierls_angle=0.0, imaginary=False): """Return a two-site native fermionic hopping gate with Peierls phase.""" + t = self.t if t is None else t def build(): if not self.spinful: gate = _spinless_hopping_gate( @@ -8594,12 +8603,13 @@ def build(): return self._cached_gate(("hopping", dt, t, peierls_angle, imaginary), build) - def density_gate(self, dt, *, V, imaginary=False): + def density_gate(self, dt, *, V=None, imaginary=False): """Return the nearest-neighbor density interaction gate. For spinless fermions this is ``V n_i n_j``. For spinful fermions it is ``V (n_up + n_down)_i (n_up + n_down)_j``. """ + V = self.V if V is None else V theta = dt * V def build(): @@ -8615,8 +8625,9 @@ def build(): return self._cached_gate(("density", dt, V, imaginary), build) - def chemical_potential_gate(self, dt, *, mu, site=None, imaginary=False): + def chemical_potential_gate(self, dt, *, mu=None, site=None, imaginary=False): """Return the chemical-potential part of an onsite gate.""" + mu = self.mu if mu is None else mu mu = mu if site is None else _node_parameter(mu, site) if self.spinful: mu_up, mu_down = _as_spin_pair(mu, name="mu") @@ -8861,16 +8872,16 @@ def gate_stream( imaginary=False, t=None, U=None, - V=0.0, - mu=0.0, + V=None, + mu=None, ): - """Return a canonical fermion gate stream with explicit couplings.""" + """Return a canonical first- or second-order fermion gate stream.""" if order not in {1, 2}: raise ValueError("order must be 1 or 2.") - if t is None: - raise TypeError("gate_stream requires explicit t=... .") - if self.spinful and U is None: - raise TypeError("gate_stream requires explicit U=... for spinful fermions.") + t = self.t if t is None else t + U = self.U if U is None else U + V = self.V if V is None else V + mu = self.mu if mu is None else mu edges = _as_edges(edges) sites = _sites_from_edges(edges, sites) @@ -8943,16 +8954,14 @@ def strang_gate_stream( imaginary=False, t=None, U=None, - V=0.0, - mu=0.0, + V=None, + mu=None, ): - """Return an edge-coloured second-order stream with explicit couplings.""" - if t is None: - raise TypeError("strang_gate_stream requires explicit t=... .") - if self.spinful and U is None: - raise TypeError( - "strang_gate_stream requires explicit U=... for spinful fermions." - ) + """Return an edge-coloured second-order native fermionic gate stream.""" + t = self.t if t is None else t + U = self.U if U is None else U + V = self.V if V is None else V + mu = self.mu if mu is None else mu edges = _as_edges(edges) sites = _sites_from_edges(edges, sites) half_dt = dt / 2 @@ -9100,23 +9109,23 @@ def hamiltonian( *, t=None, U=None, - V=0.0, - mu=0.0, + V=None, + mu=None, flat=False, to_backend=None, ): - """Validate explicit terms or build a model only from explicit couplings. + """Validate explicit terms or build the configured lattice model. The canonical form is a mapping from one-site or two-site locations to native fermionic Symmray arrays. It is checked for symmetry, physical sectors, support rank, and backend consistency before being bundled in a :class:`SymHamiltonian`. Passing lattice edges remains a compact - convenience, but requires its couplings explicitly; no coupling is - stored on :class:`Fermion`. + convenience, using this helper's model parameters unless explicit + coupling overrides are supplied. """ to_backend = self.to_backend if to_backend is None else to_backend if isinstance(terms_or_edges, Mapping): - if any(value is not None for value in (t, U)) or V != 0 or mu != 0: + if any(value is not None for value in (t, U, V, mu)): raise TypeError( "When passing explicit terms, put every coupling in the " "native arrays rather than passing t/U/V/mu again." @@ -9130,12 +9139,10 @@ def hamiltonian( parameters={}, ) - if t is None: - raise TypeError("hamiltonian(edges, ...) requires explicit t=... .") - if self.spinful and U is None: - raise TypeError( - "hamiltonian(edges, ...) requires explicit U=... for spinful fermions." - ) + t = self.t if t is None else t + U = self.U if U is None else U + V = self.V if V is None else V + mu = self.mu if mu is None else mu params = {"t": t, "V": V, "mu": mu} if self.spinful: params["U"] = U From a644a404eae742e7b47f51f6d445861f50457e21 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Tue, 28 Jul 2026 18:07:53 -0600 Subject: [PATCH 16/70] Improve fermionic VMC workflows --- docs/api/optimizers/mera.md | 17 +- docs/api/tensors/core.md | 5 +- docs/api/tensors/symmetric.md | 60 +- docs/api/vmc.md | 223 +- docs/development/plans/project.md | 17 +- examples/qmera_fermion_hubbard_2d.py | 10 +- examples/qmera_fermion_hubbard_4x4_pbc.py | 9 +- src/pepsy/backends/config.py | 4 +- src/pepsy/backends/linalg_jax.py | 109 +- src/pepsy/optimizers/mera/builders.py | 3 +- src/pepsy/optimizers/mera/fermions.py | 16 +- src/pepsy/tensors/constructors.py | 26 +- src/pepsy/tensors/symmetric.py | 109 +- src/pepsy/vmc/__init__.py | 7 + src/pepsy/vmc/api.py | 165 +- src/pepsy/vmc/netket.py | 2280 ++++++++++++++++++--- src/pepsy/vmc/torch/__init__.py | 10 + src/pepsy/vmc/torch/_core.py | 2 + src/pepsy/vmc/torch/amplitude.py | 219 +- src/pepsy/vmc/torch/benchmark.py | 254 +++ src/pepsy/vmc/torch/distributed.py | 199 ++ src/pepsy/vmc/torch/driver.py | 232 ++- src/pepsy/vmc/torch/fermion.py | 162 +- src/pepsy/vmc/torch/local_energy.py | 111 +- src/pepsy/vmc/torch/results.py | 71 +- src/pepsy/vmc/torch/sampler.py | 12 +- tests/test_backends.py | 38 + tests/test_fermion_gate_cache.py | 6 +- tests/test_netket_flat_z2.py | 437 ++++ tests/test_optimize_mera.py | 35 +- tests/test_optimize_mps.py | 4 +- tests/test_optimize_tree.py | 7 +- tests/test_package_layout.py | 2 + tests/test_symmetric_tensors.py | 148 +- tests/test_tree_sampler.py | 4 +- tests/test_vmc_api.py | 386 +++- tests/test_vmc_distributed.py | 198 ++ tests/test_vmc_importance.py | 27 +- tests/test_vmc_local_energy.py | 209 ++ 39 files changed, 5095 insertions(+), 738 deletions(-) create mode 100644 src/pepsy/vmc/torch/benchmark.py create mode 100644 src/pepsy/vmc/torch/distributed.py create mode 100644 tests/test_netket_flat_z2.py create mode 100644 tests/test_vmc_distributed.py create mode 100644 tests/test_vmc_local_energy.py diff --git a/docs/api/optimizers/mera.md b/docs/api/optimizers/mera.md index 0e6cffe..24ff6d0 100644 --- a/docs/api/optimizers/mera.md +++ b/docs/api/optimizers/mera.md @@ -132,16 +132,18 @@ from pepsy.optimizers.mera import QMeraGeometry fermion = pepsy.Fermion( spinful=True, symmetry="U1U1", - t=1.0, - U=8.0, ) edges = ((0, 1), (1, 2)) -site_terms = fermion.local_terms(edges) -gate_stream = fermion.gate_stream(edges, dt=0.01, sites=range(3)) +site_terms = fermion.local_terms(edges, t=1.0, U=8.0) +gate_stream = fermion.gate_stream( + edges, dt=0.01, sites=range(3), t=1.0, U=8.0 +) geometry = QMeraGeometry(shape=3, site_modes=("up", "down")) -qmera_terms = fermion.local_terms(geometry, layout="qmera") +qmera_terms = fermion.local_terms( + geometry, layout="qmera", t=1.0, U=8.0 +) ``` For the normal spinful Hubbard workflow, let the builder own the mode @@ -163,15 +165,16 @@ builder = QMeraBuilder( product_state_factory=backend.product_state, ) geometry = builder.geometry # inferred from the model -terms = builder.fermion_terms() # inferred from the model +terms = builder.fermion_terms(t=1.0, U=8.0) # explicit physical couplings optimizer = builder.fermion_parametric_optimizer( energy_per_site=False, + term_params={"t": 1.0, "U": 8.0}, ) ``` Passing `site_modes` or `mode_order` remains useful when testing a custom register convention. A site-layout object such as -`fermion.hamiltonian(edges)` is still a valid native MPS/PEPS Hamiltonian, but +`fermion.hamiltonian(edges, t=..., U=...)` is still a valid native MPS/PEPS Hamiltonian, but it does not by itself specify qMERA's explicit mode registers or RG schedule; use `fermion.local_terms(geometry, layout="qmera")` or the builder shortcut above for that conversion. diff --git a/docs/api/tensors/core.md b/docs/api/tensors/core.md index 85cd584..74b79c3 100644 --- a/docs/api/tensors/core.md +++ b/docs/api/tensors/core.md @@ -8,8 +8,9 @@ The SVD/QR registration helpers are also available directly from `pepsy`, e.g. `import pepsy as py; py.reg_rel_svd_torch()`. Torch exposes `reg_rel_svd_torch()`, `reg_real_svd_torch()`, `reg_complex_svd_torch()`, `reg_real_qr_torch()`, and `reg_complex_qr_torch()`. JAX exposes SVD aliases -`reg_rel_svd_jax()`, `reg_real_svd_jax()`, and `reg_complex_svd_jax()` for the -same custom-VJP SVD registration. +`reg_rel_svd_jax()`, `reg_real_svd_jax()`, and `reg_complex_svd_jax()` for a +thin-SVD custom VJP that preserves JAX's native derivative while safely +restoring cotangents from Quimb fixed-rank truncation. > API details are maintained as handwritten Markdown in this page. diff --git a/docs/api/tensors/symmetric.md b/docs/api/tensors/symmetric.md index 03d70a9..743f7c6 100644 --- a/docs/api/tensors/symmetric.md +++ b/docs/api/tensors/symmetric.md @@ -66,7 +66,7 @@ model and returns the underlying native fermionic MPS directly; callers do not need to instantiate ``SymMPS``: ```python -fh = py.Fermion(spinful=True, symmetry="U1U1", t=1.0, U=8.0) +fh = py.Fermion(spinful=True, symmetry="U1U1") psi = py.ps_to_mps(8, fermion=fh, seed=7) # Override the product-state charge pattern when needed. @@ -145,45 +145,39 @@ psi = py.hrs_to_peps( ) ``` -The direct method uses Symmray's ``PEPS_fermionic_rand`` constructor and -normalizes the returned PEPS. A unitary PEPS-growth method is not implemented -yet. Both constructors return the underlying PEPS with native fermionic -Symmray tensors; use ``SymPEPS`` only when wrapper methods or stored -Hamiltonian metadata are needed. +The direct method uses Symmray's ``PEPS_fermionic_rand`` constructor. It skips +global normalization by default, avoiding an expensive CPU boundary-MPS +contraction; this is safe for NetKet VMC because a global wavefunction scalar +cancels from sampling and local-energy ratios. Pass ``normalize=True`` only +when a globally normalized PEPS is explicitly required. A unitary PEPS-growth +method is not implemented yet. Both constructors return the underlying PEPS +with native fermionic Symmray tensors; use ``SymPEPS`` only when wrapper +methods or stored Hamiltonian metadata are needed. ## Unified native fermion helper ``Fermion`` is the model-facing helper for both one-mode spinless fermions -and four-state spinful Hubbard sites. ``spinful`` selects the local space; +and four-state spinful Hubbard sites. It owns the local space, symmetry, and +optional backend conversion; physical couplings are passed explicitly when +constructing a term, gate, or stream. ``spinful`` selects the local space; ``symmetry`` selects the conserved charge group. Spinless helpers support ``U1`` and ``Z2``; spinful helpers support ``U1``, ``Z2``, ``U1U1``, and ``Z2Z2``. ```python -spinless = py.Fermion( - spinful=False, - symmetry="U1", - t=1.0, - V=0.5, - mu=0.0, -) +spinless = py.Fermion(spinful=False, symmetry="U1") -spinful = py.Fermion( - spinful=True, - symmetry="U1U1", - t=1.0, - U=8.0, -) +spinful = py.Fermion(spinful=True, symmetry="U1U1") spinless.operator("number") spinful.operator("n_up") spinful.hopping_operator(spin="up") spinful.interaction_operator() spinful.chemical_potential_operator() -spinful.onsite_gate(dt=0.01, site=0) -spinful.gate("interaction", dt=0.01) -spinless.gate_stream(edges, dt=0.01, order=2) -spinful.local_terms(edges) # native terms for energy optimization +spinful.onsite_gate(dt=0.01, site=0, U=8.0) +spinful.gate("interaction", dt=0.01, U=8.0) +spinless.gate_stream(edges, dt=0.01, order=2, t=1.0, V=0.5, mu=0.0) +spinful.local_terms(edges, t=1.0, U=8.0) # native terms for optimization ``` The bare ``*_operator`` methods return explicit native fermionic operators, @@ -247,7 +241,7 @@ expectation routines: edges = tuple(peps.edges) # ((x0, y0), (x1, y1)) sites = tuple(peps.sites) terms = {edge: -t * fermion.hopping_operator() for edge in edges} -terms |= {site: fermion.onsite_term(site) for site in sites} +terms |= {site: fermion.onsite_term(site, U=U, mu=mu) for site in sites} ham_peps = fermion.hamiltonian(terms) ``` @@ -258,7 +252,8 @@ for spinful fermions; spinless helpers provide ``create``, ``annihilate``, ``number``, and ``parity``. All returned operators retain the selected Symmray Abelian symmetry and fermionic grading. -The site-layout ``local_terms`` mapping is keyed by edges. Its onsite +The site-layout ``local_terms(edges, t=..., U=..., mu=...)`` mapping is keyed +by edges. Its onsite Hubbard and chemical-potential pieces are divided by each site's coordination inside the incident edge tensors, so summing the mapping still includes each one-site contribution exactly once. If those pieces should be visibly @@ -330,19 +325,19 @@ gates. ```python fermions = py.SpinfulFermion( symmetry="U1U1", # use "U1" to conserve only total particle number - t=1.0, - U=8.0, ) site_charge = fermions.half_filled_site_charge(L=16) number_up = fermions.observable("number_up") pair_create = fermions.observable("pair_create") -hamiltonian = fermions.hamiltonian(edges) +hamiltonian = fermions.hamiltonian(edges, t=1.0, U=8.0) # A symmetric, edge-coloured second-order step. Each colour contains # vertex-disjoint hopping bonds; forward then reverse colours make the hopping # product formula second order as well. -gates = fermions.strang_gate_stream(edges, dt=0.01, sites=range(16)) +gates = fermions.strang_gate_stream( + edges, dt=0.01, sites=range(16), t=1.0, U=8.0 +) ``` ``py.SymmFermions`` is the companion namespace for future symmetric-fermion @@ -371,7 +366,7 @@ edges = setup.edges occupations = setup.occupations terms = {edge: -t * fermion.hopping_operator() for edge in edges} -terms |= {site: fermion.onsite_term(site) for site in sites} +terms |= {site: fermion.onsite_term(site, U=U, mu=mu) for site in sites} ham = fermion.hamiltonian(terms) gate_stream = fermion.strang_gate_stream( @@ -379,6 +374,9 @@ gate_stream = fermion.strang_gate_stream( dt=0.01, sites=sites, imaginary=True, + t=t, + U=U, + mu=mu, ) ``` diff --git a/docs/api/vmc.md b/docs/api/vmc.md index 2183582..472ff6d 100644 --- a/docs/api/vmc.md +++ b/docs/api/vmc.md @@ -42,18 +42,33 @@ Each builder returns a `NetKetPEPSVMC` setup bundle with the NetKet Hilbert space, Hamiltonian, sampler, variational state, packed PEPS ansatz, and optional SR preconditioner. The bundle exposes `n_sites`, `n_params`, `setup.expect_energy()`, and `setup.make_driver(...)` so notebooks can keep the -main path compact. It also provides a clean class-level workflow with progress -bars: +main path compact. Pass `progress=True` to `build_fermion_vmc(...)` to show its +eight setup phases and inspect `setup.build_timing` for the measured phase costs. +It also provides a clean class-level workflow with progress bars: - `setup.warmup(progress=True)` — compile the sampler/amplitude/energy kernels - up front behind a small two-stage bar so the optimization ETA is meaningful. + up front behind a small staged bar so the optimization ETA is meaningful. +- `setup.benchmark_amplitude(n_samples=1)` — time a synchronized amplitude + batch; `.compile_seconds` and `.amplitude_seconds_per_sample` separate + shape-specific JAX compilation from steady-state evaluation. +- `setup.to_peps()` — reconstruct the current optimized Flax parameter tree as + the underlying quimb/Symmray PEPS network for Pepsy-side measurements or + persistence. +- `setup.run(n_iter, ...)` — concise wrapper that warms up and optimizes, while + returning separate warmup, optimization, total, and per-step timings. - `setup.optimize(n_iter, *, learning_rate=..., driver="vmc"|"vmc_sr", energy_shift=..., per_site=..., warmup=True, progress=True)` — run VMC with a single live energy progress bar and return a `VMCOptimizeResult` (`steps/energies/errors/variances`, `.shifted_energies`, and `.plot(...)`). -- `setup.measure(observables=None)` — evaluate observables (a single operator, - a `{name: operator}` mapping, or those stored on the setup) and return their - `nk.stats.Stats`. +- `setup.sample(sampling=None, progress=True)` — collect a retained batch and + report chains, burn-in, sweep size, acceptance rate, and elapsed time in + `VMCSamples.diagnostics`. +- `setup.measure_samples(samples, observables=None)` — evaluate one operator + or a `{name: operator}` mapping on exactly that retained NetKet batch, with + no additional sampling. A mapping value may be + `NetKetEtaPairObservable(...)`; it is compiled from the PEPS lattice only + after the shared sample cache is selected. `setup.measure(..., + samples=samples)` is the equivalent convenience form. For fermions, the supported JIT configuration intentionally combines a flat `Z2` Symmray PEPS ansatz with NetKet's fixed `(N_up, N_down)` `U1U1` @@ -158,8 +173,11 @@ The shared contracts do not import Torch, JAX, Flax, or NetKet. Each adapter compiles the same symbolic terms into its own connection/operator representation and preserves its native sampling and contraction strategy. `ContractionConfig`, `SamplingConfig`, and `OptimizationConfig` use the same -validated option names, including an explicit `n_samples_per_chain` -convention. +validated option names. The canonical sampling spelling follows NetKet: +`n_samples` (total), `n_chains`, `n_discard_per_chain`, and `sweep_size`. +The older `n_samples_per_chain`, `burn_in`, and `thin` spellings remain +compatibility aliases; `n_samples` must be divisible by `n_chains` so the +chain-preserving result has equal chain lengths. `compile_operator_sum_torch(...)` lowers the common terms to Torch connection tables (using a supplied `Fermion` object for graded symbolic factors), while @@ -173,7 +191,12 @@ The shared runtime settings are consumed directly by both façades: ```python from pepsy.vmc import OptimizationConfig, SamplingConfig -sampling = SamplingConfig(n_samples_per_chain=512, n_chains=32, burn_in=64) +sampling = SamplingConfig( + n_samples=16_384, + n_chains=32, + n_discard_per_chain=64, + sweep_size=2, +) torch_samples = torch_vmc.sample(sampling) netket_samples = netket_vmc.sample() # its sampler was built from MCState @@ -347,12 +370,74 @@ larger `chi` for production accuracy. The driver also deduplicates identical connected targets across current walkers during `local_energies`, `estimate_observable(s)`, and `step`; this is especially useful for serial U1/U1U1 boundary contractions. +For native U1/U1U1 CPU measurements, pass `boundary_workers > 1` to evaluate +independent cached boundary-window closures concurrently. This is restricted +to no-grad inference and defaults to `1` to avoid accidental BLAS +oversubscription; keep it at `1` when using CUDA or a shared cotengra optimizer +that is not thread-safe. The connected profile reports `num_requests`, +`num_reused`, `num_parallel`, and `num_fallback` for the current measurement. For no-grad calls on that serial boundary route, completed amplitudes are also cached by configuration and the current torch-parameter version, so repeated walkers can reuse the full boundary contraction. Inspect `model.last_amplitude_cache_stats` (or the `profile["cache"]["amplitude"]` entry); the cache is invalidated automatically after parameter updates and is never used for gradient-enabled or custom-parameter calls. + +### Native Torch throughput probes and rank-sharded sampling + +Use `benchmark_torch_amplitudes` (or +`TorchVMCDriver.benchmark_amplitudes`) with an existing configuration batch to +compare the supported amplitude batching and chunk sizes without drawing more +Markov samples. It moves the batch to the amplitude model's device, uses +`torch.no_grad()`, synchronizes CUDA timing, and verifies each candidate +against the first result by default. Boundary-amplitude cache hits are disabled +by default so the result reflects contraction throughput; pass +`include_cache=True` to time the cache-aware serving path. The returned +`TorchAmplitudeBenchmark.executed_batching` is the route actually selected, +which makes an unavailable `"vmap"` request visible as `"serial"`. + +```python +timing = vmc.benchmark_amplitudes( + samples.configs, + chunk_sizes=(None, 32, 64), + amplitude_batchings=("serial", "auto", "vmap"), + repeats=5, +) +print(timing.best) +``` + +After the application initializes an optional `torch.distributed` process +group, `TorchVMCDriver.sample(..., distributed=True)` and the corresponding +`TorchFermionVMC.sample`, `run_measurement`, and measurement-mode `run` accept +a rank-sharded native Markov calculation. `SamplingConfig.n_chains` is the +global count and must be at least the world size; it is split deterministically +among ranks. Set `seed` or `sampler_seed` in that configuration so each rank +receives a distinct reproducible stream. Each rank retains only its local PEPS +configurations and amplitudes, while unweighted observable moments, acceptance +counts, and rank-local effective sample sizes are reduced to global estimates. + +```python +# Launch this script with torchrun after selecting each rank's CUDA device. +torch.distributed.init_process_group("nccl") +samples = vmc.sample( + sampling=pvmc.SamplingConfig( + n_samples=4096, + n_chains=128, # global count + sampler_seed=7, + ), + distributed=True, +) +energy = vmc.measure(samples) # detects the distributed sample metadata +print(energy.n_samples, energy.distributed.global_n_chains) +``` + +`TorchMCMCSamples.distributed` and `TorchVMCEnergyEstimate.distributed` record +the local and global counts; `TorchMCMCSamples.to_common()` preserves the same +metadata in `VMCSamples.diagnostics["distributed"]`. Distributed measurement +does not all-gather histories, so global R-hat is intentionally unavailable. +It currently supports only unweighted native Markov samples—not external +proposal/importance weights or multi-rank SR optimization. + After measuring an observable, `result.chain_diagnostics` reports `r_hat`, the integrated autocorrelation time, and an effective sample size when there are at least two chains and two retained samples per chain. In that case, @@ -432,6 +517,7 @@ model = pvmc.TorchPEPSBoundaryAmplitude( chi=64, cutoff=1e-10, dtype=torch.float64, + boundary_workers=4, # CPU inference only; benchmark this on your machine. ) ``` @@ -640,16 +726,21 @@ observable. `estimate_energy(...)` remains as a compatibility alias. For compatibility with an existing lower-level sweep loop, coordinate-labelled PEPS can still initialize `TorchFermionVMC` in the constructor. New code -should prefer the first-run recipe below instead. Pass `fermion` to generate -the default Hamiltonian, or omit it when supplying explicit `terms`: +should prefer the first-run recipe below instead. Supply an explicit native +Hamiltonian (or its terms) because `Fermion` intentionally stores no model +couplings: ```python from pepsy import Fermion -fermion = Fermion(spinful=True, symmetry="U1U1", t=t, U=U) +fermion = Fermion(spinful=True, symmetry="U1U1") +edges = ((0, 1), (1, 2), (2, 3)) +terms = {edge: -t * fermion.hopping_operator() for edge in edges} +terms |= {site: fermion.onsite_term(site, U=U) for site in range(4)} vmc = pvmc.TorchFermionVMC( peps, - fermion, + fermion=fermion, + terms=terms, n_walkers=128, contraction="boundary", # or "exact" / "ctmrg" / "hotrg" chi=32, # required for an approximate contraction @@ -667,11 +758,12 @@ For a native fermionic PEPS measurement, construct `TorchFermionVMC` from the state and native Fermion terms only. The first `sample` (or `warmup`) owns both the chain recipe and PEPS contraction recipe: -1. `SamplingConfig` owns the number of chains, retained samples, burn-in, - thinning, and RNG seeds. In the native Torch sampler, `burn_in` counts - discarded thinning intervals, so the `Metropolis` total is - `(burn_in + n_samples_per_chain) * thin` batched sweeps; every batched - sweep advances all chains once. +1. `SamplingConfig` owns the total retained samples, number of chains, + per-chain discard count, sweep spacing, and RNG seeds. The canonical + keywords are `n_samples`, `n_chains`, `n_discard_per_chain`, and + `sweep_size`. In the native Torch sampler, the Metropolis total is + `(n_discard_per_chain + n_samples_per_chain) * sweep_size` batched + sweeps; every batched sweep advances all chains once. 2. `contraction_opts` is a single mapping with `method`, `chi`, `cutoff`, and any backend options such as `mode`. It is consumed when the first operation builds the amplitude model, then remains fixed with the Markov state. @@ -681,10 +773,10 @@ the chain recipe and PEPS contraction recipe: ```python sampling = pvmc.SamplingConfig( - n_samples_per_chain=256, + n_samples=8192, n_chains=32, - burn_in=64, - thin=2, + n_discard_per_chain=64, + sweep_size=2, seed=7, ) contraction_opts = { @@ -1031,26 +1123,48 @@ driver = setup.make_driver( ### General fermion models and observables `build_fermion_vmc(...)` lifts the Fermi-Hubbard specialization to any spinful -fermion model. Pass a `pepsy.Fermion` (its hopping `t`, on-site `U`, -nearest-neighbor density `V`, and chemical potential `mu` are turned into a -NetKet fermion operator over the lattice edges), an explicit list of symbolic -terms, or a ready NetKet operator. Optional observables are stored on the setup -and evaluated with `setup.measure(...)`. +fermion model. Pass a `pepsy.Fermion` together with an explicit native +Hamiltonian/term mapping, an explicit list of symbolic terms, or a ready +NetKet operator. `Fermion` owns only local symmetry and backend conventions; +the hopping `t`, on-site `U`, density `V`, and chemical potential `mu` live in +the explicit term arrays. Optional observables are stored on the setup and +evaluated with `setup.measure(...)`. Like `TorchFermionVMC`, it can also infer the lattice geometry directly from a native Pepsy Hamiltonian. Pass the `SymHamiltonian` from `fermion.hamiltonian(...)` as `hamiltonian=`, or its coordinate-keyed `.terms` mapping as `terms=`, together with `fermion=`; the builder reads the -integer edges and periodic axes (`pbc`) from those terms and rebuilds the -matching NetKet Hamiltonian. Explicit `edges` / `graph` / `pbc` still take -precedence. +integer edges and periodic axes (`pbc`) from those terms and compiles the +supplied native local operators to NetKet. Explicit `edges` / `graph` / `pbc` +still take precedence. + +`SamplingConfig(n_samples=..., n_chains=..., n_discard_per_chain=..., +sweep_size=..., chunk_size=...)` controls retained samples, chains, per-chain +discard, sweep spacing, and batching. NetKet's Metropolis sampler also performs +`sweep_size` proposals +between retained samples; pass `sampler_sweep_size=...` to make that cost +explicit (the NetKet default is the Hilbert-space size). + +The generic fermion builder defaults to `conserving=False`, which constructs +the exact ordinary NetKet fermion operator without a first-use conversion +compile. Set `conserving="auto"` when the specialized conserving operator is +worth the one-time conversion cost for a long production run. + +For a compact construction call, put these numerical settings in +`NetKetVMCConfig(...)` and pass it as `config=`. Boundary/PEPS contraction +settings belong there because changing `chi` changes the compiled amplitude +model; retained sample counts and burn-in can instead be overridden at +`setup.sample(...)` time. ```python import pepsy as py import pepsy.vmc as pvmc -fermion = py.Fermion(spinful=True, symmetry="U1U1", t=1.0, U=8.0, V=0.5) -ham = fermion.hamiltonian(terms) # native SymHamiltonian over the lattice +fermion = py.Fermion(spinful=True, symmetry="U1U1") +terms = {edge: -1.0 * fermion.hopping_operator() for edge in edges} +terms |= {site: fermion.onsite_term(site, U=8.0) for site in sites} +terms |= {edge: 0.5 * fermion.density_operator() for edge in edges} +ham = fermion.hamiltonian(terms) # authoritative native Hamiltonian # Clean API: infer edges + PBC from the native terms, build the NetKet model. setup = pvmc.build_fermion_vmc( @@ -1060,10 +1174,14 @@ setup = pvmc.build_fermion_vmc( ``` ```python -fermion = py.Fermion(spinful=True, symmetry="U1U1", t=1.0, U=8.0, V=0.5) +fermion = py.Fermion(spinful=True, symmetry="U1U1") +terms = {edge: -1.0 * fermion.hopping_operator() for edge in edges} +terms |= {site: fermion.onsite_term(site, U=8.0) for site in sites} +ham = fermion.hamiltonian(terms) setup = pvmc.build_fermion_vmc( peps, fermion=fermion, + hamiltonian=ham, Lx=4, Ly=4, contraction="ctmrg", @@ -1083,7 +1201,23 @@ result = setup.optimize( per_site=setup.n_sites, ) result.plot(per_site=setup.n_sites) -stats = setup.measure() # {name: nk.stats.Stats} + +# Sample once, then measure each operator on those exact Markov chains. +samples = setup.sample() +stats = setup.measure_samples( + samples, + { + "energy": setup.hamiltonian, + # The coordinate lattice is inferred from setup.ansatz.orbital_sites. + "eta_pair": pvmc.NetKetEtaPairObservable( + 1, + 0, + periodic=True, + staggered=True, + ), + **pvmc.standard_fermion_observables(setup.hilbert), + }, +) # {name: nk.stats.Stats} ``` Two lower-level helpers back this path: @@ -1094,13 +1228,17 @@ Two lower-level helpers back this path: `sz` in `{+1, -1, None}`. It doubles as a Hamiltonian or an observable, and can optionally convert to NetKet's particle-number/spin-conserving operator for cheaper local energies. -- `fermion_model_terms(fermion, edges, *, n_sites=None)` returns the symbolic - hopping / Hubbard / density / chemical-potential terms for a spinful - `pepsy.Fermion`, so custom models can start from the standard terms and add - their own. +- `fermion_model_terms(fermion, edges, *, t, U, V=0.0, mu=0.0, n_sites=None)` + returns symbolic uniform Fermi-Hubbard terms from explicit couplings. For + non-uniform models, pass the authoritative native term mapping directly to + `build_fermion_vmc(...)`. - `standard_fermion_observables(hilbert)` returns common observables (`n_up`, `n_down`, `n_total`, `double_occupancy`) as a ready `{name: operator}` mapping for `setup.measure(...)`. +- `NetKetEtaPairObservable(dx, dy, *, periodic=True, staggered=False)` is a + declarative eta-pair correlator for `measure_samples(...)`. Its zero-offset + form is the mean double occupancy; nonzero offsets measure + `Delta_i^dag Delta_j + h.c.` and use either periodic or in-bounds pairs. the PEPS once, keep `flat=True` Symmray data for JIT-friendly leaves, and evaluate many spin-orbital occupation rows with one compiled function. In the currently tested Symmray stack this flat path is the `Z2` route; sparse @@ -1173,8 +1311,15 @@ The approximate contraction choices are Quimb's finite 2D contractions: `contraction="hotrg"` calls `contract_hotrg`, `contraction="ctmrg"` calls `contract_ctmrg`, and `contraction="boundary"` or `"mps"` calls `contract_boundary(mode="mps")`. Pass `contraction_opts={...}` for lower-level -Quimb options such as `sequence`, `max_separation`, `canonize`, or a non-default -boundary `mode`. +Quimb options. These options are method-specific: `sequence`, +`max_separation`, and `canonize` are boundary-MPS controls, while CTMRG's +stable default is `mode="projector"`. + +For flat fermionic Symmray tensors under JAX, `max_separation=0` is accepted +but has a guarded compatibility fallback to `1` if Quimb reaches its upstream +empty-boundary-axis or block-matmul path. This applies to both boundary-MPS +and CTMRG. The requested `sequence`, `chi`, and `canonize` settings remain +active; dense/non-Symmray contractions are not modified. For larger lattices use an approximate contraction with `chi=...`, and treat `chi`, `chunk_size`, `chunk_size_bwd`, `n_samples`, and the SR setting as diff --git a/docs/development/plans/project.md b/docs/development/plans/project.md index 541e50b..a2af2e9 100644 --- a/docs/development/plans/project.md +++ b/docs/development/plans/project.md @@ -453,6 +453,9 @@ Symmray optional (the `vmc` extra), so the core package stays lightweight. `recommend_netket_vmc_settings`, `choose_netket_chunk_size`, plus the `SpinOrbitalColumns` / `PackedFermionicPEPS` / `NetKetVMCSettings` / `NetKetChunkSettings` / `NetKetFermiHubbardVMC` dataclasses. +- `NetKetPEPSVMC.to_peps()` reconstructs the current optimized NetKet/Flax + variables as the underlying quimb/Symmray PEPS network, so downstream Pepsy + measurement or persistence can resume from an optimized VMC state. - Validated end-to-end in `../pepsy_examples/fermi_hubbard/fermi_hubbard_vmc.ipynb` (2×2, `t=1`, `U=8`, `D=4`, `Z2`, half filling): @@ -470,18 +473,16 @@ Symmray optional (the `vmc` extra), so the core package stays lightweight. ### Plan (next Pepsy cuts) -1. `to_peps(variables)` / `update_peps_from_variables(...)` to convert an - optimized NetKet state back into a `SymPEPS`. -2. `benchmark_log_amplitude_batch(...)` timing batch / chunk sizes and +1. `benchmark_log_amplitude_batch(...)` timing batch / chunk sizes and `exact` vs `hotrg` without a full VMC optimization. -3. `chi_sweep_energy(...)` for HOTRG `chi` convergence of sampled energies. -4. A GPU example script/notebook that measures after warmup and reports device, +2. `chi_sweep_energy(...)` for HOTRG `chi` convergence of sampled energies. +3. A GPU example script/notebook that measures after warmup and reports device, memory, batch size, and compile time separately. -5. `U1U1` and odd-parity `Z2` setup helpers so 3×3 / odd half-filled systems +4. `U1U1` and odd-parity `Z2` setup helpers so 3×3 / odd half-filled systems are not awkward. -6. Scalable SR: diagonal / iterative / minSR-style paths before dense SR on +5. Scalable SR: diagonal / iterative / minSR-style paths before dense SR on large PEPS parameter counts. -7. Keep CI tiny and CPU-only; large GPU timing stays in examples. +6. Keep CI tiny and CPU-only; large GPU timing stays in examples. Development rule of thumb (from the notebook): reusable NetKet/Symmray glue moves into `pepsy.vmc`; physics sanity checks and timing experiments stay in the diff --git a/examples/qmera_fermion_hubbard_2d.py b/examples/qmera_fermion_hubbard_2d.py index 37270f3..2e0d110 100644 --- a/examples/qmera_fermion_hubbard_2d.py +++ b/examples/qmera_fermion_hubbard_2d.py @@ -67,16 +67,10 @@ def product_state_factory(schedule, sites, **kwargs): param_scale=0.02, product_state_factory=product_state_factory, ) - fermion = py.Fermion( - spinful=True, - symmetry="U1U1", - t=0.2, - U=4.0, - mu=0.1, - ) + fermion = py.Fermion(spinful=True, symmetry="U1U1") schedule = builder.build_schedule() parameters = builder.initialize_parameters(schedule) - terms = builder.fermion_terms(fermion) + terms = builder.fermion_terms(fermion, t=0.2, U=4.0, mu=0.1) # Symmray operators are already native; preserve them with # convert_terms=False. The cache reuses paths for repeated local cones. diff --git a/examples/qmera_fermion_hubbard_4x4_pbc.py b/examples/qmera_fermion_hubbard_4x4_pbc.py index d5c9b6b..8a5c670 100644 --- a/examples/qmera_fermion_hubbard_4x4_pbc.py +++ b/examples/qmera_fermion_hubbard_4x4_pbc.py @@ -57,11 +57,10 @@ def main(): ) schedule = builder.build_schedule() - # The Fermion helper supplies the onsite U and chemical-potential terms; - # the qMERA unitary above supplies the native number-conserving hopping - # layer. Keeping these roles separate makes the U1U1 convention explicit. - fermion = py.Fermion(spinful=True, symmetry="U1U1", t=0.2, U=4.0, mu=0.1) - terms = builder.fermion_terms(fermion) + # Fermion fixes the local U1U1 convention only; physical couplings remain + # explicit at term construction so every simulation path sees the same model. + fermion = py.Fermion(spinful=True, symmetry="U1U1") + terms = builder.fermion_terms(fermion, t=0.2, U=4.0, mu=0.1) print("RG register sizes:", [len(layer.input_sites) for layer in schedule.layers], "->", len(schedule.top_sites)) print("first-layer isometry blocks:", len(schedule.layers[0].isometry_blocks)) diff --git a/src/pepsy/backends/config.py b/src/pepsy/backends/config.py index b56b771..1856918 100644 --- a/src/pepsy/backends/config.py +++ b/src/pepsy/backends/config.py @@ -435,7 +435,7 @@ def reg_complex_svd_jax(): def reg_rel_svd_jax(): - """Register JAX SVD with Pepsy's custom VJP rule in autoray.""" + """Register JAX SVD with Pepsy's truncation-safe VJP rule in autoray.""" try: __import__("jax") except ImportError as exc: # pragma: no cover - exercised in no-jax CI @@ -450,7 +450,7 @@ def reg_rel_svd_jax(): def reg_real_svd_jax(): - """Register JAX SVD custom-VJP rule for real-valued SVD workloads.""" + """Register JAX SVD's truncation-safe VJP rule for real workloads.""" try: __import__("jax") except ImportError as exc: # pragma: no cover - exercised in no-jax CI diff --git a/src/pepsy/backends/linalg_jax.py b/src/pepsy/backends/linalg_jax.py index 47c3fb1..c989409 100644 --- a/src/pepsy/backends/linalg_jax.py +++ b/src/pepsy/backends/linalg_jax.py @@ -1,76 +1,101 @@ -"""JAX-side linalg registrations with custom VJP rules.""" +"""JAX-side linalg registrations with truncation-safe VJP rules.""" import autoray as ar +import jax import jax.numpy as jnp from jax import custom_vjp @custom_vjp def svd_jax(A): - """JAX SVD primitive wrapped with a custom VJP definition.""" + """Thin JAX SVD with a Quimb-truncation-safe backward rule. + + Quimb's ``svd_truncated`` can pass cotangents only for the singular-vector + columns it retained. The custom VJP restores those leading columns to the + full thin-SVD output shape before delegating the actual derivative to + JAX's maintained SVD pullback. This is important for approximate tensor + contractions, where a fixed ``max_bond`` is the normal JIT-compatible + path. + """ return jnp.linalg.svd(A, full_matrices=False) - -def _safe_reciprocal(x, epsilon=1.0e-12): - """Regularized reciprocal used by the JAX SVD backward expressions.""" - return x / (x * x + epsilon) - - def h(x): """Return the conjugate transpose of ``x`` (Hermitian transpose).""" return jnp.conj(jnp.transpose(x)) -def jaxsvd_fwd(A): - """Forward rule for :func:`svd_jax` custom VJP.""" - u, s, v = svd_jax(A) - return (u, s, v), (u, s, v) - - -def jaxsvd_bwd(residual, tangents): - """Backward rule for :func:`svd_jax` custom VJP.""" - U, S, V = residual - du, ds, dv = tangents - - dU = jnp.conj(du) - dS = jnp.conj(ds) - dV = jnp.transpose(dv) - - ms = jnp.diag(S) - ms1 = jnp.diag(_safe_reciprocal(S)) - dAs = U @ jnp.diag(dS) @ V +def _restore_truncated_tangent(tangent, full, *, axis): + """Pad a leading-rank Quimb cotangent to a thin-SVD output shape.""" + axis %= full.ndim + if tangent is None: + return jnp.zeros_like(full) + if tangent.shape == full.shape: + return tangent + + if tangent.ndim != full.ndim: + raise TypeError( + "SVD cotangent rank does not match the corresponding thin-SVD " + f"output: got {tangent.shape!r}, expected {full.shape!r}." + ) + for dim, (actual, expected) in enumerate(zip(tangent.shape, full.shape)): + if dim != axis and actual != expected: + raise TypeError( + "SVD cotangent shape is incompatible with the corresponding " + f"thin-SVD output: got {tangent.shape!r}, expected " + f"{full.shape!r}." + ) + if tangent.shape[axis] > full.shape[axis]: + raise TypeError( + "SVD cotangent has more singular-vector components than the " + f"thin-SVD output: got {tangent.shape!r}, expected " + f"{full.shape!r}." + ) + + slices = [slice(None)] * full.ndim + slices[axis] = slice(0, tangent.shape[axis]) + return jnp.zeros_like(full).at[tuple(slices)].set(tangent) - F = S * S - (S * S)[:, None] - F = _safe_reciprocal(F) - jnp.diag(jnp.diag(_safe_reciprocal(F))) - J = F * (h(U) @ dU) - dAu = U @ (J + h(J)) @ ms @ V - - K = F * (V @ dV) - dAv = U @ ms @ (K + h(K)) @ V +def jaxsvd_fwd(A): + """Forward rule for :func:`svd_jax`, retaining full thin-SVD shapes.""" + outputs = jnp.linalg.svd(A, full_matrices=False) + return outputs, (A, outputs) - O = h(dU) @ U @ ms1 - dAc = -1 / 2.0 * U @ (jnp.diag(jnp.diag(O - jnp.conj(O)))) @ V - dAv = dAv + U @ ms1 @ h(dV) @ (jnp.eye(jnp.size(V[1, :])) - h(V) @ V) - dAu = dAu + (jnp.eye(jnp.size(U[:, 1])) - U @ h(U)) @ dU @ ms1 @ V - grad_a = jnp.conj(dAv + dAu + dAs + dAc) - return (grad_a,) +def jaxsvd_bwd(residual, tangents): + """Differentiate a thin SVD after restoring Quimb's truncated tangents.""" + A, outputs = residual + U, S, Vh = outputs + dU, dS, dVh = tangents + cotangents = ( + _restore_truncated_tangent(dU, U, axis=-1), + _restore_truncated_tangent(dS, S, axis=-1), + _restore_truncated_tangent(dVh, Vh, axis=-2), + ) + _, pullback = jax.vjp( + lambda matrix: jnp.linalg.svd(matrix, full_matrices=False), + A, + ) + cotangent_tree = jax.tree_util.tree_unflatten( + jax.tree_util.tree_structure(outputs), + cotangents, + ) + return pullback(cotangent_tree) svd_jax.defvjp(jaxsvd_fwd, jaxsvd_bwd) def reg_complex_svd_jax(): - """Register the custom JAX SVD implementation in autoray.""" + """Register the truncation-safe JAX thin-SVD implementation in autoray.""" ar.register_function("jax", "linalg.svd", svd_jax) def reg_rel_svd_jax(): - """Register the custom JAX SVD implementation in autoray.""" + """Register the truncation-safe JAX thin-SVD implementation in autoray.""" reg_complex_svd_jax() def reg_real_svd_jax(): - """Register the custom JAX SVD implementation in autoray.""" + """Register the truncation-safe JAX thin-SVD implementation in autoray.""" reg_complex_svd_jax() diff --git a/src/pepsy/optimizers/mera/builders.py b/src/pepsy/optimizers/mera/builders.py index 450a3d1..22f6dc6 100644 --- a/src/pepsy/optimizers/mera/builders.py +++ b/src/pepsy/optimizers/mera/builders.py @@ -397,7 +397,8 @@ def fermion_terms(self, fermion=None, **params): on the builder keeps the representation choice in one place while allowing the regular local-term and optimizer machinery to handle the result. If the builder was created with ``fermion=...``, the argument - can be omitted. + can be omitted. Pass physical couplings (for example ``t=`` and + ``U=``) explicitly; ``Fermion`` deliberately does not store them. """ fermion = self.fermion if fermion is None else fermion if fermion is None: diff --git a/src/pepsy/optimizers/mera/fermions.py b/src/pepsy/optimizers/mera/fermions.py index cd6a024..fb322ef 100644 --- a/src/pepsy/optimizers/mera/fermions.py +++ b/src/pepsy/optimizers/mera/fermions.py @@ -467,7 +467,8 @@ def qmera_symmray_fermi_hubbard_terms(geometry, *, fermion=None, **kwargs): geometry : QMeraGeometry Geometry whose physical sites are expanded into two-state modes. fermion : pepsy.Fermion, optional - Unified model helper supplying ``symmetry``, ``t``, ``U``, and ``mu``. + Unified model helper supplying only the local symmetry convention. + Pass ``t=``, ``U=``, and optional ``mu=`` explicitly. qMERA deliberately remains mode-native, so this adapter accepts only a spinful ``U1U1`` helper and does not turn a four-state site tensor into a qMERA register tensor. @@ -479,12 +480,13 @@ def qmera_symmray_fermi_hubbard_terms(geometry, *, fermion=None, **kwargs): raise ValueError( "qMERA Hubbard terms currently require Fermion(symmetry='U1U1')." ) - kwargs = { - "t": getattr(fermion, "t", 1.0), - "U": getattr(fermion, "U", 8.0), - "mu": getattr(fermion, "mu", 0.0), - **kwargs, - } + missing = [name for name in ("t", "U") if name not in kwargs] + if missing: + raise TypeError( + "qMERA Fermi-Hubbard terms require explicit " + + ", ".join(f"{name}=" for name in missing) + + ". Fermion does not store couplings." + ) backend = kwargs.pop("backend", None) if backend is None: backend = QMeraSymmrayFermionBackend( diff --git a/src/pepsy/tensors/constructors.py b/src/pepsy/tensors/constructors.py index 64d6b14..e148cac 100644 --- a/src/pepsy/tensors/constructors.py +++ b/src/pepsy/tensors/constructors.py @@ -1218,6 +1218,7 @@ def hrs_to_peps( subsizes="maximal", contraction_opt="auto-hq", to_backend=None, + normalize=False, ): """Create a random product or Fermion-symmetric PEPS. @@ -1225,7 +1226,8 @@ def hrs_to_peps( With ``fermion``, construct a native charge-preserving random PEPS instead: ``method="direct"`` uses Symmray's direct block-filled random PEPS, with ``chi`` controlling the virtual bond dimension. The direct state is - normalized before it is returned. A unitary PEPS-growth method is not yet + returned without a global norm by default; pass ``normalize=True`` when a + normalized PEPS is required. A unitary PEPS-growth method is not yet implemented. In the fermionic branch, ``haar_params`` and ``perturb`` do not apply. @@ -1260,13 +1262,20 @@ def hrs_to_peps( Advanced override for the per-site charge pattern. method : {"direct"}, optional Fermion-aware random-state construction. ``"direct"`` fills allowed - Symmray blocks using ``PEPS_fermionic_rand`` and normalizes the result. + Symmray blocks using ``PEPS_fermionic_rand``. Global normalization is + optional and disabled by default. subsizes : object, optional Symmray charge-sector sizing policy used by ``method="direct"``. contraction_opt : object, optional Contraction optimizer stored by the internal symmetric wrapper. to_backend : callable, optional Backend mapper applied to Fermion-aware Symmray blocks. + normalize : bool, optional + Whether to globally normalize a Fermion-aware PEPS before returning + it. Defaults to ``False``. ``normalize=True`` uses a CPU boundary-MPS + contraction and is only needed when the caller requires a global + physical norm. NetKet VMC uses the default safely: a global PEPS + scalar cancels from the Metropolis probability and local-energy ratios. Returns ------- @@ -1288,6 +1297,8 @@ def hrs_to_peps( method = str(method).strip().lower().replace("-", "_") if method not in {"direct", "unitary"}: raise ValueError("method must be 'direct' or 'unitary'.") + if not isinstance(normalize, bool): + raise TypeError("normalize must be a bool.") if fermion is not None: from .symmetric import ( # pylint: disable=import-outside-toplevel @@ -1347,7 +1358,16 @@ def hrs_to_peps( contraction_opt=contraction_opt, to_backend=None, ) - state.normalize() + if normalize: + try: + state.normalize() + except Exception as exc: + raise RuntimeError( + "Fermionic PEPS global normalization failed in Quimb's " + "boundary-MPS decomposition. For a NetKet VMC initial " + "state, pass normalize=False: its global amplitude scale " + "cancels from sampling and local-energy ratios." + ) from exc if to_backend is not None: state.apply_to_arrays(to_backend) return state.peps diff --git a/src/pepsy/tensors/symmetric.py b/src/pepsy/tensors/symmetric.py index 2132423..f3b9e09 100644 --- a/src/pepsy/tensors/symmetric.py +++ b/src/pepsy/tensors/symmetric.py @@ -7314,12 +7314,13 @@ def fermion_hopping_param_gen( class Fermion: """Native spinless or spinful fermion observables, gates, and streams. - The helper owns the local fermionic space, symmetry convention, optional - model parameters, and backend conversion. Explicit native terms can still - be passed to :meth:`hamiltonian`; model parameters are used for compact - lattice and qMERA workflows. It is intended for direct Symmray-backed - fermionic MPS or PEPS workflows and does not introduce a qubit or - Jordan-Wigner circuit representation. + The helper owns only the local fermionic space, symmetry convention, and + optional backend conversion. Hamiltonian couplings are deliberately not + stored here: construct them as explicit native terms, then validate and + bundle them with :meth:`hamiltonian`. This prevents a native Hamiltonian, + a gate stream, and a VMC adapter from silently using different couplings. + It is intended for direct Symmray-backed fermionic MPS or PEPS workflows; + it does not introduce a qubit or Jordan-Wigner circuit representation. ``strang_gate_stream`` uses a deterministic edge colouring and a forward/reverse half-step sequence. Consequently its hopping layers are @@ -7329,13 +7330,9 @@ class Fermion: """ symmetry: str | None = None - t: object = 1.0 - U: object = 8.0 dtype: object = "complex128" to_backend: object = None spinful: bool = True - V: object = 0.0 - mu: object = 0.0 _dense_ops: dict = field(default_factory=dict, init=False, repr=False) _observable_cache: dict = field(default_factory=dict, init=False, repr=False) _gate_cache: dict = field(default_factory=dict, init=False, repr=False) @@ -8266,13 +8263,12 @@ def hopping_operator(self, *, spin=None, peierls_angle=0.0): peierls_angle=peierls_angle, ) - def hopping_term(self, edge, *, spin=None, t=None, peierls_angle=0.0): + def hopping_term(self, edge, *, spin=None, t, peierls_angle=0.0): """Return ``-t`` times the hopping operator on ``edge``.""" try: left, right = tuple(edge) except (TypeError, ValueError) as exc: raise ValueError("edge must contain exactly two site labels.") from exc - t = self.t if t is None else t t = _edge_parameter(t, left, right) return self._hopping_operator_on_sites( left, @@ -8295,9 +8291,8 @@ def interaction_operator(self): sites=(0,), ) - def interaction_term(self, site, *, U=None): + def interaction_term(self, site, *, U): """Return ``U n_up n_down`` on one physical site.""" - U = self.U if U is None else U U = _node_parameter(U, site) return self.operator_term( [(U, ((site, "double"),))], @@ -8315,9 +8310,8 @@ def chemical_potential_operator(self): terms = [(1.0, ((0, "number"),))] return self.operator_term(terms, sites=(0,)) - def chemical_potential_term(self, site, *, mu=None): + def chemical_potential_term(self, site, *, mu): """Return ``-mu n`` on one physical site.""" - mu = self.mu if mu is None else mu if self.spinful: mu = _node_parameter(mu, site) mu_up, mu_down = _as_spin_pair(mu, name="mu") @@ -8329,12 +8323,12 @@ def chemical_potential_term(self, site, *, mu=None): terms = [(-_node_parameter(mu, site), ((site, "number"),))] return self.operator_term(terms, sites=(site,)) - def onsite_term(self, site, *, U=None, mu=None): + def onsite_term(self, site, *, U=None, mu=0.0): """Return ``U n_up n_down - mu n`` on one site.""" - U = self.U if U is None else U - mu = self.mu if mu is None else mu terms = [] if self.spinful: + if U is None: + raise TypeError("onsite_term requires explicit U=... for spinful fermions.") terms.append((_node_parameter(U, site), ((site, "double"),))) mu = _node_parameter(mu, site) mu_up, mu_down = _as_spin_pair(mu, name="mu") @@ -8361,13 +8355,12 @@ def density_operator(self): ] return self.operator_term(terms, sites=(0, 1)) - def density_term(self, edge, *, V=None): + def density_term(self, edge, *, V): """Return ``V n_i n_j`` on a physical edge.""" try: left, right = tuple(edge) except (TypeError, ValueError) as exc: raise ValueError("edge must contain exactly two site labels.") from exc - V = self.V if V is None else V V = _edge_parameter(V, left, right) if self.spinful: names = ("number_up", "number_down") @@ -8515,7 +8508,7 @@ def heisenberg_gate(self, theta, *, edge=None, imaginary=False): szz_gate = spin_z_correlator_gate xy_gate = xy_exchange_gate - def interaction_gate(self, dt, *, site=None, U=None, imaginary=False): + def interaction_gate(self, dt, *, site=None, U, imaginary=False): """Return the exact onsite interaction gate. With a site-dependent ``U`` mapping or callable, pass ``site`` so the @@ -8527,7 +8520,6 @@ def interaction_gate(self, dt, *, site=None, U=None, imaginary=False): "Spinless fermions have no onsite doublon interaction; use " "density_gate(...) for the nearest-neighbor V interaction." ) - U = self.U if U is None else U U = U if site is None else _node_parameter(U, site) theta = dt * U @@ -8541,20 +8533,20 @@ def build(): return self._cached_gate(("interaction", dt, site, U, imaginary), build) - def onsite_gate(self, dt, *, site=None, U=None, mu=None, imaginary=False): + def onsite_gate(self, dt, *, site=None, U=None, mu=0.0, imaginary=False): """Return the complete one-site Hubbard gate. The generated gate represents ``U n_up n_down - mu n`` for spinful fermions and ``-mu n`` for spinless fermions. ``U`` and ``mu`` may be site-dependent mappings or callables when ``site`` is supplied. """ - U = self.U if U is None else U - mu = self.mu if mu is None else mu if site is not None: U = _node_parameter(U, site) mu = _node_parameter(mu, site) if self.spinful: + if U is None: + raise TypeError("onsite_gate requires explicit U=... for spinful fermions.") mu_up, mu_down = _as_spin_pair(mu, name="mu") U_site = U diagonal = ( @@ -8579,9 +8571,8 @@ def build(): return self._cached_gate(("onsite", dt, site, U, mu, imaginary), build) - def hopping_gate(self, dt, *, t=None, peierls_angle=0.0, imaginary=False): + def hopping_gate(self, dt, *, t, peierls_angle=0.0, imaginary=False): """Return a two-site native fermionic hopping gate with Peierls phase.""" - t = self.t if t is None else t def build(): if not self.spinful: gate = _spinless_hopping_gate( @@ -8603,13 +8594,12 @@ def build(): return self._cached_gate(("hopping", dt, t, peierls_angle, imaginary), build) - def density_gate(self, dt, *, V=None, imaginary=False): + def density_gate(self, dt, *, V, imaginary=False): """Return the nearest-neighbor density interaction gate. For spinless fermions this is ``V n_i n_j``. For spinful fermions it is ``V (n_up + n_down)_i (n_up + n_down)_j``. """ - V = self.V if V is None else V theta = dt * V def build(): @@ -8625,9 +8615,8 @@ def build(): return self._cached_gate(("density", dt, V, imaginary), build) - def chemical_potential_gate(self, dt, *, mu=None, site=None, imaginary=False): + def chemical_potential_gate(self, dt, *, mu, site=None, imaginary=False): """Return the chemical-potential part of an onsite gate.""" - mu = self.mu if mu is None else mu mu = mu if site is None else _node_parameter(mu, site) if self.spinful: mu_up, mu_down = _as_spin_pair(mu, name="mu") @@ -8649,7 +8638,7 @@ def build(): return self._cached_gate(("chemical", dt, site, mu, imaginary), build) def gate(self, name, dt, *, site=None, where=None, imaginary=False, **params): - """Build a named native gate using the model's local conventions.""" + """Build a named native gate using the local fermionic conventions.""" edge = params.pop("edge", where) del where # Gate locations belong to the stream entry, not the tensor. name = str(name).lower().replace("-", "_") @@ -8674,7 +8663,7 @@ def gate(self, name, dt, *, site=None, where=None, imaginary=False, **params): dt, site=site, U=params.pop("U", None), - mu=params.pop("mu", None), + mu=params.pop("mu", 0.0), imaginary=imaginary, ) if name in {"interaction", "onsite_interaction", "doublon"}: @@ -8872,16 +8861,16 @@ def gate_stream( imaginary=False, t=None, U=None, - V=None, - mu=None, + V=0.0, + mu=0.0, ): - """Return a canonical first- or second-order fermion gate stream.""" + """Return a canonical fermion gate stream with explicit couplings.""" if order not in {1, 2}: raise ValueError("order must be 1 or 2.") - t = self.t if t is None else t - U = self.U if U is None else U - V = self.V if V is None else V - mu = self.mu if mu is None else mu + if t is None: + raise TypeError("gate_stream requires explicit t=... .") + if self.spinful and U is None: + raise TypeError("gate_stream requires explicit U=... for spinful fermions.") edges = _as_edges(edges) sites = _sites_from_edges(edges, sites) @@ -8954,14 +8943,16 @@ def strang_gate_stream( imaginary=False, t=None, U=None, - V=None, - mu=None, + V=0.0, + mu=0.0, ): - """Return an edge-coloured second-order native fermionic gate stream.""" - t = self.t if t is None else t - U = self.U if U is None else U - V = self.V if V is None else V - mu = self.mu if mu is None else mu + """Return an edge-coloured second-order stream with explicit couplings.""" + if t is None: + raise TypeError("strang_gate_stream requires explicit t=... .") + if self.spinful and U is None: + raise TypeError( + "strang_gate_stream requires explicit U=... for spinful fermions." + ) edges = _as_edges(edges) sites = _sites_from_edges(edges, sites) half_dt = dt / 2 @@ -9109,23 +9100,23 @@ def hamiltonian( *, t=None, U=None, - V=None, - mu=None, + V=0.0, + mu=0.0, flat=False, to_backend=None, ): - """Validate explicit terms or build the configured lattice model. + """Validate explicit terms or build a model only from explicit couplings. The canonical form is a mapping from one-site or two-site locations to native fermionic Symmray arrays. It is checked for symmetry, physical sectors, support rank, and backend consistency before being bundled in a :class:`SymHamiltonian`. Passing lattice edges remains a compact - convenience, using this helper's model parameters unless explicit - coupling overrides are supplied. + convenience, but requires its couplings explicitly; no coupling is + stored on :class:`Fermion`. """ to_backend = self.to_backend if to_backend is None else to_backend if isinstance(terms_or_edges, Mapping): - if any(value is not None for value in (t, U, V, mu)): + if any(value is not None for value in (t, U)) or V != 0 or mu != 0: raise TypeError( "When passing explicit terms, put every coupling in the " "native arrays rather than passing t/U/V/mu again." @@ -9139,10 +9130,12 @@ def hamiltonian( parameters={}, ) - t = self.t if t is None else t - U = self.U if U is None else U - V = self.V if V is None else V - mu = self.mu if mu is None else mu + if t is None: + raise TypeError("hamiltonian(edges, ...) requires explicit t=... .") + if self.spinful and U is None: + raise TypeError( + "hamiltonian(edges, ...) requires explicit U=... for spinful fermions." + ) params = {"t": t, "V": V, "mu": mu} if self.spinful: params["U"] = U diff --git a/src/pepsy/vmc/__init__.py b/src/pepsy/vmc/__init__.py index c509403..f0a1a45 100644 --- a/src/pepsy/vmc/__init__.py +++ b/src/pepsy/vmc/__init__.py @@ -28,6 +28,9 @@ "NetKetChunkSettings": ".netket", "NetKetBuildTiming": ".netket", "NetKetAmplitudeTiming": ".netket", + "NetKetGPUUsage": ".netket", + "NetKetResourceUsage": ".netket", + "NetKetEtaPairObservable": ".netket", "NetKetPEPSVMC": ".netket", "NetKetVMCSetup": ".netket", "NetKetVMCConfig": ".netket", @@ -38,6 +41,8 @@ "PackedFermionicPEPS": ".netket", "SpinOrbitalColumns": ".netket", "TorchConnections": ".torch", + "TorchAmplitudeBenchmark": ".torch", + "TorchAmplitudeBenchmarkRun": ".torch", "TorchFermionVMC": ".torch", "TorchFermionVMCMetadata": ".torch", "TorchChainDiagnostics": ".torch", @@ -45,6 +50,7 @@ "TorchVMCConvergenceReport": ".torch", "TorchImportanceSamples": ".torch", "TorchMCMCSamples": ".torch", + "TorchDistributedMetadata": ".torch", "TorchSampleProvenance": ".torch", "TorchMetropolisResult": ".torch", "TorchMetropolisSampler": ".torch", @@ -67,6 +73,7 @@ "VMCSamples": ".api", "VMCWarning": ".api", "apply_torch_sr_update": ".torch", + "benchmark_torch_amplitudes": ".torch", "build_heisenberg_vmc": ".netket", "build_ising_vmc": ".netket", "build_fermi_hubbard_vmc": ".netket", diff --git a/src/pepsy/vmc/api.py b/src/pepsy/vmc/api.py index 497fb20..32830fc 100644 --- a/src/pepsy/vmc/api.py +++ b/src/pepsy/vmc/api.py @@ -85,6 +85,9 @@ class NumericalStabilityWarning(VMCWarning): } +_SAMPLING_UNSET = object() + + def _positive_int(name, value, *, allow_none=False): if value is None and allow_none: return None @@ -147,49 +150,130 @@ def _resolve_contraction_config(contraction, chi=None, cutoff=None, options=None ) -@dataclass(frozen=True) +@dataclass(frozen=True, init=False) class SamplingConfig: """Shared chain-preserving sampling settings. - ``burn_in`` is the number of discarded *thinning intervals* per chain. - Thus the native Torch sampler advances each chain - ``(burn_in + n_samples_per_chain) * thin`` Metropolis sweeps: it discards - ``burn_in * thin`` sweeps, then retains one configuration after every - ``thin`` further sweeps. The returned batch has shape - ``(n_samples_per_chain, n_chains, n_sites)``. + The canonical constructor follows NetKet:: + + ``n_samples`` is the total requested across chains, + ``n_discard_per_chain`` is the discarded prefix of every chain, and + ``sweep_size`` is the number of Metropolis sweeps between retained + configurations. The total must be divisible by ``n_chains`` so both + backends can return a fixed chain-shaped batch. + + The former ``n_samples_per_chain``, ``burn_in``, and ``thin`` keywords are + retained as compatibility aliases. They are also exposed as properties + for callers that still use the original Pepsy spelling. """ - n_samples_per_chain: int = 128 - n_chains: int = 16 - burn_in: int = 0 - thin: int = 1 - seed: int | None = None - sampler_seed: int | None = None - chunk_size: int | None = None - proposal: str | None = None + n_samples_per_chain: int + n_chains: int + burn_in: int + thin: int + seed: int | None + sampler_seed: int | None + chunk_size: int | None + proposal: str | None - def __post_init__(self): - n_samples = _positive_int("n_samples_per_chain", self.n_samples_per_chain) - n_chains = _positive_int("n_chains", self.n_chains) - if isinstance(self.burn_in, bool) or not isinstance(self.burn_in, int) or self.burn_in < 0: - raise ValueError("burn_in must be a non-negative integer.") - thin = _positive_int("thin", self.thin) - chunk_size = _positive_int("chunk_size", self.chunk_size, allow_none=True) - if self.seed is not None and isinstance(self.seed, bool): + def __init__( + self, + n_samples_per_chain=_SAMPLING_UNSET, + n_chains=16, + burn_in=_SAMPLING_UNSET, + thin=_SAMPLING_UNSET, + seed=None, + sampler_seed=None, + chunk_size=None, + proposal=None, + *, + n_samples=_SAMPLING_UNSET, + n_discard_per_chain=_SAMPLING_UNSET, + sweep_size=_SAMPLING_UNSET, + ): + if ( + n_samples is not _SAMPLING_UNSET + and n_samples_per_chain is not _SAMPLING_UNSET + ): + raise ValueError( + "Pass either n_samples or n_samples_per_chain, not both." + ) + if ( + n_discard_per_chain is not _SAMPLING_UNSET + and burn_in is not _SAMPLING_UNSET + ): + raise ValueError( + "Pass either n_discard_per_chain or burn_in, not both." + ) + if sweep_size is not _SAMPLING_UNSET and thin is not _SAMPLING_UNSET: + raise ValueError("Pass either sweep_size or thin, not both.") + + n_chains = _positive_int("n_chains", n_chains) + if n_samples is _SAMPLING_UNSET: + n_samples_per_chain = ( + 128 + if n_samples_per_chain is _SAMPLING_UNSET + else n_samples_per_chain + ) + n_samples_per_chain = _positive_int( + "n_samples_per_chain", + n_samples_per_chain, + ) + n_samples = n_samples_per_chain * n_chains + else: + n_samples = _positive_int("n_samples", n_samples) + if n_samples % n_chains: + raise ValueError( + "n_samples must be divisible by n_chains so every chain " + "has the same retained length." + ) + n_samples_per_chain = n_samples // n_chains + + if burn_in is _SAMPLING_UNSET: + burn_in = ( + 0 + if n_discard_per_chain is _SAMPLING_UNSET + else n_discard_per_chain + ) + if ( + isinstance(burn_in, bool) + or not isinstance(burn_in, int) + or burn_in < 0 + ): + raise ValueError("n_discard_per_chain must be a non-negative integer.") + if thin is _SAMPLING_UNSET: + thin = 1 if sweep_size is _SAMPLING_UNSET else sweep_size + thin = _positive_int("sweep_size", thin) + chunk_size = _positive_int("chunk_size", chunk_size, allow_none=True) + if seed is not None and isinstance(seed, bool): raise ValueError("seed must be an integer or None.") - if self.sampler_seed is not None and isinstance(self.sampler_seed, bool): + if sampler_seed is not None and isinstance(sampler_seed, bool): raise ValueError("sampler_seed must be an integer or None.") - object.__setattr__(self, "n_samples_per_chain", n_samples) + + object.__setattr__(self, "n_samples_per_chain", n_samples_per_chain) object.__setattr__(self, "n_chains", n_chains) - object.__setattr__(self, "burn_in", int(self.burn_in)) + object.__setattr__(self, "burn_in", int(burn_in)) object.__setattr__(self, "thin", thin) + object.__setattr__(self, "seed", seed) + object.__setattr__(self, "sampler_seed", sampler_seed) object.__setattr__(self, "chunk_size", chunk_size) + object.__setattr__(self, "proposal", proposal) @property def n_samples(self): """Total requested samples across all chains.""" return self.n_samples_per_chain * self.n_chains + @property + def n_discard_per_chain(self): + """Number of discarded samples/intervals per chain.""" + return self.burn_in + + @property + def sweep_size(self): + """Number of Metropolis sweeps between retained configurations.""" + return self.thin + def torch_kwargs(self): """Return the canonical keyword mapping for Torch samplers. @@ -204,7 +288,7 @@ def torch_kwargs(self): "n_samples": self.n_samples, "n_chains": self.n_chains, "n_discard_per_chain": self.burn_in, - "n_thin": self.thin, + "sweep_size": self.thin, "seed": self.seed, "sampler_seed": self.sampler_seed, } @@ -594,8 +678,8 @@ class MCState: ``MCState`` owns the variational ansatz and the sampling specification; :class:`VMC` attaches a Hamiltonian and builds the selected backend. The familiar ``n_samples`` value is the total over all chains, exactly as in - NetKet. ``sampling=SamplingConfig(...)`` remains available when the - per-chain convention is more convenient. + NetKet. ``SamplingConfig`` accepts the same total-sample spelling, while + retaining its older per-chain aliases for compatibility. """ peps: Any @@ -613,6 +697,7 @@ def __init__( n_chains=None, n_discard_per_chain=None, thin=None, + sweep_size=None, seed=None, sampler_seed=None, chunk_size=None, @@ -629,6 +714,7 @@ def __init__( "n_chains": n_chains, "n_discard_per_chain": n_discard_per_chain, "thin": thin, + "sweep_size": sweep_size, "seed": seed, "sampler_seed": sampler_seed, "chunk_size": chunk_size, @@ -650,11 +736,19 @@ def __init__( "n_samples must be divisible by n_chains so every chain " "has the same retained length." ) + if thin is not None and sweep_size is not None: + raise ValueError("Pass either thin or sweep_size, not both.") sampling = SamplingConfig( - n_samples_per_chain=n_samples // n_chains, + n_samples=n_samples, n_chains=n_chains, - burn_in=0 if n_discard_per_chain is None else n_discard_per_chain, - thin=1 if thin is None else thin, + n_discard_per_chain=( + 0 if n_discard_per_chain is None else n_discard_per_chain + ), + sweep_size=( + 1 + if thin is None and sweep_size is None + else thin if sweep_size is None else sweep_size + ), seed=seed, sampler_seed=sampler_seed, chunk_size=chunk_size, @@ -693,7 +787,12 @@ def n_chains(self): @property def n_discard_per_chain(self): """Per-chain burn-in, using NetKet's naming convention.""" - return self.sampling.burn_in + return self.sampling.n_discard_per_chain + + @property + def sweep_size(self): + """Metropolis sweeps between retained samples.""" + return self.sampling.sweep_size def to_problem(self, hamiltonian, *, observables=None): """Make the compatibility :class:`VMCProblem` representation.""" diff --git a/src/pepsy/vmc/netket.py b/src/pepsy/vmc/netket.py index b17d168..baf31fb 100644 --- a/src/pepsy/vmc/netket.py +++ b/src/pepsy/vmc/netket.py @@ -6,9 +6,11 @@ from __future__ import annotations +from collections.abc import Mapping from dataclasses import dataclass from numbers import Integral import os +import time from typing import Any import warnings @@ -18,6 +20,12 @@ __all__ = [ "NetKetLocalConfigMap", "NetKetChunkSettings", + "NetKetBuildTiming", + "NetKetAmplitudeTiming", + "NetKetGPUUsage", + "NetKetResourceUsage", + "NetKetVMCConfig", + "NetKetEtaPairObservable", "NetKetPEPSVMC", "NetKetFermiHubbardVMC", "NetKetVMCSetup", @@ -173,6 +181,176 @@ class NetKetChunkSettings: chunk_size_bwd: int | None +@dataclass(frozen=True) +class NetKetBuildTiming: + """Wall-clock breakdown for :func:`build_fermion_vmc`. + + The setup phase is separate from :meth:`NetKetPEPSVMC.warmup`: it creates + the Hilbert space, packs the PEPS, builds the JAX model, and constructs the + sampler/state. ``warmup`` then triggers the lazy sampler/amplitude/energy + compilations. Keeping these timings separate makes the first-run cost + visible instead of attributing it all to a VMC iteration. + """ + + settings_seconds: float + geometry_seconds: float + hamiltonian_seconds: float + peps_seconds: float + model_seconds: float + sampler_seconds: float + total_seconds: float + preconditioner_seconds: float = 0.0 + state_seconds: float = 0.0 + + def as_dict(self): + """Return a JSON-friendly phase-to-seconds mapping.""" + return { + "settings_seconds": float(self.settings_seconds), + "geometry_seconds": float(self.geometry_seconds), + "hamiltonian_seconds": float(self.hamiltonian_seconds), + "peps_seconds": float(self.peps_seconds), + "model_seconds": float(self.model_seconds), + "sampler_seconds": float(self.sampler_seconds), + "total_seconds": float(self.total_seconds), + "preconditioner_seconds": float(self.preconditioner_seconds), + "state_seconds": float(self.state_seconds), + } + + @property + def slowest_phase(self): + """Return ``(name, seconds)`` for the slowest setup phase.""" + phases = self.as_dict() + phases.pop("total_seconds") + name = max(phases, key=phases.get) + return name, phases[name] + + +@dataclass(frozen=True) +class NetKetAmplitudeTiming: + """Wall-clock timing for a synchronized PEPS amplitude batch.""" + + n_samples: int + amplitude_seconds: float + compile_seconds: float | None = None + + @property + def amplitude_seconds_per_sample(self): + """Return the measured average time per amplitude in the batch.""" + return float(self.amplitude_seconds) / int(self.n_samples) + + def as_dict(self): + """Return a JSON-friendly timing mapping.""" + return { + "n_samples": int(self.n_samples), + "amplitude_seconds": float(self.amplitude_seconds), + "amplitude_seconds_per_sample": float( + self.amplitude_seconds_per_sample + ), + "compile_seconds": ( + None + if self.compile_seconds is None + else float(self.compile_seconds) + ), + } + + +@dataclass(frozen=True) +class NetKetGPUUsage: + """One ``nvidia-smi`` snapshot for a GPU visible to the current process.""" + + index: int + name: str + memory_used_mib: int | None + memory_total_mib: int | None + utilization_percent: int | None + memory_utilization_percent: int | None + process_memory_mib: int | None = None + + def as_dict(self): + """Return a JSON-friendly representation of this GPU snapshot.""" + return { + "index": self.index, + "name": self.name, + "memory_used_mib": self.memory_used_mib, + "memory_total_mib": self.memory_total_mib, + "utilization_percent": self.utilization_percent, + "memory_utilization_percent": self.memory_utilization_percent, + "process_memory_mib": self.process_memory_mib, + } + + +@dataclass(frozen=True) +class NetKetResourceUsage: + """Host and GPU resource snapshots collected around one VMC operation. + + ``gpu_peak`` is sampled while the operation runs. GPU utilization is an + instantaneous NVML/``nvidia-smi`` quantity, so its peak is useful for + observing activity, but is not a time-averaged utilization percentage. + ``process_memory_mib`` distinguishes this notebook kernel's allocation + from memory consumed by other processes sharing the device. + """ + + elapsed_seconds: float + host_rss_before_mib: float | None + host_rss_after_mib: float | None + host_rss_peak_mib: float | None + gpu_before: tuple[NetKetGPUUsage, ...] = () + gpu_after: tuple[NetKetGPUUsage, ...] = () + gpu_peak: tuple[NetKetGPUUsage, ...] = () + host_monitor: str | None = None + + def as_dict(self): + """Return a JSON-friendly representation suitable for sample metadata.""" + return { + "elapsed_seconds": float(self.elapsed_seconds), + "host_rss_before_mib": self.host_rss_before_mib, + "host_rss_after_mib": self.host_rss_after_mib, + "host_rss_peak_mib": self.host_rss_peak_mib, + "gpu_before": tuple(item.as_dict() for item in self.gpu_before), + "gpu_after": tuple(item.as_dict() for item in self.gpu_after), + "gpu_peak": tuple(item.as_dict() for item in self.gpu_peak), + "host_monitor": self.host_monitor, + } + + def summary(self, label="VMC resources"): + """Format a compact human-readable summary for notebooks and logs.""" + def gib(value): + return "n/a" if value is None else f"{value / 1024:.2f} GiB" + + parts = [ + f"{label}: {self.elapsed_seconds:.1f}s", + "host RSS " + f"{gib(self.host_rss_before_mib)} -> {gib(self.host_rss_after_mib)} " + f"(peak {gib(self.host_rss_peak_mib)})", + ] + after = {item.index: item for item in self.gpu_after} + before = {item.index: item for item in self.gpu_before} + peak = {item.index: item for item in self.gpu_peak} + for index in sorted(set(before) | set(after) | set(peak)): + current = after.get(index) or before.get(index) or peak[index] + highest = peak.get(index, current) + process_memory = current.process_memory_mib + process_peak = highest.process_memory_mib + current_utilization = current.utilization_percent + peak_utilization = highest.utilization_percent + device_memory = ( + "n/a" + if current.memory_used_mib is None + or current.memory_total_mib is None + else f"{current.memory_used_mib / 1024:.2f}/" + f"{current.memory_total_mib / 1024:.2f} GiB" + ) + parts.append( + f"GPU {index} ({current.name}): process {gib(process_memory)} " + f"(peak {gib(process_peak)}), device {device_memory}, " + f"util {current_utilization if current_utilization is not None else 'n/a'}% " + f"(peak {peak_utilization if peak_utilization is not None else 'n/a'}%)" + ) + if not after and not before and not peak: + parts.append("GPU metrics unavailable (nvidia-smi not visible)") + return "; ".join(parts) + + @dataclass(frozen=True) class NetKetVMCSettings: """Conservative large-run NetKet settings suggested by Pepsy.""" @@ -189,6 +367,80 @@ class NetKetVMCSettings: notes: tuple[str, ...] +@dataclass(frozen=True) +class NetKetVMCConfig: + """Validated construction/runtime settings for :func:`build_fermion_vmc`. + + This bundle keeps physics inputs (PEPS, ``Fermion``, and Hamiltonian) + separate from numerical settings without removing the builder's legacy + keyword arguments. ``sampling`` controls the MCState defaults; call + ``setup.sample(...)`` later to request a different retained batch. + """ + + contraction: Any = "exact" + sampling: Any = None + sampler_sweep_size: int | None = None + conserving: bool | str = False + use_sr: bool | str = False + param_dtype: Any | None = None + verify_columns: bool = False + progress: bool = False + + def __post_init__(self): + from .api import SamplingConfig + + if self.sampling is not None and not isinstance(self.sampling, SamplingConfig): + raise TypeError("sampling must be a SamplingConfig or None.") + if self.sampler_sweep_size is not None: + _check_positive_int("sampler_sweep_size", self.sampler_sweep_size) + if self.conserving not in {False, True, "auto"}: + raise ValueError("conserving must be False, True, or 'auto'.") + if self.use_sr not in {False, True, "auto"}: + raise ValueError("use_sr must be False, True, or 'auto'.") + if not isinstance(self.verify_columns, bool): + raise TypeError("verify_columns must be a bool.") + if not isinstance(self.progress, bool): + raise TypeError("progress must be a bool.") + + +@dataclass(frozen=True) +class NetKetEtaPairObservable: + r"""Declarative eta-pair observable resolved by NetKet measurement. + + This represents + + .. math:: + + P_\eta(dx, dy) = \frac{1}{N}\sum_i + \left(\Delta_i^\dagger \Delta_{i + (dx, dy)} + \mathrm{h.c.}\right). + + Pass an instance as a value in the ``observables`` mapping given to + :meth:`NetKetPEPSVMC.measure_samples`. The operator is compiled there + from the packed PEPS lattice order, so it shares the retained NetKet + configurations with every other requested observable. At zero + displacement it instead measures the mean double occupancy. + + ``periodic=False`` keeps only in-bounds pairs and normalizes by their + number. ``staggered=True`` multiplies each nonzero-displacement pair by + ``(-1)**(x_i + y_i + x_j + y_j)``. + """ + + dx: int + dy: int + periodic: bool = True + staggered: bool = False + + def __post_init__(self): + for name in ("dx", "dy"): + value = getattr(self, name) + if isinstance(value, bool) or not isinstance(value, Integral): + raise TypeError(f"{name} must be an integer.") + object.__setattr__(self, name, int(value)) + for name in ("periodic", "staggered"): + if not isinstance(getattr(self, name), bool): + raise TypeError(f"{name} must be a bool.") + + def _make_progress_bar(*, total=None, desc=None, enabled=True): """Return a notebook-friendly ``tqdm`` bar, or ``None`` when unavailable.""" if not enabled: @@ -200,6 +452,326 @@ def _make_progress_bar(*, total=None, desc=None, enabled=True): return tqdm(total=total, desc=desc, leave=True, dynamic_ncols=True) +def _block_until_ready(value): + """Synchronize a JAX value or pytree without importing JAX eagerly.""" + wait = getattr(value, "block_until_ready", None) + if callable(wait): + wait() + return + if isinstance(value, Mapping): + for item in value.values(): + _block_until_ready(item) + return + if isinstance(value, (tuple, list)): + for item in value: + _block_until_ready(item) + return + # ``np.asarray`` synchronizes JAX arrays and remains harmless for the + # small NumPy/scalar stand-ins used in optional-dependency tests. + _ = np.asarray(value) + + +def _tree_isfinite(value): + """Return whether every numerical leaf in a JAX/Flax pytree is finite.""" + if isinstance(value, Mapping): + return all(_tree_isfinite(item) for item in value.values()) + if isinstance(value, (tuple, list)): + return all(_tree_isfinite(item) for item in value) + array = np.asarray(value) + try: + return bool(np.all(np.isfinite(array))) + except TypeError: + # Test doubles and framework metadata are occasionally carried next + # to numerical leaves; only numerical parameter arrays need checking. + return True + + +def _process_rss_mib(): + """Return this Python process's resident memory, if ``psutil`` is present.""" + try: + import psutil + + return psutil.Process().memory_info().rss / (1 << 20) + except Exception: + return None + + +def _nvidia_smi_integer(value): + """Parse a possibly unavailable ``nvidia-smi`` integer field.""" + value = value.strip() + if value.lower() in {"", "n/a", "[not supported]"}: + return None + try: + return int(value) + except ValueError: + return None + + +def _nvidia_smi_gpu_usage(): + """Return per-GPU memory/utilization snapshots without a Python NVML dep.""" + import subprocess + + try: + gpu_process = subprocess.run( + [ + "nvidia-smi", + "--query-gpu=index,uuid,name,memory.used,memory.total," + "utilization.gpu,utilization.memory", + "--format=csv,noheader,nounits", + ], + check=False, + capture_output=True, + text=True, + timeout=2.0, + ) + except (OSError, subprocess.SubprocessError): + return () + if gpu_process.returncode or not gpu_process.stdout.strip(): + return () + + process_memory_by_uuid = {} + try: + applications = subprocess.run( + [ + "nvidia-smi", + "--query-compute-apps=pid,gpu_uuid,used_memory", + "--format=csv,noheader,nounits", + ], + check=False, + capture_output=True, + text=True, + timeout=2.0, + ) + if applications.returncode == 0: + for line in applications.stdout.splitlines(): + fields = [field.strip() for field in line.split(",")] + if len(fields) != 3 or _nvidia_smi_integer(fields[0]) != os.getpid(): + continue + process_memory_by_uuid[fields[1]] = _nvidia_smi_integer(fields[2]) + except (OSError, subprocess.SubprocessError): + pass + + result = [] + for line in gpu_process.stdout.splitlines(): + fields = [field.strip() for field in line.split(",")] + if len(fields) != 7: + continue + index = _nvidia_smi_integer(fields[0]) + if index is None: + continue + result.append( + NetKetGPUUsage( + index=index, + name=fields[2], + memory_used_mib=_nvidia_smi_integer(fields[3]), + memory_total_mib=_nvidia_smi_integer(fields[4]), + utilization_percent=_nvidia_smi_integer(fields[5]), + memory_utilization_percent=_nvidia_smi_integer(fields[6]), + process_memory_mib=process_memory_by_uuid.get(fields[1]), + ) + ) + return tuple(result) + + +class _NetKetResourceMonitor: + """Sample host RSS and ``nvidia-smi`` while one blocking VMC call runs.""" + + def __init__(self, *, interval=0.25): + interval = float(interval) + if interval <= 0: + raise ValueError("resource_interval must be positive.") + self.interval = interval + self._started = None + self._host_before = None + self._gpu_before = () + self._gpu_peak = {} + self._host_monitor = None + self._host_monitor_name = None + self._stop_event = None + self._thread = None + + def _record_gpu_usage(self, usage): + for item in usage: + previous = self._gpu_peak.get(item.index) + if previous is None: + self._gpu_peak[item.index] = item + continue + self._gpu_peak[item.index] = NetKetGPUUsage( + index=item.index, + name=item.name, + memory_used_mib=max( + value + for value in (previous.memory_used_mib, item.memory_used_mib) + if value is not None + ) + if any( + value is not None + for value in (previous.memory_used_mib, item.memory_used_mib) + ) + else None, + memory_total_mib=item.memory_total_mib + or previous.memory_total_mib, + utilization_percent=max( + value + for value in (previous.utilization_percent, item.utilization_percent) + if value is not None + ) + if any( + value is not None + for value in ( + previous.utilization_percent, + item.utilization_percent, + ) + ) + else None, + memory_utilization_percent=max( + value + for value in ( + previous.memory_utilization_percent, + item.memory_utilization_percent, + ) + if value is not None + ) + if any( + value is not None + for value in ( + previous.memory_utilization_percent, + item.memory_utilization_percent, + ) + ) + else None, + process_memory_mib=max( + value + for value in ( + previous.process_memory_mib, + item.process_memory_mib, + ) + if value is not None + ) + if any( + value is not None + for value in ( + previous.process_memory_mib, + item.process_memory_mib, + ) + ) + else None, + ) + + def _poll(self): + while not self._stop_event.wait(self.interval): + self._record_gpu_usage(_nvidia_smi_gpu_usage()) + + def start(self): + """Start non-intrusive resource sampling.""" + import threading + + self._host_before = _process_rss_mib() + self._gpu_before = _nvidia_smi_gpu_usage() + self._record_gpu_usage(self._gpu_before) + # xyzpy already supplies a low-overhead RSS peak sampler. It remains + # optional so the NetKet bridge has no new hard runtime dependency. + try: + import xyzpy + + # Keep host RSS sampling responsive at teardown; GPU polling still + # follows ``resource_interval`` and is the expensive operation. + self._host_monitor = xyzpy.MemoryMonitor( + interval=min(self.interval, 0.05) + ) + self._host_monitor.start() + self._host_monitor_name = "xyzpy.MemoryMonitor" + except Exception: + self._host_monitor = None + # Exclude snapshot/monitor setup from the reported operation time. + self._started = time.perf_counter() + if self._gpu_before: + self._stop_event = threading.Event() + self._thread = threading.Thread(target=self._poll, daemon=True) + self._thread.start() + return self + + def stop(self): + """Stop sampling and return the collected resource report.""" + if self._started is None: + raise RuntimeError("Resource monitor has not been started.") + ended = time.perf_counter() + if self._stop_event is not None: + self._stop_event.set() + if self._thread is not None: + self._thread.join(timeout=3.0) + gpu_after = _nvidia_smi_gpu_usage() + self._record_gpu_usage(gpu_after) + host_after = _process_rss_mib() + host_peak = max( + value + for value in (self._host_before, host_after) + if value is not None + ) if any(value is not None for value in (self._host_before, host_after)) else None + if self._host_monitor is not None: + self._host_monitor.stop() + if self._host_monitor.peak is not None: + host_peak = max(host_peak or 0.0, self._host_monitor.peak * 1024) + return NetKetResourceUsage( + elapsed_seconds=ended - self._started, + host_rss_before_mib=self._host_before, + host_rss_after_mib=host_after, + host_rss_peak_mib=host_peak, + gpu_before=self._gpu_before, + gpu_after=gpu_after, + gpu_peak=tuple( + self._gpu_peak[index] for index in sorted(self._gpu_peak) + ), + host_monitor=self._host_monitor_name, + ) + + +def _netket_sampling_diagnostics(vstate, sampler): + """Read portable sampling metadata from a NetKet MCState.""" + state = getattr(vstate, "sampler_state", None) + + def scalar(value): + if value is None: + return None + try: + return float(np.asarray(value)) + except (TypeError, ValueError): + return None + + acceptance = scalar(getattr(state, "acceptance", None)) + n_steps = scalar(getattr(state, "n_steps", None)) + n_accepted = scalar(getattr(state, "n_accepted", None)) + return { + "n_samples": int(getattr(vstate, "n_samples", 0)), + "n_chains": int(getattr(sampler, "n_chains", 0)), + "burn_in": int(getattr(vstate, "n_discard_per_chain", 0) or 0), + "sweep_size": getattr(sampler, "sweep_size", None), + "chunk_size": getattr(vstate, "chunk_size", None), + "acceptance_rate": acceptance, + "n_steps": None if n_steps is None else int(n_steps), + "n_accepted": None if n_accepted is None else int(n_accepted), + } + + +def _format_sampling_postfix(diagnostics): + """Format compact chain diagnostics for a progress-bar postfix.""" + parts = [] + acceptance = diagnostics.get("acceptance_rate") + if acceptance is not None: + parts.append(f"acc={acceptance:.2f}") + n_samples = diagnostics.get("n_samples") + n_chains = diagnostics.get("n_chains") + if n_samples and n_chains: + parts.append(f"samples={n_samples}/{n_chains}ch") + burn_in = diagnostics.get("burn_in") + sweep_size = diagnostics.get("sweep_size") + if burn_in is not None: + parts.append(f"burn={burn_in}") + if sweep_size is not None: + parts.append(f"sweep={sweep_size}") + return " | ".join(parts) + + @dataclass(frozen=True) class VMCOptimizeResult: """Energy-optimization history returned by :meth:`NetKetPEPSVMC.optimize`. @@ -208,7 +780,9 @@ class VMCOptimizeResult: (real energy mean, error of the mean, and sample variance). ``energy_shift`` is added to every energy by :attr:`shifted_energies` and :meth:`plot` so a convention offset (for example ``-U/4`` for Fermi-Hubbard) can be applied - without mutating the raw samples. + without mutating the raw samples. ``compile_seconds`` measures the staged + sampler/amplitude/energy warmup; ``optimization_seconds`` and + ``total_seconds`` are wall-clock measurements for the run. """ steps: Any @@ -219,6 +793,20 @@ class VMCOptimizeResult: final_error: float compile_seconds: float | None = None energy_shift: float = 0.0 + optimization_seconds: float | None = None + total_seconds: float | None = None + + @property + def warmup_seconds(self): + """Alias for the one-time sampler/JAX compilation duration.""" + return self.compile_seconds + + @property + def optimization_seconds_per_step(self): + """Return the measured optimization wall time per recorded step.""" + if self.optimization_seconds is None or len(np.asarray(self.steps)) == 0: + return None + return float(self.optimization_seconds) / len(np.asarray(self.steps)) @property def shifted_energies(self): @@ -263,6 +851,7 @@ def __init__(self, n_iter, *, enabled=True, energy_shift=0.0, per_site=None): self.energy_shift = float(energy_shift) self.per_site = per_site self._bar = None + self._first_update_started = None self.steps = [] self.energies = [] self.errors = [] @@ -275,7 +864,44 @@ def _ensure_bar(self): ) return self._bar + def start(self, status="first update: sampling and compiling gradients"): + """Show the VMC bar before NetKet's first (potentially JIT-heavy) step.""" + self._first_update_started = time.perf_counter() + bar = self._ensure_bar() + if bar is not None: + bar.set_description_str(f"VMC 0/{self.n_iter}: preparing") + bar.set_postfix_str(str(status)) + + def set_status(self, status): + """Update the pre-first-step status without advancing the bar.""" + bar = self._ensure_bar() + if bar is not None and not self.steps: + bar.set_postfix_str(str(status)) + def __call__(self, step, log_data, driver): + # NetKet invokes legacy callbacks after ``update_parameters``, but JAX + # may still have the backward pass and optimizer update queued on the + # device. Synchronize the updated parameter pytree before advancing + # tqdm: otherwise the first energy statistic can tick while the GPU is + # still working on that update, and the apparent stall moves to the + # following bar item. VMC updates are parameter-dependent, so this + # does not remove useful inter-step GPU parallelism; it makes elapsed + # time and ETA honest. + vstate = getattr(driver, "variational_state", None) + if vstate is None: + vstate = getattr(driver, "_variational_state", None) + if vstate is not None: + parameters = getattr(vstate, "parameters", None) + if parameters is not None: + _block_until_ready(parameters) + if parameters is not None and not _tree_isfinite(parameters): + raise FloatingPointError( + "NetKet VMC produced non-finite PEPS parameters after " + f"update {int(step)}. Reduce the learning rate, restart " + "from the pre-update state, and verify a short run before " + "using a larger sample count." + ) + name = getattr(driver, "_loss_name", "Energy") stats = None if isinstance(log_data, dict): @@ -286,6 +912,13 @@ def __call__(self, step, log_data, driver): mean = float(np.real(np.asarray(getattr(stats, "mean", stats)))) err = float(getattr(stats, "error_of_mean", np.nan)) var = float(getattr(stats, "variance", np.nan)) + if not np.isfinite(mean): + raise FloatingPointError( + "NetKet VMC produced a non-finite local-energy estimate " + f"at update {int(step)}. The PEPS/sampler state is no " + "longer numerically usable; reduce the learning rate and " + "restart from the initial PEPS." + ) self.steps.append(int(step)) self.energies.append(mean) self.errors.append(err) @@ -294,7 +927,20 @@ def __call__(self, step, log_data, driver): if bar is not None: scale = 1.0 if not self.per_site else float(self.per_site) shown = (mean + self.energy_shift) / scale - bar.set_postfix_str(f"E={shown:.6f}\u00b1{err / scale:.1e}") + details = [f"E={shown:.6f}\u00b1{err / scale:.1e}"] + if self._first_update_started is not None: + first_seconds = time.perf_counter() - self._first_update_started + details.insert(0, f"first update {first_seconds:.1f}s") + self._first_update_started = None + if vstate is not None: + sampling = _netket_sampling_diagnostics( + vstate, getattr(vstate, "sampler", None) + ) + sampling_text = _format_sampling_postfix(sampling) + if sampling_text: + details.append(sampling_text) + bar.set_description_str("VMC energy") + bar.set_postfix_str(" | ".join(details)) bar.update(1) return True @@ -303,7 +949,12 @@ def close(self): self._bar.close() self._bar = None - def result(self, compile_seconds=None): + def result( + self, + compile_seconds=None, + optimization_seconds=None, + total_seconds=None, + ): energies = np.asarray(self.energies, dtype=float) errors = np.asarray(self.errors, dtype=float) return VMCOptimizeResult( @@ -314,6 +965,8 @@ def result(self, compile_seconds=None): final_energy=float(energies[-1]) if energies.size else float("nan"), final_error=float(errors[-1]) if errors.size else float("nan"), compile_seconds=compile_seconds, + optimization_seconds=optimization_seconds, + total_seconds=total_seconds, energy_shift=self.energy_shift, ) @@ -331,6 +984,7 @@ class NetKetPEPSVMC: ansatz: PackedPEPS config_map: NetKetLocalConfigMap | None preconditioner: Any | None + build_timing: NetKetBuildTiming | None = None @property def n_sites(self): @@ -346,6 +1000,43 @@ def expect_energy(self): """Return NetKet's expectation estimate for this setup Hamiltonian.""" return self.vstate.expect(self.hamiltonian) + def _mc_diagnostic_method(self, name): + """Return a recent NetKet MC diagnostic or raise a clear version error.""" + method = getattr(self.vstate, name, None) + if callable(method): + return method + raise RuntimeError( + f"NetKet MC diagnostic {name}() is unavailable in this NetKet " + "version. Install NetKet >= 3.22 to use it." + ) + + def check_mc_convergence(self, hamiltonian=None, **kwargs): + """Diagnose final-state mixing without mutating the VMC state. + + This delegates to NetKet's ``MCState.check_mc_convergence``. It is + intentionally an explicit post-optimization diagnostic because it + draws long chains to estimate :math:`\\hat R` and autocorrelation time. + """ + if hamiltonian is None: + hamiltonian = self.hamiltonian + return self._mc_diagnostic_method("check_mc_convergence")( + hamiltonian, **kwargs + ) + + def thermalise(self, hamiltonian=None, **kwargs): + """Advance chains in place until NetKet's mixing criterion is met.""" + if hamiltonian is None: + hamiltonian = self.hamiltonian + return self._mc_diagnostic_method("thermalise")(hamiltonian, **kwargs) + + def expect_to_precision(self, observable=None, **kwargs): + """Sample an observable until NetKet reaches a requested tolerance.""" + if observable is None: + observable = self.hamiltonian + return self._mc_diagnostic_method("expect_to_precision")( + observable, **kwargs + ) + def make_driver(self, **kwargs): """Create a NetKet VMC driver for this setup. @@ -353,15 +1044,116 @@ def make_driver(self, **kwargs): """ return make_netket_vmc_driver(self, **kwargs) - def warmup(self, *, progress=True, hamiltonian=None): + def warmup( + self, + *, + progress=True, + hamiltonian=None, + verbose=None, + resource_monitor=False, + resource_interval=0.25, + ): """Compile the VMC kernels up front with a small staged progress bar. Thin wrapper over :func:`warmup_netket_vmc`; returns the elapsed compile seconds so the following optimization ETA is meaningful. """ - return warmup_netket_vmc(self, hamiltonian=hamiltonian, progress=progress) + return warmup_netket_vmc( + self, + hamiltonian=hamiltonian, + progress=progress, + verbose=verbose, + resource_monitor=resource_monitor, + resource_interval=resource_interval, + ) + + def to_peps(self, variables=None, *, device_get=True): + """Return the current NetKet parameters as a quimb PEPS-like network. + + With no explicit ``variables``, this reconstructs the PEPS from the + setup's current ``MCState`` parameters, including parameters updated by + a NetKet VMC driver. Pass either the full Flax variables mapping or + its ``"params"`` collection to inspect a checkpoint without mutating + the VMC setup. By default leaves are copied from JAX devices to NumPy, + yielding a regular quimb/Symmray network suitable for Pepsy methods; + set ``device_get=False`` to retain JAX-backed leaves. + + The result is the underlying quimb network (the ``.tn`` of a + :class:`pepsy.SymPEPS`), preserving the packed skeleton's tensor + topology, indices, and fermionic metadata. + """ + if variables is None: + variables = getattr(self.vstate, "variables", None) + if variables is None: + variables = {"params": self.vstate.parameters} + params = _peps_params_from_netket_variables( + self.ansatz, + variables, + device_get=device_get, + ) + return qtn.unpack(params, self.ansatz.skeleton) + + def benchmark_amplitude(self, configs=None, *, n_samples=1): + """Time one synchronized jitted PEPS amplitude batch. + + If ``configs`` is omitted, the first ``n_samples`` configurations from + the current NetKet sample cache are used. Sampling itself is outside + the timed region, so the result isolates NetKet's jitted log-amplitude + evaluation. The first synchronized call is treated as a compile probe + and is reported separately in ``compile_seconds``. The returned + ``amplitude_seconds`` is a second call with the same batch shape, so + it measures steady-state evaluation rather than shape-specific JAX + compilation. Choose ``n_samples`` to match a real sampler or forward + chunk; timing an arbitrary small batch does not predict VMC throughput. + """ + n_samples = int(n_samples) + if n_samples <= 0: + raise ValueError("n_samples must be positive.") + if configs is None: + native = self.vstate.samples + configs = native.reshape((-1, native.shape[-1])) + else: + configs = np.asarray(configs) + if configs.ndim == 1: + configs = configs.reshape((1, -1)) + if configs.ndim != 2: + raise ValueError( + "configs must have shape (n_samples, n_orbitals)." + ) + if int(configs.shape[0]) < n_samples: + raise ValueError( + f"configs contains {int(configs.shape[0])} rows, " + f"but n_samples={n_samples} was requested." + ) + configs = configs[:n_samples] + log_value = getattr(self.vstate, "log_value", None) + if not callable(log_value): + raise TypeError( + "NetKet amplitude benchmarking requires an MCState with a " + "log_value(configs) method." + ) + compile_started = time.perf_counter() + log_amplitudes = log_value(configs) + _block_until_ready(log_amplitudes) + compile_seconds = time.perf_counter() - compile_started + started = time.perf_counter() + log_amplitudes = log_value(configs) + _block_until_ready(log_amplitudes) + return NetKetAmplitudeTiming( + n_samples=n_samples, + amplitude_seconds=time.perf_counter() - started, + compile_seconds=compile_seconds, + ) - def sample(self, sampling=None): + def sample( + self, + sampling=None, + *, + fresh=False, + progress=False, + resource_monitor=False, + resource_interval=0.25, + ): """Collect samples using the shared :class:`SamplingConfig` contract. NetKet stores samples as ``(n_chains, n_samples_per_chain, sites)``; @@ -369,55 +1161,103 @@ def sample(self, sampling=None): ``(n_samples_per_chain, n_chains, sites)`` layout used by Torch. The sampler's chain count is fixed when the setup is built, so a config requesting another count raises a clear error instead of silently - returning a different ensemble. + returning a different ensemble. With ``fresh=True``, reset NetKet's + retained-sample cache before reading it, forcing one new batch from the + current chain state. With ``progress=True``, the bar shows + retained samples, burn-in, sweep size, elapsed time, and NetKet's + acceptance rate after sampling. Set ``resource_monitor=True`` to also + print host RSS plus this kernel's GPU-memory and GPU-utilization + snapshots; the serializable report is retained in + ``VMCSamples.diagnostics["resources"]``. """ from .api import BackendCapabilityWarning, SamplingConfig, VMCSamples if sampling is not None and not isinstance(sampling, SamplingConfig): raise TypeError("sampling must be a SamplingConfig or None.") - if sampling is None: - native = self.vstate.samples - requested_chunk = None - else: - actual_chains = getattr(self.sampler, "n_chains", None) - if actual_chains is not None and int(actual_chains) != sampling.n_chains: - raise ValueError( - "sampling.n_chains does not match the setup sampler: " - f"expected {int(actual_chains)}, got {sampling.n_chains}. " - "Rebuild the setup with n_chains=... to change it." - ) - if sampling.thin != 1: - warnings.warn( - "NetKet MCState.sample has no per-sample thinning option; " - "sampling.thin is ignored.", - BackendCapabilityWarning, - stacklevel=2, - ) - if sampling.proposal is not None: - warnings.warn( - "sampling.proposal is selected when the NetKet sampler is " - "built and is ignored by MCState.sample.", - BackendCapabilityWarning, - stacklevel=2, - ) - if sampling.seed is not None or sampling.sampler_seed is not None: - warnings.warn( - "NetKet MCState.sample cannot reseed an existing sampler; " - "seed settings are ignored. Rebuild the setup to reseed.", - BackendCapabilityWarning, - stacklevel=2, - ) - kwargs = sampling.netket_kwargs() - kwargs.pop("n_chains", None) - requested_chunk = sampling.chunk_size - old_chunk = getattr(self.vstate, "chunk_size", None) - if requested_chunk is not None: - self.vstate.chunk_size = requested_chunk - try: - native = self.vstate.sample(**kwargs) - finally: + if not isinstance(fresh, bool): + raise TypeError("fresh must be a bool.") + bar = _make_progress_bar(total=1, desc="NetKet sampling", enabled=progress) + monitor = ( + _NetKetResourceMonitor(interval=resource_interval) + if resource_monitor + else None + ) + if monitor is not None: + monitor.start() + # Keep the sampling time independent of optional telemetry setup. + started = time.perf_counter() + try: + if fresh: + self.vstate.reset() + if sampling is None: + native = self.vstate.samples + else: + actual_chains = getattr(self.sampler, "n_chains", None) + if ( + actual_chains is not None + and int(actual_chains) != sampling.n_chains + ): + raise ValueError( + "sampling.n_chains does not match the setup sampler: " + f"expected {int(actual_chains)}, got {sampling.n_chains}. " + "Rebuild the setup with n_chains=... to change it." + ) + if sampling.sweep_size != 1: + warnings.warn( + "NetKet MCState.sample has no per-sample thinning option; " + "sampling.sweep_size is ignored.", + BackendCapabilityWarning, + stacklevel=2, + ) + if sampling.proposal is not None: + warnings.warn( + "sampling.proposal is selected when the NetKet sampler is " + "built and is ignored by MCState.sample.", + BackendCapabilityWarning, + stacklevel=2, + ) + if sampling.seed is not None or sampling.sampler_seed is not None: + warnings.warn( + "NetKet MCState.sample cannot reseed an existing sampler; " + "seed settings are ignored. Rebuild the setup to reseed.", + BackendCapabilityWarning, + stacklevel=2, + ) + kwargs = sampling.netket_kwargs() + kwargs.pop("n_chains", None) + requested_chunk = sampling.chunk_size + old_chunk = getattr(self.vstate, "chunk_size", None) if requested_chunk is not None: - self.vstate.chunk_size = old_chunk + self.vstate.chunk_size = requested_chunk + try: + native = self.vstate.sample(**kwargs) + finally: + if requested_chunk is not None: + self.vstate.chunk_size = old_chunk + finally: + resource_usage = monitor.stop() if monitor is not None else None + + diagnostics = _netket_sampling_diagnostics(self.vstate, self.sampler) + # ``MCState.sample`` accepts a temporary sample count/discard value, + # but does not consistently copy those values back onto the state. + # Keep the progress text and returned metadata faithful to the request + # rather than reporting the build-time defaults in that case. + if sampling is not None: + diagnostics["n_samples"] = sampling.n_samples + diagnostics["n_chains"] = sampling.n_chains + diagnostics["burn_in"] = sampling.n_discard_per_chain + if sampling.chunk_size is not None: + diagnostics["chunk_size"] = sampling.chunk_size + diagnostics["elapsed_seconds"] = time.perf_counter() - started + if resource_usage is not None: + diagnostics["resources"] = resource_usage.as_dict() + print(resource_usage.summary("NetKet sampling resources"), flush=True) + if bar is not None: + postfix = _format_sampling_postfix(diagnostics) + if postfix: + bar.set_postfix_str(postfix) + bar.update(1) + bar.close() shape = getattr(native, "shape", ()) if len(shape) == 3: @@ -432,11 +1272,14 @@ def sample(self, sampling=None): "NetKet samples must have shape (chains, samples, sites) or " f"(samples, sites), got {shape}." ) + diagnostics["n_samples_per_chain"] = n_samples_per_chain return VMCSamples( configs=configs, n_samples_per_chain=n_samples_per_chain, n_chains=n_chains, native=native, + acceptance_rate=diagnostics.get("acceptance_rate"), + diagnostics=diagnostics, ) def optimize( @@ -510,6 +1353,20 @@ def optimize( if energy_shift is None: energy_shift = 0.0 driver_options = {} if driver_options is None else dict(driver_options) + total_started = time.perf_counter() + cb = _VMCProgressCallback( + n_iter, + enabled=progress, + energy_shift=energy_shift, + per_site=per_site, + ) + callbacks = [cb] + if extra_callbacks: + callbacks.extend(extra_callbacks) + # When a separate warmup has already completed (the usual notebook + # path), even driver setup would otherwise be invisible. + if not warmup: + cb.start("building VMC driver") run_driver = make_netket_vmc_driver( self, optimizer=optimizer, @@ -520,15 +1377,10 @@ def optimize( compile_seconds = None if warmup: compile_seconds = self.warmup(progress=progress) - cb = _VMCProgressCallback( - n_iter, - enabled=progress, - energy_shift=energy_shift, - per_site=per_site, - ) - callbacks = [cb] - if extra_callbacks: - callbacks.extend(extra_callbacks) + # NetKet invokes callbacks only after its first update. Create this + # bar now so the JIT-heavy first gradient/sampling pass is visible. + cb.start() + optimization_started = time.perf_counter() try: run_driver.run( n_iter, @@ -538,16 +1390,92 @@ def optimize( ) finally: cb.close() - return cb.result(compile_seconds=compile_seconds) + optimization_seconds = time.perf_counter() - optimization_started + total_seconds = time.perf_counter() - total_started + return cb.result( + compile_seconds=compile_seconds, + optimization_seconds=optimization_seconds, + total_seconds=total_seconds, + ) - def measure(self, observables=None): - """Measure observables on the current variational state. + def run( + self, + n_iter, + *, + learning_rate=0.02, + driver="vmc", + optimizer=None, + warmup=True, + progress=True, + energy_shift=0.0, + per_site=None, + sr_mode="real", + sr_diag_shift=0.01, + use_sr=False, + driver_options=None, + **run_kwargs, + ): + """Run NetKet VMC with one timing-aware public entry point. + + This is the concise native setup API: contraction and sampler choices + belong to :func:`build_fermion_vmc`, while optimization choices are + passed directly here. Standard ``driver='vmc'`` uses no SR + preconditioner unless ``use_sr=True`` or an explicit driver option is + supplied. The returned :class:`VMCOptimizeResult` reports + one-time warmup/JAX compilation, optimization wall time, total wall + time, and the measured time per optimization step. + """ + options = {} if driver_options is None else dict(driver_options) + driver_name = str(driver).replace("-", "_").lower() + if driver_name in {"vmc_sr", "vmcsr", "sr"}: + options.setdefault("sr_mode", sr_mode) + options.setdefault("sr_diag_shift", sr_diag_shift) + elif "use_sr" not in options and "preconditioner" not in options: + options["use_sr"] = use_sr + return self.optimize( + n_iter, + learning_rate=learning_rate, + driver=driver, + optimizer=optimizer, + progress=progress, + warmup=warmup, + energy_shift=energy_shift, + per_site=per_site, + driver_options=options, + **run_kwargs, + ) - ``observables`` may be a single NetKet operator, a ``{name: operator}`` - mapping, or ``None`` to use any observables stored on the setup (for - example those passed to :func:`build_fermion_vmc`). Returns a single - ``nk.stats.Stats`` for one operator, otherwise a ``{name: Stats}`` dict. + def measure_samples(self, samples, observables=None): + """Measure observables from a retained NetKet sample batch. + + ``samples`` must be the :class:`~pepsy.vmc.VMCSamples` returned by + :meth:`sample` (or its unchanged native NetKet array). NetKet's local + estimators consume the ``MCState`` sample cache rather than accepting + configurations as an argument, so this method verifies that the batch + is still that cache. It then evaluates every observable on precisely + the same Markov-chain configurations, without drawing another batch. + + A fresh ``sample(...)`` call is required after a parameter update or + any other operation that invalidates NetKet's cache. ``observables`` + may be a single NetKet operator, a ``{name: operator}`` mapping, a + :class:`NetKetEtaPairObservable` specification, or ``None`` to use + those stored on the setup. """ + from .api import VMCBackendCapabilityError + + native_samples = getattr(samples, "native", samples) + cached_samples = getattr(self.vstate, "_samples", None) + if native_samples is None: + raise VMCBackendCapabilityError( + "NetKet measurement needs a native sample batch. Call " + "setup.sample(...) and pass its returned VMCSamples object." + ) + if cached_samples is None or native_samples is not cached_samples: + raise VMCBackendCapabilityError( + "NetKet can measure only the current MCState sample cache. " + "Call setup.sample(...) and pass that returned batch; external " + "or stale configurations cannot be installed safely." + ) if observables is None: observables = getattr(self, "observables", None) if observables is None: @@ -557,10 +1485,48 @@ def measure(self, observables=None): ) if isinstance(observables, dict): return { - name: self.vstate.expect(op) + name: self.vstate.expect( + _resolve_netket_measurement_observable( + self.hilbert, + self.ansatz, + op, + ) + ) for name, op in observables.items() } - return self.vstate.expect(observables) + return self.vstate.expect( + _resolve_netket_measurement_observable( + self.hilbert, + self.ansatz, + observables, + ) + ) + + def measure( + self, + observables=None, + *, + samples=None, + sampling=None, + progress=False, + ): + """Measure observables, optionally from an explicitly retained batch. + + ``setup.measure_samples(samples, observables)`` is the explicit + sample-once/measure-many form. This convenience wrapper preserves the + direct API: it uses the current NetKet cache when available, or calls + :meth:`sample` once when a batch must be drawn. Pass ``samples=`` to + make the sampling boundary explicit. + """ + if samples is not None and sampling is not None: + raise ValueError("Pass either samples or sampling, not both.") + if samples is None: + cached_samples = getattr(self.vstate, "_samples", None) + if sampling is not None or cached_samples is None: + samples = self.sample(sampling, progress=progress) + else: + samples = cached_samples + return self.measure_samples(samples, observables) @dataclass(frozen=True) @@ -594,7 +1560,36 @@ def n_sites(self): def n_params(self): return self.setup.n_params - def sample(self, sampling=None): + @property + def build_timing(self): + """Return the setup-phase timing breakdown, when available.""" + return self.setup.build_timing + + def check_mc_convergence(self, hamiltonian=None, **kwargs): + """Forward NetKet's post-optimization mixing diagnostic.""" + return self.setup.check_mc_convergence(hamiltonian, **kwargs) + + def thermalise(self, hamiltonian=None, **kwargs): + """Forward NetKet's in-place chain thermalisation helper.""" + return self.setup.thermalise(hamiltonian, **kwargs) + + def expect_to_precision(self, observable=None, **kwargs): + """Forward NetKet's precision-targeted expectation helper.""" + return self.setup.expect_to_precision(observable, **kwargs) + + def benchmark_amplitude(self, configs=None, *, n_samples=1): + """Time a synchronized amplitude batch on the native NetKet setup.""" + return self.setup.benchmark_amplitude(configs, n_samples=n_samples) + + def sample( + self, + sampling=None, + *, + fresh=False, + progress=False, + resource_monitor=False, + resource_interval=0.25, + ): """Collect samples as backend-neutral :class:`VMCSamples`. Chain count and seeds belong to an MCState at build time in NetKet. @@ -604,12 +1599,17 @@ def sample(self, sampling=None): from .api import SamplingConfig, VMCBackendCapabilityError if sampling is None: - return self.setup.sample() + return self.setup.sample( + fresh=fresh, + progress=progress, + resource_monitor=resource_monitor, + resource_interval=resource_interval, + ) if not isinstance(sampling, SamplingConfig): raise TypeError("sampling must be a SamplingConfig or None.") unsupported = [] - if sampling.thin != 1: - unsupported.append("thin") + if sampling.sweep_size != 1: + unsupported.append("sweep_size (formerly thin)") if sampling.seed is not None or sampling.sampler_seed is not None: unsupported.append("seed/sampler_seed") if sampling.proposal is not None: @@ -621,7 +1621,13 @@ def sample(self, sampling=None): "construction. Supply it while building the native sampler, " "or use the native NetKet API explicitly." ) - return self.setup.sample(sampling) + return self.setup.sample( + sampling, + fresh=fresh, + progress=progress, + resource_monitor=resource_monitor, + resource_interval=resource_interval, + ) def measure( self, @@ -635,14 +1641,16 @@ def measure( """Measure energy and optional observables on the current VMC state.""" from .api import VMCBackendCapabilityError, VMCMeasurement - if samples is not None or weights is not None or proposal_log_probs is not None: + if weights is not None or proposal_log_probs is not None: raise VMCBackendCapabilityError( - "NetKet's portable adapter does not yet accept externally " - "supplied weighted sample batches. Use its MCState sampling " - "path, or use the Torch adapter for importance sampling." + "NetKet's portable adapter does not accept weighted or " + "proposal-distribution sample batches. Use the Torch adapter " + "for importance sampling." ) - - samples = self.sample(sampling) if sampling is not None else None + if samples is not None and sampling is not None: + raise ValueError("Pass either samples or sampling, not both.") + if samples is None: + samples = self.sample(sampling) if observables is None: extra = dict(getattr(self.setup, "observables", None) or {}) else: @@ -655,7 +1663,10 @@ def measure( "'energy' is reserved for problem.hamiltonian; use a different " "observable name." ) - native = self.setup.measure({"energy": self.setup.hamiltonian, **extra}) + native = self.setup.measure_samples( + samples, + {"energy": self.setup.hamiltonian, **extra}, + ) energy = native["energy"] return VMCMeasurement( energy_mean=getattr(energy, "mean", energy), @@ -718,7 +1729,16 @@ def optimize(self, optimization=None, *, n_steps=None, **kwargs): per_site=per_site, diagnostics={ "backend": self.backend, + "build_timing": ( + self.setup.build_timing.as_dict() + if self.setup.build_timing is not None + else None + ), "compile_seconds": native.compile_seconds, + "warmup_seconds": native.warmup_seconds, + "optimization_seconds": native.optimization_seconds, + "total_seconds": native.total_seconds, + "optimization_seconds_per_step": native.optimization_seconds_per_step, }, native=native, ) @@ -857,8 +1877,8 @@ def sample(self, sampling=None): ) result = sampler.sample( n_samples=sampling.n_samples, - n_discard_per_chain=sampling.burn_in, - n_thin=sampling.thin, + n_discard_per_chain=sampling.n_discard_per_chain, + sweep_size=sampling.sweep_size, ) object.__setattr__(self, "configs", sampler.configs) object.__setattr__(self, "amplitudes", sampler.amplitudes) @@ -902,12 +1922,15 @@ def configure_jax_for_vmc( preallocate=False, mem_fraction=0.65, platform=None, + compilation_cache_dir=None, disable_netket_tips=True, ): """Set JAX/NetKet environment defaults for notebook VMC runs. Call this before importing ``jax`` or ``netket``. Existing environment - values are preserved. + values are preserved. When ``compilation_cache_dir`` is set, it must be a + private, trusted location: JAX treats a persistent compilation cache as + executable trusted input. """ os.environ.setdefault( "XLA_PYTHON_CLIENT_PREALLOCATE", @@ -917,6 +1940,11 @@ def configure_jax_for_vmc( os.environ.setdefault("XLA_PYTHON_CLIENT_MEM_FRACTION", str(mem_fraction)) if platform is not None: os.environ.setdefault("JAX_PLATFORMS", str(platform)) + if compilation_cache_dir is not None: + compilation_cache_dir = os.fspath(compilation_cache_dir) + if not compilation_cache_dir: + raise ValueError("compilation_cache_dir must not be empty.") + os.environ.setdefault("JAX_COMPILATION_CACHE_DIR", compilation_cache_dir) if disable_netket_tips: os.environ.setdefault("NETKET_NO_TIPS", "1") @@ -1026,13 +2054,12 @@ def _require_static_cutoff_for_jit(name, contraction, cutoff): def _maybe_register_stable_jax_svd(contraction): - """Install Pepsy's regularized JAX SVD backward rule for VMC gradients. + """Install Pepsy's truncation-safe JAX SVD backward rule for VMC gradients. - HOTRG/CTMRG/boundary-MPS contractions compress with SVDs whose naive - reverse-mode rule diverges on near-degenerate singular values, which - destabilizes gradient-based VMC. Registering the relative-broadened custom - VJP (the same rule ``PepsEnergyOptimizer`` uses) stabilizes those - gradients. Exact contraction has no SVD, so registration is skipped. This + HOTRG/CTMRG/boundary-MPS contractions compress with SVDs. Quimb may retain + only the leading singular-vector columns, so Pepsy's registered thin-SVD + pullback restores the omitted zero cotangents before using JAX's native + derivative. Exact contraction has no SVD, so registration is skipped. This is a global autoray side effect and a soft no-op when JAX is unavailable. """ from .api import ContractionConfig @@ -1053,6 +2080,107 @@ def _contraction_options(contraction_opts): return {} if contraction_opts is None else dict(contraction_opts) +_FLAT_SYMMRAY_BOUNDARY_FALLBACK_WARNED = False +_FLAT_SYMMRAY_CTMRG_FALLBACK_WARNED = False + + +def _is_flat_symmray_network(tn): + """Return whether ``tn`` contains flat Symmray tensor data.""" + for tensor in tn: + data = getattr(tensor, "data", None) + if _is_symmray_array(data): + return "Flat" in type(data).__name__ + return False + + +def _contract_boundary_for_vmc(tn, *, max_bond, cutoff, method_opts): + """Contract a PEPS boundary, with a flat-Symmray compatibility retry. + + Quimb's ``max_separation=0`` path can ask its canonizer to unfuse an + empty boundary axis. Current flat fermionic Symmray arrays cannot represent + that intermediate operation, which surfaces as an ``align_axes`` or + ``unfuse`` exception during JAX tracing. The requested sequence, chi, + canonization, and other options are still honored; only this one stopping + threshold is relaxed to Quimb's stable ``1`` fallback for flat Symmray. + Dense/non-Symmray networks and all other failures are re-raised unchanged. + """ + global _FLAT_SYMMRAY_BOUNDARY_FALLBACK_WARNED + kwargs = dict(method_opts) + try: + return tn.contract_boundary( + max_bond=max_bond, + cutoff=cutoff, + strip_exponent=True, + **kwargs, + ) + except (AttributeError, TypeError, ValueError): + if ( + kwargs.get("max_separation", 1) == 0 + and _is_flat_symmray_network(tn) + ): + if not _FLAT_SYMMRAY_BOUNDARY_FALLBACK_WARNED: + warnings.warn( + "Flat Symmray JAX boundary contraction does not support " + "the max_separation=0 intermediate axis path; retrying " + "with max_separation=1. The requested sequence, chi, and " + "canonization options remain active.", + RuntimeWarning, + stacklevel=3, + ) + _FLAT_SYMMRAY_BOUNDARY_FALLBACK_WARNED = True + kwargs["max_separation"] = 1 + return tn.contract_boundary( + max_bond=max_bond, + cutoff=cutoff, + strip_exponent=True, + **kwargs, + ) + raise + + +def _contract_ctmrg_for_vmc(tn, *, max_bond, cutoff, method_opts): + """Contract CTMRG, retrying flat-Symmray ``max_separation=0`` safely. + + CTMRG delegates boundary compression to Quimb internally. On flat + fermionic Symmray arrays, its zero-separation intermediate can produce a + JAX block-matmul shape error (for example an environment axis of size + ``chi`` paired with a Z2 block axis). Keep all requested options active and + relax only this stopping threshold to Quimb's stable value ``1``. + """ + global _FLAT_SYMMRAY_CTMRG_FALLBACK_WARNED + kwargs = dict(method_opts) + try: + return tn.contract_ctmrg( + max_bond=max_bond, + cutoff=cutoff, + strip_exponent=True, + **kwargs, + ) + except (AttributeError, TypeError, ValueError): + if ( + kwargs.get("max_separation", 1) == 0 + and _is_flat_symmray_network(tn) + ): + if not _FLAT_SYMMRAY_CTMRG_FALLBACK_WARNED: + warnings.warn( + "Flat Symmray JAX CTMRG does not support the " + "max_separation=0 intermediate axis path; retrying " + "with max_separation=1. The requested sequence, chi, " + "and canonization options remain active.", + RuntimeWarning, + stacklevel=3, + ) + _FLAT_SYMMRAY_CTMRG_FALLBACK_WARNED = True + kwargs["max_separation"] = 1 + return tn.contract_ctmrg( + max_bond=max_bond, + cutoff=cutoff, + strip_exponent=True, + **kwargs, + ) + raise + + def _resolve_netket_contraction(contraction, chi, cutoff, contraction_opts): """Resolve a common ContractionConfig for NetKet public builders.""" from .api import ContractionConfig @@ -1112,7 +2240,7 @@ def _resolve_sampling_build_config( return ( sampling.n_samples, sampling.n_chains, - sampling.burn_in, + sampling.n_discard_per_chain, sampling.chunk_size if sampling.chunk_size is not None else chunk_size, sampling.seed if sampling.seed is not None else seed, sampling.sampler_seed if sampling.sampler_seed is not None else sampler_seed, @@ -1227,7 +2355,7 @@ def prepare_fermionic_peps_for_netket(peps, *, device=None): sr = _require_symmray() try: converted = data.to_flat() - except RuntimeError: + except (RuntimeError, ValueError): converted = _z2_flat_padded(data, sr) padded_sites.append(site) tn[site].modify(data=converted) @@ -1487,7 +2615,13 @@ def make_netket_autochunk_callback( ): """Create NetKet's auto-chunk callback for sampler/forward/backward OOMs.""" nk = _require_netket() - return nk.callbacks.AutoChunkSize( + callback_cls = getattr(getattr(nk, "callbacks", None), "AutoChunkSize", None) + if callback_cls is None: + raise RuntimeError( + "NetKet AutoChunkSize requires NetKet >= 3.22; " + f"found {getattr(nk, '__version__', 'an unknown version')!r}." + ) + return callback_cls( sampler_chunk_size=sampler_chunk_size, chunk_size=chunk_size, chunk_size_bwd=chunk_size_bwd, @@ -1627,6 +2761,33 @@ def _infer_lattice_shape_from_peps(peps): return Lx, Ly +def _infer_lattice_shape_from_fermi_terms(terms): + """Infer ``(Lx, Ly)`` from coordinate-keyed native fermion terms. + + Integer site labels do not contain enough information to distinguish a + rectangular lattice from a one-dimensional indexing scheme, so this + fallback deliberately accepts only coordinate-keyed native terms. The + normal public path still prefers the PEPS geometry when it is available. + """ + if terms is None or not hasattr(terms, "keys"): + return None + coordinates = [] + for key in terms.keys(): + if _is_coordinate_edge_key(key): + coordinates.extend(tuple(site) for site in key) + elif isinstance(key, tuple) and len(key) == 2: + try: + coordinates.append((int(key[0]), int(key[1]))) + except (TypeError, ValueError): + continue + if not coordinates: + return None + return ( + max(int(site[0]) for site in coordinates) + 1, + max(int(site[1]) for site in coordinates) + 1, + ) + + def _site_index_for_lattice(site, Lx, Ly): if isinstance(site, Integral): site = int(site) @@ -1713,9 +2874,18 @@ def _infer_pbc_from_fermi_terms(terms, Lx, Ly): pbc_x = False pbc_y = False for key in terms.keys(): - if not _is_coordinate_edge_key(key): - continue - (i0, j0), (i1, j1) = key + if _is_coordinate_edge_key(key): + (i0, j0), (i1, j1) = key + else: + try: + left, right = tuple(key) + left, right = int(left), int(right) + except (TypeError, ValueError): + continue + if not (0 <= left < Lx * Ly and 0 <= right < Lx * Ly): + continue + i0, j0 = divmod(left, Ly) + i1, j1 = divmod(right, Ly) i0, j0, i1, j1 = int(i0), int(j0), int(i1), int(j1) pbc_x |= {i0, i1} == {0, Lx - 1} and j0 == j1 and Lx > 2 pbc_y |= {j0, j1} == {0, Ly - 1} and i0 == i1 and Ly > 2 @@ -2031,18 +3201,18 @@ def contract_mantissa_exponent(tnx): **method_opts, ) if contraction == "ctmrg": - return tnx.contract_ctmrg( + return _contract_ctmrg_for_vmc( + tnx, max_bond=chi, cutoff=cutoff, - strip_exponent=True, - **method_opts, + method_opts=method_opts, ) if contraction == "boundary": - return tnx.contract_boundary( + return _contract_boundary_for_vmc( + tnx, max_bond=chi, cutoff=cutoff, - strip_exponent=True, - **method_opts, + method_opts=method_opts, ) amp = tnx.contract(all) return amp, jnp.zeros((), dtype=real_dtype) @@ -2118,18 +3288,18 @@ def evaluate_one(tn, phys): **method_opts, ) elif contraction == "ctmrg": - mantissa, exponent = tnx.contract_ctmrg( + mantissa, exponent = _contract_ctmrg_for_vmc( + tnx, max_bond=chi, cutoff=cutoff, - strip_exponent=True, - **method_opts, + method_opts=method_opts, ) elif contraction == "boundary": - mantissa, exponent = tnx.contract_boundary( + mantissa, exponent = _contract_boundary_for_vmc( + tnx, max_bond=chi, cutoff=cutoff, - strip_exponent=True, - **method_opts, + method_opts=method_opts, ) else: mantissa = tnx.contract(all) @@ -2272,6 +3442,58 @@ def __call__(self, x): return PEPSLogAmplitude() +def _netket_variables_from_ansatz(ansatz, param_dtype=None): + """Build Flax variables directly from an already-packed PEPS ansatz. + + ``MCState`` otherwise calls ``model.init`` with a dummy configuration. + For these PEPS models that dummy call needlessly traces the contraction + once, even though the packed leaves are already the desired parameters. + """ + _, jnp = _require_jax() + dtype = None if param_dtype is None else param_dtype + return { + "params": { + f"t{k}": jnp.asarray(leaf, dtype=dtype) + for k, leaf in enumerate(ansatz.leaves) + } + } + + +def _peps_params_from_netket_variables(ansatz, variables, *, device_get=True): + """Rebuild quimb's packed parameter tree from Flax/NetKet variables. + + The PEPS Flax models deliberately name their leaves ``t0``, ``t1``, ... + in :func:`_netket_variables_from_ansatz`. Keeping the inverse conversion + here makes that private model detail explicit and gives the public setup + handoff a useful validation error instead of an opaque ``qtn.unpack`` + failure. + """ + if not hasattr(variables, "get"): + raise TypeError( + "variables must be a Flax variables mapping or a parameters mapping." + ) + parameters = variables.get("params", variables) + if not hasattr(parameters, "__getitem__"): + raise TypeError("variables['params'] must be a mapping of PEPS leaves.") + + expected = tuple(f"t{k}" for k in range(len(ansatz.leaves))) + try: + leaves = tuple(parameters[name] for name in expected) + except KeyError as error: + raise ValueError( + "NetKet PEPS parameters are missing leaf " + f"{error.args[0]!r}; expected {expected!r}." + ) from error + + if device_get: + jax, _ = _require_jax() + leaves = tuple(np.asarray(jax.device_get(leaf)) for leaf in leaves) + else: + leaves = tuple(leaves) + jax, _ = _require_jax() + return jax.tree_util.tree_unflatten(ansatz.treedef, leaves) + + def _make_fermionic_peps_batched_amplitude_apply( ansatz, columns, @@ -2335,18 +3557,18 @@ def contract_mantissa_exponent(tnx): **method_opts, ) if contraction == "ctmrg": - return tnx.contract_ctmrg( + return _contract_ctmrg_for_vmc( + tnx, max_bond=chi, cutoff=cutoff, - strip_exponent=True, - **method_opts, + method_opts=method_opts, ) if contraction == "boundary": - return tnx.contract_boundary( + return _contract_boundary_for_vmc( + tnx, max_bond=chi, cutoff=cutoff, - strip_exponent=True, - **method_opts, + method_opts=method_opts, ) amp = tnx.contract(all) return amp, jnp.zeros((), dtype=real_dtype) @@ -2427,18 +3649,18 @@ def evaluate_one(tn, phys): **method_opts, ) elif contraction == "ctmrg": - mantissa, exponent = tnx.contract_ctmrg( + mantissa, exponent = _contract_ctmrg_for_vmc( + tnx, max_bond=chi, cutoff=cutoff, - strip_exponent=True, - **method_opts, + method_opts=method_opts, ) elif contraction == "boundary": - mantissa, exponent = tnx.contract_boundary( + mantissa, exponent = _contract_boundary_for_vmc( + tnx, max_bond=chi, cutoff=cutoff, - strip_exponent=True, - **method_opts, + method_opts=method_opts, ) else: mantissa = tnx.contract(all) @@ -2677,6 +3899,7 @@ def _build_netket_peps_vmc( vstate = _require_netket().vqs.MCState( sampler, model, + variables=_netket_variables_from_ansatz(ansatz, param_dtype), n_samples=n_samples, n_discard_per_chain=n_discard_per_chain, chunk_size=_check_positive_int("chunk_size", chunk_size), @@ -3118,7 +4341,10 @@ def compile_operator_sum_netket(hilbert, terms, *, site_order=None, conserving=F Symbolic fermion products are lowered to the existing :func:`netket_fermion_operator` primitive. Local matrix terms are lowered to ``nk.operator.LocalOperator``. The identity constant is included in the - returned native operator, so callers must not add it again. + returned native operator, so callers must not add it again. Pass + ``conserving=\"auto\"`` to use NetKet's reduced fixed-sector fermion + operator when the symbolic terms preserve particle number and spin; terms + that do not preserve the sector retain the generic operator. """ from .api import ( LocalMatrixTerm, @@ -3207,19 +4433,29 @@ def map_site(site): return total -def fermion_model_terms(fermion, edges, *, n_sites=None): +def fermion_model_terms( + fermion, + edges, + *, + t, + U, + V=0.0, + mu=0.0, + n_sites=None, +): r"""Return symbolic Hamiltonian terms for a spinful :class:`pepsy.Fermion`. Reconstructs the hopping (:math:`-t`), on-site Hubbard (:math:`U\,n_\uparrow n_\downarrow`), nearest-neighbor density (:math:`V\,n_i n_j`) and chemical-potential (:math:`-\mu\,n`) terms of - ``fermion`` over the integer ``edges`` as a list of ``(coefficient, ops)`` - pairs (see :func:`netket_fermion_operator`). This lets any spinful - :class:`pepsy.Fermion` model drive a NetKet VMC run, not only plain - Fermi-Hubbard. - - Only spinful fermions with uniform scalar parameters are supported; supply - per-edge/per-site couplings directly as explicit terms. + explicit coefficients over the integer ``edges`` as a list of + ``(coefficient, ops)`` pairs (see :func:`netket_fermion_operator`). The + coefficients are arguments rather than state on ``fermion``, so the + symbolic NetKet operator cannot diverge from the native Hamiltonian. + + Only uniform scalar parameters are supported; supply non-uniform couplings + as a native ``fermion.hamiltonian({...})`` mapping to + :func:`build_fermion_vmc`. """ if not getattr(fermion, "spinful", True): raise NotImplementedError( @@ -3231,20 +4467,19 @@ def fermion_model_terms(fermion, edges, *, n_sites=None): n_sites = 1 + max((max(i, j) for i, j in edges), default=-1) sites = range(int(n_sites)) - def _scalar(name, default): - value = getattr(fermion, name, default) + def _scalar(name, value): try: return float(value) except (TypeError, ValueError) as exc: raise TypeError( - f"fermion.{name} must be a real scalar for fermion_model_terms; " + f"{name} must be a real scalar for fermion_model_terms; " f"got {value!r}. Pass explicit terms for non-uniform couplings." ) from exc - t = _scalar("t", 1.0) - U = _scalar("U", 0.0) - V = _scalar("V", 0.0) - mu = _scalar("mu", 0.0) + t = _scalar("t", t) + U = _scalar("U", U) + V = _scalar("V", V) + mu = _scalar("mu", mu) terms = [] for i, j in edges: @@ -3286,8 +4521,256 @@ def _scalar(name, default): return terms +def _native_term_support(where, *, coordinate_sites): + """Return the ordered one- or two-site support encoded by a term key.""" + if coordinate_sites and ( + isinstance(where, (tuple, list)) + and len(where) == 2 + and all(isinstance(value, Integral) for value in where) + ): + return (tuple(int(value) for value in where),) + if isinstance(where, (tuple, list)): + support = tuple(where) + else: + support = (where,) + if len(support) not in {1, 2}: + raise ValueError( + "Native Fermion Hamiltonian terms must have one-site or two-site keys; " + f"got {where!r}." + ) + return support + + +def _native_term_to_numpy(term): + """Transfer one small native local term to a host dense matrix.""" + dense = term.to_dense() if hasattr(term, "to_dense") else term + detach = getattr(dense, "detach", None) + if callable(detach): + dense = detach() + cpu = getattr(dense, "cpu", None) + if callable(cpu): + dense = cpu() + return np.asarray(dense, dtype=np.complex128) + + +def _project_native_term(matrix, candidates, *, where): + """Expand a native local matrix in a small Fermi-Hubbard operator basis.""" + columns = np.stack([candidate.reshape(-1) for candidate in candidates], axis=1) + coefficients, _, _, _ = np.linalg.lstsq(columns, matrix.reshape(-1), rcond=None) + reconstructed = (columns @ coefficients).reshape(matrix.shape) + residual = np.linalg.norm(matrix - reconstructed) + scale = max(1.0, float(np.linalg.norm(matrix))) + if residual > 5.0e-6 * scale: + raise ValueError( + "Native term at " + f"{where!r} is not in the supported spinful Fermi-Hubbard local " + "operator span. Supply an explicit symbolic OperatorSum to NetKet " + "for a custom interaction." + ) + return coefficients + + +def _native_fermi_hubbard_terms_to_netket(fermion, terms, *, site_order): + """Compile explicit native Hubbard terms into NetKet fermion monomials. + + The native Symmray terms are the authoritative Hamiltonian. Their local + matrices are decomposed into the neutral Fermi-Hubbard basis (identity, + spin-resolved number, doublon, hopping, and density terms), then emitted + as NetKet creation/annihilation products. This retains NetKet's fermionic + Jordan-Wigner signs while allowing the native and VMC paths to share the + exact same explicit term mapping. + """ + if fermion is None or not getattr(fermion, "spinful", False): + raise TypeError( + "Compiling native terms for NetKet requires a spinful Fermion helper." + ) + terms = _native_fermion_terms_mapping(terms) + coordinate_sites = any( + isinstance(where, (tuple, list)) + and len(where) == 2 + and all( + isinstance(site, (tuple, list)) + and len(site) == 2 + and all(isinstance(value, Integral) for value in site) + for site in where + ) + for where in terms + ) + site_order = tuple(site_order) + site_to_orbital = {site: orbital for orbital, site in enumerate(site_order)} + + def orbital(site): + if site in site_to_orbital: + return site_to_orbital[site] + if isinstance(site, Integral) and 0 <= int(site) < len(site_order): + return int(site) + raise ValueError( + f"Native term site {site!r} is not present in the PEPS site order." + ) + + # Build basis operators on CPU: native input terms can live on JAX/Torch, + # but compilation only needs tiny 4x4 / 16x16 host matrices. + from ..tensors import Fermion # pylint: disable=import-outside-toplevel + + reference = Fermion( + spinful=True, + symmetry=fermion.symmetry, + dtype=fermion.dtype, + ) + one_site_basis = ( + _native_term_to_numpy(reference.observable("identity")), + _native_term_to_numpy(reference.observable("number_up")), + _native_term_to_numpy(reference.observable("number_down")), + _native_term_to_numpy(reference.interaction_operator()), + ) + two_site_basis = ( + _native_term_to_numpy(reference.hopping_operator(spin="up")), + _native_term_to_numpy( + reference.hopping_operator(spin="up", peierls_angle=np.pi / 2) + ), + _native_term_to_numpy(reference.hopping_operator(spin="down")), + _native_term_to_numpy( + reference.hopping_operator(spin="down", peierls_angle=np.pi / 2) + ), + _native_term_to_numpy(reference.density_operator()), + _native_term_to_numpy( + reference.operator_term( + [(1.0, ((0, "number_up"),))], sites=(0, 1) + ) + ), + _native_term_to_numpy( + reference.operator_term( + [(1.0, ((0, "number_down"),))], sites=(0, 1) + ) + ), + _native_term_to_numpy( + reference.operator_term( + [(1.0, ((1, "number_up"),))], sites=(0, 1) + ) + ), + _native_term_to_numpy( + reference.operator_term( + [(1.0, ((1, "number_down"),))], sites=(0, 1) + ) + ), + _native_term_to_numpy( + reference.operator_term( + [(1.0, ((0, "double"),))], sites=(0, 1) + ) + ), + _native_term_to_numpy( + reference.operator_term( + [(1.0, ((1, "double"),))], sites=(0, 1) + ) + ), + _native_term_to_numpy( + reference.operator_term([(1.0, ())], sites=(0, 1)) + ), + ) + + symbolic = [] + constant = 0.0j + for where, term in dict(terms).items(): + support = _native_term_support(where, coordinate_sites=coordinate_sites) + matrix = _native_term_to_numpy(term) + if len(support) == 1: + (site,) = support + coeff_identity, coeff_up, coeff_down, coeff_double = _project_native_term( + matrix, + one_site_basis, + where=where, + ) + constant += coeff_identity + target = orbital(site) + if abs(coeff_up) > 1.0e-10: + symbolic.append((coeff_up, ((target, 1, True), (target, 1, False)))) + if abs(coeff_down) > 1.0e-10: + symbolic.append((coeff_down, ((target, -1, True), (target, -1, False)))) + if abs(coeff_double) > 1.0e-10: + symbolic.append( + ( + coeff_double, + ( + (target, 1, True), + (target, 1, False), + (target, -1, True), + (target, -1, False), + ), + ) + ) + continue + + left, right = (orbital(site) for site in support) + ( + up_real, + up_imag, + down_real, + down_imag, + density, + left_up, + left_down, + right_up, + right_down, + left_double, + right_double, + identity, + ) = _project_native_term(matrix, two_site_basis, where=where) + constant += identity + for sz, real, imag in ( + (1, up_real, up_imag), + (-1, down_real, down_imag), + ): + forward = real + 1.0j * imag + backward = real - 1.0j * imag + if abs(forward) > 1.0e-10: + symbolic.append((forward, ((left, sz, True), (right, sz, False)))) + if abs(backward) > 1.0e-10: + symbolic.append((backward, ((right, sz, True), (left, sz, False)))) + if abs(density) > 1.0e-10: + for left_sz in (1, -1): + for right_sz in (1, -1): + symbolic.append( + ( + density, + ( + (left, left_sz, True), + (left, left_sz, False), + (right, right_sz, True), + (right, right_sz, False), + ), + ) + ) + for target, sz, coefficient in ( + (left, 1, left_up), + (left, -1, left_down), + (right, 1, right_up), + (right, -1, right_down), + ): + if abs(coefficient) > 1.0e-10: + symbolic.append( + (coefficient, ((target, sz, True), (target, sz, False))) + ) + for target, coefficient in ( + (left, left_double), + (right, right_double), + ): + if abs(coefficient) > 1.0e-10: + symbolic.append( + ( + coefficient, + ( + (target, 1, True), + (target, 1, False), + (target, -1, True), + (target, -1, False), + ), + ) + ) + return symbolic, constant + + def standard_fermion_observables(hilbert): - """Return common spinful-fermion observables for :meth:`NetKetPEPSVMC.measure`. + """Return common spinful observables for :meth:`NetKetPEPSVMC.measure_samples`. The returned ``{name: operator}`` mapping contains the total particle number per spin (``"n_up"``, ``"n_down"``), the total particle number @@ -3302,15 +4785,144 @@ def standard_fermion_observables(hilbert): for i in range(n) ] return { - "n_up": netket_fermion_operator(hilbert, n_up), - "n_down": netket_fermion_operator(hilbert, n_down), - "n_total": netket_fermion_operator(hilbert, n_up + n_down), - "double_occupancy": netket_fermion_operator(hilbert, doub), + "n_up": netket_fermion_operator(hilbert, n_up, conserving="auto"), + "n_down": netket_fermion_operator(hilbert, n_down, conserving="auto"), + "n_total": netket_fermion_operator( + hilbert, n_up + n_down, conserving="auto" + ), + "double_occupancy": netket_fermion_operator( + hilbert, doub, conserving="auto" + ), } +def _eta_pair_lattice_sites(ansatz): + """Return a rectangular zero-origin coordinate lattice from an ansatz.""" + sites = tuple(getattr(ansatz, "orbital_sites", ())) + if not sites: + raise ValueError( + "Eta-pair measurement needs a packed PEPS with coordinate " + "orbital_sites." + ) + if len(set(sites)) != len(sites): + raise ValueError("Eta-pair measurement requires unique orbital_sites.") + if any( + not isinstance(site, tuple) + or len(site) != 2 + or any( + isinstance(coord, bool) or not isinstance(coord, Integral) + for coord in site + ) + for site in sites + ): + raise ValueError( + "Eta-pair measurement requires two-dimensional integer " + "orbital_sites." + ) + + sites = tuple((int(x), int(y)) for x, y in sites) + xs = {x for x, _ in sites} + ys = {y for _, y in sites} + if min(xs) != 0 or min(ys) != 0: + raise ValueError( + "Eta-pair measurement requires a zero-origin rectangular lattice." + ) + Lx, Ly = max(xs) + 1, max(ys) + 1 + expected = {(x, y) for x in range(Lx) for y in range(Ly)} + if set(sites) != expected: + raise ValueError( + "Eta-pair measurement requires orbital_sites to cover a complete " + "rectangular lattice." + ) + return sites, Lx, Ly + + +def _netket_eta_pair_operator(hilbert, ansatz, specification): + """Compile a declarative eta-pair specification for a packed PEPS.""" + sites, Lx, Ly = _eta_pair_lattice_sites(ansatz) + site_to_orbital = {site: orbital for orbital, site in enumerate(sites)} + + if specification.dx == 0 and specification.dy == 0: + coefficient = 1.0 / len(sites) + terms = [ + ( + coefficient, + ( + (orbital, 1, True), + (orbital, 1, False), + (orbital, -1, True), + (orbital, -1, False), + ), + ) + for orbital in range(len(sites)) + ] + return netket_fermion_operator(hilbert, terms, conserving="auto") + + pairs = [] + site_set = set(sites) + for left in sites: + x, y = left + if specification.periodic: + right = ((x + specification.dx) % Lx, (y + specification.dy) % Ly) + else: + right = (x + specification.dx, y + specification.dy) + if right not in site_set: + continue + pairs.append((left, right)) + if not pairs: + raise ValueError("The requested eta-pair displacement has no valid pairs.") + + normalizer = len(sites) if specification.periodic else len(pairs) + terms = [] + for left, right in pairs: + phase = ( + -1.0 + if specification.staggered and (sum(left) + sum(right)) % 2 + else 1.0 + ) + coefficient = phase / normalizer + left_orbital = site_to_orbital[left] + right_orbital = site_to_orbital[right] + terms.extend( + ( + ( + coefficient, + ( + (left_orbital, 1, True), + (left_orbital, -1, True), + (right_orbital, -1, False), + (right_orbital, 1, False), + ), + ), + ( + coefficient, + ( + (right_orbital, 1, True), + (right_orbital, -1, True), + (left_orbital, -1, False), + (left_orbital, 1, False), + ), + ), + ) + ) + return netket_fermion_operator(hilbert, terms, conserving="auto") + + +def _resolve_netket_measurement_observable(hilbert, ansatz, observable): + """Compile declarative NetKet measurement observables on demand.""" + if isinstance(observable, NetKetEtaPairObservable): + return _netket_eta_pair_operator(hilbert, ansatz, observable) + return observable + + def _normalize_fermion_observables(hilbert, observables, *, site_order=None): - """Resolve an ``{name: operator_or_terms}`` mapping to NetKet operators.""" + """Resolve an ``{name: operator_or_terms}`` mapping to NetKet operators. + + Declarative symbolic operators request NetKet's conserving fermion + specialization when possible. Non-conserving observables safely fall back + to the generic operator, so pairing or spin-changing measurements retain + their physical meaning. + """ from .api import OperatorSum if observables is None: return None @@ -3327,9 +4939,14 @@ def _normalize_fermion_observables(hilbert, observables, *, site_order=None): hilbert, spec, site_order=site_order, + conserving="auto", ) else: - resolved[str(name)] = netket_fermion_operator(hilbert, spec) + resolved[str(name)] = netket_fermion_operator( + hilbert, + spec, + conserving="auto", + ) return resolved @@ -3391,13 +5008,13 @@ def build_fermi_hubbard_vmc( """Create NetKet VMC objects for a fixed-sector Fermi-Hubbard PEPS. ``fermion`` and ``terms`` may be the native Pepsy objects used during - imaginary-time evolution. ``terms`` is inspected only to infer the + imaginary-time evolution. ``terms`` is inspected only to infer the lattice edges and which Hamiltonian axes are periodic; its operator coefficients (including any chemical-potential or on-site shifts) are NOT used to build the Hamiltonian. The Hamiltonian is always NetKet's ``FermiHubbardJax`` with the resolved ``t``/``U`` (see below), so pass the - hopping/interaction through ``t``, ``U`` (or ``fermion``) rather than - through ``terms``. ``sector`` or + hopping/interaction through explicit ``t`` and ``U`` rather than through + ``terms``. ``sector`` or ``n_fermions_per_spin`` accepts either ``(N_up, N_down)`` or Pepsy's ``setup.spin_occupations`` mapping. The evolved PEPS is prepared for JAX internally, and its variational parameters use quimb's native @@ -3408,11 +5025,9 @@ def build_fermi_hubbard_vmc( setting ``MU=0`` once the spin sector is fixed. When ``register_stable_svd`` is True (default) and ``contraction`` is an - SVD-based approximation (``hotrg``/``ctmrg``/``boundary``), Pepsy's - regularized JAX SVD backward rule is installed globally via - :func:`pepsy.reg_rel_svd_jax` so VMC gradients through the compression SVDs - stay finite near degenerate singular values. Set it False to keep your own - autoray SVD registration. + SVD-based approximation (``hotrg``/``ctmrg``/``boundary``), Pepsy installs + a JAX thin-SVD pullback that safely handles Quimb's fixed-rank truncation. + Set it False to keep your own autoray SVD registration. """ ( n_samples, @@ -3449,11 +5064,10 @@ def build_fermi_hubbard_vmc( n_fermions_per_spin = tuple(int(value) for value in n_fermions_per_spin) if len(n_fermions_per_spin) != 2: raise ValueError("n_fermions_per_spin must contain (N_up, N_down).") - if fermion is not None: - if t is None: - t = getattr(fermion, "t", 1.0) - if U is None: - U = getattr(fermion, "U", 8.0) + if fermion is not None and not getattr(fermion, "spinful", True): + raise NotImplementedError( + "build_fermi_hubbard_vmc supports spinful fermions only." + ) t = 1.0 if t is None else t U = 8.0 if U is None else U if pbc is None: @@ -3511,6 +5125,7 @@ def build_fermi_hubbard_vmc( vstate = nk.vqs.MCState( sampler, model, + variables=_netket_variables_from_ansatz(ansatz, param_dtype), n_samples=n_samples, n_discard_per_chain=n_discard_per_chain, chunk_size=_check_positive_int("chunk_size", chunk_size), @@ -3549,6 +5164,7 @@ def build_fermion_vmc( hamiltonian=None, terms=None, observables=None, + config=None, Lx=None, Ly=None, n_fermions_per_spin=None, @@ -3556,7 +5172,7 @@ def build_fermion_vmc( pbc=None, edges=None, graph=None, - conserving="auto", + conserving=False, contraction="exact", chi=None, cutoff=0.0, @@ -3568,9 +5184,10 @@ def build_fermion_vmc( chunk_size=256, sampling=None, sampler_chunk_size=None, + sampler_sweep_size=None, seed=None, sampler_seed=None, - use_sr="auto", + use_sr=False, max_sr_params=5_000, sr_diag_shift=0.01, sr_diag_scale=None, @@ -3579,22 +5196,19 @@ def build_fermion_vmc( sr_solver_restart=False, param_dtype=None, verify_columns=False, + progress=False, ): """Create a NetKet VMC setup for a general spinful-fermion model. Unlike :func:`build_fermi_hubbard_vmc`, the Hamiltonian is not restricted to NetKet's ``FermiHubbardJax``. Define the model in one of these ways: - * ``fermion`` -- a :class:`pepsy.Fermion` whose hopping / interaction / - density / chemical-potential parameters are turned into a NetKet fermion - operator over ``edges`` (see :func:`fermion_model_terms`). This covers - Fermi-Hubbard, Hubbard + nearest-neighbor ``V``, and a chemical potential. * ``fermion`` **plus** native ``terms=`` / ``hamiltonian=`` -- pass the - native Pepsy ``SymHamiltonian`` (from ``fermion.hamiltonian(...)``) or its - coordinate-keyed ``.terms`` mapping to let the builder infer the integer - lattice ``edges`` and periodic axes (``pbc``) directly from the terms, - then rebuild the matching NetKet Hamiltonian from ``fermion``. This mirrors - the ergonomics of :class:`pepsy.vmc.TorchFermionVMC`. + authoritative native Pepsy ``SymHamiltonian`` (from + ``fermion.hamiltonian(...)``) or its coordinate-keyed ``.terms`` mapping. + The builder infers integer lattice ``edges`` and periodic axes (``pbc``) + directly from those terms, then compiles the supplied local operators to + a matching NetKet fermion operator. Couplings never live on ``fermion``. * ``terms`` -- an explicit list of symbolic ``(coefficient, ops)`` terms (see :func:`netket_fermion_operator`) for a custom fermionic model. * ``OperatorSum`` -- the backend-neutral term representation shared with @@ -3604,6 +5218,36 @@ def build_fermion_vmc( ``edges`` / ``graph`` / ``pbc``, when given explicitly, take precedence over any geometry inferred from native terms. + When omitted, ``Lx``/``Ly`` are inferred from the PEPS rectangular site + layout, with a coordinate-keyed native Hamiltonian as a fallback. Native + coordinate or row-major integer edge keys provide the graph and periodic + boundary inference. Pass :class:`~pepsy.vmc.ContractionConfig` and + :class:`~pepsy.vmc.SamplingConfig` through ``contraction=`` and + ``sampling=`` to keep those settings in one validated object. + + ``progress=True`` shows an eight-stage setup bar (settings, Hilbert/graph, + Hamiltonian, PEPS packing, JAX model, sampler, MCState, and optional SR + preconditioner). The returned + setup stores the corresponding :class:`NetKetBuildTiming` in + ``build_timing``. This setup timing is distinct from the lazy JAX + sampler/amplitude/energy compilation reported by :meth:`warmup`. + + ``conserving="auto"`` option asks NetKet to convert the operator to its + experimental particle-number/spin-conserving representation. That can + trigger a one-time Numba compilation during the first build; it is an + optional runtime optimization, not a physics change. The default + ``conserving=False`` builds the ordinary exact NetKet fermion operator + immediately. ``sampler_sweep_size`` is passed to NetKet's Metropolis sampler as the + number of proposals between retained samples. When omitted, NetKet uses + the Hilbert-space size (``2 * Lx * Ly`` for this spinful model), so set it + explicitly when you want the Markov-chain work to be obvious and + reproducible. ``SamplingConfig.burn_in`` remains the per-chain discard + count before retained samples. + + For a compact call, pass a :class:`NetKetVMCConfig` as ``config``. Its + numerical fields override the corresponding legacy keywords, while the + explicit ``fermion``/``hamiltonian`` inputs remain the model definition. + ``observables`` is an optional ``{name: operator_or_terms}`` mapping stored on the returned setup; call :meth:`NetKetPEPSVMC.measure` to evaluate them (see :func:`standard_fermion_observables` for common choices). All @@ -3611,6 +5255,35 @@ def build_fermion_vmc( :func:`build_fermi_hubbard_vmc`, and the returned setup exposes the same :meth:`NetKetPEPSVMC.warmup` and :meth:`NetKetPEPSVMC.optimize` helpers. """ + if config is not None: + if not isinstance(config, NetKetVMCConfig): + raise TypeError("config must be a NetKetVMCConfig or None.") + contraction = config.contraction + sampling = config.sampling + sampler_sweep_size = config.sampler_sweep_size + conserving = config.conserving + use_sr = config.use_sr + param_dtype = config.param_dtype + verify_columns = config.verify_columns + progress = config.progress + + build_started = time.perf_counter() + build_bar = _make_progress_bar( + total=8, desc="Build NetKet VMC", enabled=progress + ) + stage_started = build_started + stage_seconds = {} + + def mark_stage(name): + nonlocal stage_started + now = time.perf_counter() + elapsed = now - stage_started + stage_seconds[name] = elapsed + stage_started = now + if build_bar is not None: + build_bar.set_postfix_str(f"{name}: {elapsed:.1f}s") + build_bar.update(1) + ( n_samples, n_chains, @@ -3627,7 +5300,14 @@ def build_fermion_vmc( seed=seed, sampler_seed=sampler_seed, ) + contraction, chi, cutoff, contraction_opts = _resolve_netket_contraction( + contraction, + chi, + cutoff, + contraction_opts, + ) nk = _require_netket() + mark_stage("settings") from .api import OperatorSum common_hamiltonian = hamiltonian if isinstance(hamiltonian, OperatorSum) else None common_terms = terms if isinstance(terms, OperatorSum) else None @@ -3643,10 +5323,41 @@ def build_fermion_vmc( "build_fermion_vmc supports spinful fermions; use the sparse/torch " "path for spinless models." ) + + # Classify native inputs before resolving the lattice shape. This lets a + # coordinate-keyed SymHamiltonian provide geometry when the PEPS wrapper + # does not expose rectangular ``sites`` metadata. + prebuilt_operator = None + symbolic_terms = None + native_terms = None + if common_operator_sum is not None: + native_terms = common_operator_sum + elif hamiltonian is not None: + if hasattr(hamiltonian, "hilbert"): + prebuilt_operator = hamiltonian + elif _looks_like_native_fermion_terms(hamiltonian): + native_terms = _native_fermion_terms_mapping(hamiltonian) + else: + raise ValueError( + "hamiltonian must be a NetKet operator or a native Pepsy " + "SymHamiltonian / coordinate-keyed terms mapping." + ) + if common_operator_sum is None and terms is not None: + if _looks_like_native_fermion_terms(terms): + native_terms = _native_fermion_terms_mapping(terms) + else: + symbolic_terms = terms + if (Lx is None) != (Ly is None): raise ValueError("Lx and Ly must be supplied together or both omitted.") if Lx is None: - Lx, Ly = _infer_lattice_shape_from_peps(peps) + try: + Lx, Ly = _infer_lattice_shape_from_peps(peps) + except ValueError as peps_error: + inferred_shape = _infer_lattice_shape_from_fermi_terms(native_terms) + if inferred_shape is None: + raise peps_error + Lx, Ly = inferred_shape Lx, Ly = int(Lx), int(Ly) n_sites = Lx * Ly if sector is not None: @@ -3668,32 +5379,6 @@ def build_fermion_vmc( n_fermions_per_spin=n_fermions_per_spin, ) - # Classify how the model was supplied. ``terms``/``hamiltonian`` may be a - # native Pepsy SymHamiltonian or coordinate-keyed terms mapping (used to - # infer the lattice edges and periodicity, with the NetKet Hamiltonian - # rebuilt from ``fermion``), a symbolic ``(coefficient, ops)`` list, or an - # already-built NetKet operator. - prebuilt_operator = None - symbolic_terms = None - native_terms = None - if common_operator_sum is not None: - native_terms = common_operator_sum - elif hamiltonian is not None: - if hasattr(hamiltonian, "hilbert"): - prebuilt_operator = hamiltonian - elif _looks_like_native_fermion_terms(hamiltonian): - native_terms = _native_fermion_terms_mapping(hamiltonian) - else: - raise ValueError( - "hamiltonian must be a NetKet operator or a native Pepsy " - "SymHamiltonian / coordinate-keyed terms mapping." - ) - if common_operator_sum is None and terms is not None: - if _looks_like_native_fermion_terms(terms): - native_terms = _native_fermion_terms_mapping(terms) - else: - symbolic_terms = terms - # Infer lattice geometry (edges) and periodicity from native terms when # they were not given explicitly, mirroring build_fermi_hubbard_vmc. if pbc is None: @@ -3714,6 +5399,7 @@ def build_fermion_vmc( graph = nk.graph.Graph(edges=tuple(edges), n_nodes=n_sites) elif edges is None: edges = tuple(tuple(edge) for edge in graph.edges()) + mark_stage("Hilbert/graph") # Build the NetKet Hamiltonian. if common_operator_sum is not None: @@ -3729,17 +5415,36 @@ def build_fermion_vmc( hamiltonian = netket_fermion_operator( hilbert, symbolic_terms, conserving=conserving ) - elif fermion is not None: - model_terms = fermion_model_terms(fermion, edges, n_sites=n_sites) + elif native_terms is not None: + if fermion is None: + raise ValueError( + "Native Pepsy terms require fermion=... so their local " + "symmetry and basis can be validated for NetKet." + ) + native_symbolic_terms, native_constant = _native_fermi_hubbard_terms_to_netket( + fermion, + native_terms, + site_order=_row_major_sites(Lx, Ly), + ) hamiltonian = netket_fermion_operator( - hilbert, model_terms, conserving=conserving + hilbert, + native_symbolic_terms, + constant=native_constant, + conserving=conserving, + ) + elif fermion is not None: + raise ValueError( + "build_fermion_vmc requires explicit native hamiltonian=... or " + "terms=... with fermion=.... Fermion stores local symmetry and " + "backend conventions, not t/U/V/mu couplings." ) else: raise ValueError( - "Provide fermion=... (optionally with native terms=/hamiltonian= " - "for geometry), a symbolic terms=..., or an already-built NetKet " - "hamiltonian=... to define the model for build_fermion_vmc." + "Provide fermion=... plus native terms=/hamiltonian=..., a " + "symbolic terms=..., or an already-built NetKet hamiltonian=... " + "to define the model for build_fermion_vmc." ) + mark_stage("Hamiltonian") observable_ops = _normalize_fermion_observables( hilbert, @@ -3757,6 +5462,7 @@ def build_fermion_vmc( peps = prepare_fermionic_peps_for_netket(peps) ansatz = pack_fermionic_peps_ansatz(peps, lattice_shape=(Lx, Ly)) _warn_flat_z2_ansatz_fixed_u1u1_sector(ansatz, n_fermions_per_spin) + mark_stage("PEPS packing") model = make_fermionic_peps_log_amplitude_model( ansatz, columns, @@ -3766,6 +5472,7 @@ def build_fermion_vmc( contraction_opts=contraction_opts, param_dtype=param_dtype, ) + mark_stage("JAX model") sampler_kwargs = { "graph": graph, "n_chains": n_chains, @@ -3776,16 +5483,24 @@ def build_fermion_vmc( "sampler_chunk_size", sampler_chunk_size, ) + if sampler_sweep_size is not None: + sampler_kwargs["sweep_size"] = _check_positive_int( + "sampler_sweep_size", + sampler_sweep_size, + ) sampler = nk.sampler.MetropolisFermionHop(hilbert, **sampler_kwargs) + mark_stage("sampler") vstate = nk.vqs.MCState( sampler, model, + variables=_netket_variables_from_ansatz(ansatz, param_dtype), n_samples=n_samples, n_discard_per_chain=n_discard_per_chain, chunk_size=_check_positive_int("chunk_size", chunk_size), seed=seed, sampler_seed=sampler_seed, ) + mark_stage("MCState") preconditioner = _maybe_make_sr_preconditioner( ansatz, use_sr=use_sr, @@ -3796,6 +5511,23 @@ def build_fermion_vmc( sr_solver=sr_solver, sr_solver_restart=sr_solver_restart, ) + mark_stage("SR preconditioner") + if build_bar is not None: + build_bar.set_postfix_str( + f"done: {time.perf_counter() - build_started:.1f}s" + ) + build_bar.close() + build_timing = NetKetBuildTiming( + settings_seconds=stage_seconds["settings"], + geometry_seconds=stage_seconds["Hilbert/graph"], + hamiltonian_seconds=stage_seconds["Hamiltonian"], + peps_seconds=stage_seconds["PEPS packing"], + model_seconds=stage_seconds["JAX model"], + sampler_seconds=stage_seconds["sampler"], + total_seconds=time.perf_counter() - build_started, + preconditioner_seconds=stage_seconds["SR preconditioner"], + state_seconds=stage_seconds["MCState"], + ) return NetKetFermiHubbardVMC( hilbert=hilbert, graph=graph, @@ -3806,6 +5538,7 @@ def build_fermion_vmc( ansatz=ansatz, config_map=None, preconditioner=preconditioner, + build_timing=build_timing, columns=columns, observables=observable_ops, ) @@ -3849,8 +5582,8 @@ def build_netket_vmc( ) if sampling is not None: unsupported = [] - if sampling.thin != 1: - unsupported.append("thin") + if sampling.sweep_size != 1: + unsupported.append("sweep_size (formerly thin)") if sampling.proposal is not None: unsupported.append("proposal") if unsupported: @@ -4169,15 +5902,23 @@ def make_netket_vmc_driver( ) -def warmup_netket_vmc(setup, *, hamiltonian=None, progress=True, verbose=None): +def warmup_netket_vmc( + setup, + *, + hamiltonian=None, + progress=True, + verbose=None, + resource_monitor=False, + resource_interval=0.25, +): """Force XLA compilation of a NetKet VMC setup before ``driver.run(...)``. - The first optimization step compiles the sampler, the log-amplitude model - (including the CTMRG / boundary-MPS contraction), and the local-energy - kernel. That one-time compile cost is folded into the first ``tqdm`` tick, - so the NetKet progress-bar ETA is misleading until it clears. Calling this - once runs a single sample + energy evaluation so compilation happens up - front and the reported ETA is meaningful. + The first optimization step compiles the sampler, the jitted + log-amplitude model (including the CTMRG / boundary-MPS contraction), and + the local-energy/gradient kernel. That one-time compile cost is folded + into the first ``tqdm`` tick, so the NetKet progress-bar ETA is misleading + until it clears. Calling this once runs those same paths up front without + updating parameters, so the reported ETA is meaningful. Parameters ---------- @@ -4188,11 +5929,21 @@ def warmup_netket_vmc(setup, *, hamiltonian=None, progress=True, verbose=None): hamiltonian: Operator to evaluate; defaults to ``setup.hamiltonian``. progress: - When True (default), show a small two-stage ``tqdm`` bar - (sampler, then amplitude+energy) while compiling. + When True (default), show a small three-stage ``tqdm`` bar (sampler, + one jitted log-amplitude chunk, then local energy plus gradient) while + compiling. verbose: Print a short text message instead of / in addition to the bar. When ``None`` (default) it prints only if the progress bar is unavailable. + resource_monitor: + When True, use :class:`xyzpy.MemoryMonitor` when available to sample + host RSS, and ``nvidia-smi`` to sample GPU memory and utilization. + A compact report is printed after warmup. This remains opt-in because + GPU utilization sampling launches a lightweight subprocess. + resource_interval: + Seconds between GPU samples while ``resource_monitor=True``. Host RSS + is sampled more frequently by ``xyzpy`` so peak tracking has a quick + teardown. Returns ------- @@ -4210,36 +5961,107 @@ def warmup_netket_vmc(setup, *, hamiltonian=None, progress=True, verbose=None): "setup exposing a .hamiltonian attribute." ) bar = _make_progress_bar( - total=2, desc="Compiling VMC kernels", enabled=progress + total=3, desc="Warmup 1/3: sampler", enabled=progress ) if verbose is None: verbose = bar is None if verbose: print( - "Compiling NetKet VMC kernels (sampler, amplitude, energy)...", + "Compiling NetKet VMC kernels (sampler, log amplitude, energy gradient)...", flush=True, ) + monitor = ( + _NetKetResourceMonitor(interval=resource_interval) + if resource_monitor + else None + ) + if monitor is not None: + monitor.start() started = time.perf_counter() - vstate.reset() - # Stage 1: compile the Metropolis sampler. - _ = vstate.samples - if bar is not None: - bar.set_postfix_str("sampler") - bar.update(1) - # Stage 2: compile the amplitude + local-energy kernels. - stats = vstate.expect(hamiltonian) - # Resolve any lazy device arrays so compilation is finished before timing. - mean = getattr(stats, "mean", stats) - _ = float(np.asarray(mean).real) - elapsed = time.perf_counter() - started - if bar is not None: - bar.set_postfix_str(f"{elapsed:.1f}s") - bar.update(1) - bar.close() + try: + vstate.reset() + # Stage 1: compile the Metropolis sampler and collect the retained + # batch used by the following two warmup stages. + stage_started = time.perf_counter() + samples = vstate.samples + sampler_seconds = time.perf_counter() - stage_started + n_amplitude_rows = int(np.prod(samples.shape[:-1])) + sample_chains = int(samples.shape[0]) if samples.ndim >= 3 else 1 + samples_per_chain = n_amplitude_rows // sample_chains + sample_summary = ( + f"{n_amplitude_rows} retained = {sample_chains} chains x " + f"{samples_per_chain}/chain" + ) + if bar is not None: + bar.set_postfix_str(f"{sampler_seconds:.1f}s | {sample_summary}") + bar.update(1) + # Stage 2: compile NetKet's public JIT log-amplitude route. Use the + # configured forward chunk shape rather than the whole retained batch: + # VMC's chunked local-energy/gradient kernels use that shape in + # production. + chunk_size = getattr(vstate, "chunk_size", None) + if chunk_size is None: + amplitude_rows = n_amplitude_rows + else: + amplitude_rows = min(n_amplitude_rows, int(chunk_size)) + amplitude_configs = samples.reshape((-1, samples.shape[-1]))[:amplitude_rows] + if bar is not None: + bar.set_description_str( + f"Warmup 2/3: JIT log amplitudes ({amplitude_rows} rows)" + ) + amplitude_started = time.perf_counter() + log_value = getattr(vstate, "log_value", None) + if not callable(log_value): + raise TypeError( + "warmup_netket_vmc requires an MCState with a " + "log_value(configs) method." + ) + _block_until_ready(log_value(amplitude_configs)) + amplitude_seconds = time.perf_counter() - amplitude_started + if bar is not None: + bar.set_postfix_str( + f"{amplitude_seconds:.1f}s | {amplitude_rows}-row JIT chunk | " + f"{sample_summary}" + ) + bar.update(1) + # Stage 3: compile the driver-dominant local-energy and gradient route. + # ``expect_and_grad`` is what NetKet's ordinary VMC driver calls before + # updating parameters, so this deliberately leaves the PEPS unchanged. + if bar is not None: + bar.set_description_str("Warmup 3/3: local energy + gradient") + energy_started = time.perf_counter() + stats, gradient = vstate.expect_and_grad(hamiltonian) + # Resolve lazy device arrays so compilation is finished before timing. + mean = getattr(stats, "mean", stats) + _block_until_ready(mean) + _block_until_ready(gradient) + energy_seconds = time.perf_counter() - energy_started + elapsed = time.perf_counter() - started + if bar is not None: + energy_text = ( + f"{energy_seconds:.1f}s | E={float(np.asarray(mean).real):+.6f} | " + "local-energy and gradient kernels" + ) + bar.set_postfix_str(energy_text) + bar.update(1) + finally: + # Resource polling must not outlive a failed JIT/compilation attempt. + try: + if bar is not None: + bar.close() + finally: + resource_usage = monitor.stop() if monitor is not None else None if verbose: print( - f"Compiled and warmed up in {elapsed:.1f} s; " - "progress-bar ETA is now meaningful.", + "Warmup complete: " + f"{elapsed:.1f}s total (sampler {sampler_seconds:.1f}s, " + f"JIT log amplitudes {amplitude_seconds:.1f}s, energy + gradient " + f"{energy_seconds:.1f}s); {sample_summary}; " + f"stage 2 compiles a representative {amplitude_rows}-row forward " + "chunk. Stage 3 uses Hamiltonian-connected configurations and " + "compiles the backward pass used by VMC.", flush=True, ) + if resource_usage is not None: + print(resource_usage.summary("NetKet warmup resources"), flush=True) return elapsed diff --git a/src/pepsy/vmc/torch/__init__.py b/src/pepsy/vmc/torch/__init__.py index c685261..09b65e4 100644 --- a/src/pepsy/vmc/torch/__init__.py +++ b/src/pepsy/vmc/torch/__init__.py @@ -15,6 +15,11 @@ TorchPEPSBoundaryAmplitude, make_torch_peps_amplitude_model, ) +from .benchmark import ( + TorchAmplitudeBenchmark, + TorchAmplitudeBenchmarkRun, + benchmark_torch_amplitudes, +) from .connections import TorchConnections, compile_operator_sum_torch, torch_hamiltonian_connections from .driver import TorchVMCDriver from .fermion import ( @@ -44,6 +49,7 @@ TorchMetropolisResult, TorchImportanceSamples, TorchMCMCSamples, + TorchDistributedMetadata, TorchSampleProvenance, TorchChainDiagnostics, TorchVMCConvergenceEstimate, @@ -63,10 +69,13 @@ "TorchFermionVMCMetadata", "TorchPEPSAmplitude", "TorchPEPSBoundaryAmplitude", + "TorchAmplitudeBenchmark", + "TorchAmplitudeBenchmarkRun", "TorchConnections", "TorchMetropolisResult", "TorchImportanceSamples", "TorchMCMCSamples", + "TorchDistributedMetadata", "TorchSampleProvenance", "TorchChainDiagnostics", "TorchVMCConvergenceEstimate", @@ -84,6 +93,7 @@ "TorchSRResult", "TorchSquareLattice", "apply_torch_sr_update", + "benchmark_torch_amplitudes", "count_spinful_particles", "heisenberg_connections", "local_energy_from_connections", diff --git a/src/pepsy/vmc/torch/_core.py b/src/pepsy/vmc/torch/_core.py index ba96a86..64e403b 100644 --- a/src/pepsy/vmc/torch/_core.py +++ b/src/pepsy/vmc/torch/_core.py @@ -34,6 +34,7 @@ TorchVMCConvergenceReport, TorchImportanceSamples, TorchMCMCSamples, + TorchDistributedMetadata, TorchMetropolisResult, TorchSampleProvenance, TorchVMCImportanceEstimate, @@ -124,6 +125,7 @@ "TorchMetropolisResult", "TorchImportanceSamples", "TorchMCMCSamples", + "TorchDistributedMetadata", "TorchSampleProvenance", "TorchChainDiagnostics", "TorchVMCConvergenceEstimate", diff --git a/src/pepsy/vmc/torch/amplitude.py b/src/pepsy/vmc/torch/amplitude.py index 4094702..fdeb861 100644 --- a/src/pepsy/vmc/torch/amplitude.py +++ b/src/pepsy/vmc/torch/amplitude.py @@ -848,6 +848,11 @@ class TorchPEPSBoundaryAmplitude(TorchPEPSAmplitude): quimb PEPS using boundary-MPS environments around each parent walker. For a local update, only the touched row or column window is recontracted. + ``boundary_workers`` optionally evaluates independent cached-window + closures concurrently during no-grad CPU measurements. It defaults to one; + use a small value such as two or four only after checking for BLAS and + contraction-optimizer oversubscription. + Unsupported PEPS geometries or non-boundary contractions fall back to the base implementation. """ @@ -869,6 +874,7 @@ def __init__( boundary_cache_size=128, proposal_batching="auto", proposal_vmap_min_batch=8, + boundary_workers=1, ): super().__init__( peps, @@ -895,6 +901,14 @@ def __init__( "proposal_vmap_min_batch", proposal_vmap_min_batch, ) + # Native U1/U1U1 boundary contractions are not reliably vmappable. + # Independent cached-window closures can nevertheless be evaluated in + # parallel during no-grad CPU measurements. Keep this opt-in so the + # default remains deterministic and avoids BLAS/thread oversubscription. + self.boundary_workers = _check_positive_int( + "boundary_workers", + boundary_workers, + ) self._proposal_vmap_enabled = callable( getattr(_require_torch(), "vmap", None) ) @@ -1506,11 +1520,16 @@ def connected_amplitudes( torch = _require_torch() configs = _as_long_matrix(configs) amplitudes = torch.as_tensor(amplitudes, device=configs.device) + # A previous parent/fallback amplitude call must not be mistaken for + # the cache statistics of this connected-target measurement. + self.last_amplitude_cache_stats = None if connections.configs.numel() == 0: self.last_connected_reuse_stats = { + "num_requests": 0, "num_diagonal": 0, "num_reused": 0, "num_batched": 0, + "num_parallel": 0, "num_groups": 0, "num_grouped_connections": 0, "num_strip_cache_hits": 0, @@ -1534,9 +1553,11 @@ def connected_amplitudes( reuse_diagonal=reuse_diagonal, ) self.last_connected_reuse_stats = { + "num_requests": int(connections.configs.shape[0]), "num_diagonal": num_diagonal, "num_reused": 0, "num_batched": 0, + "num_parallel": 0, "num_groups": 0, "num_grouped_connections": 0, "num_strip_cache_hits": 0, @@ -1574,9 +1595,11 @@ def connected_amplitudes( finally: self._vmap_forward_enabled = previous_vmap_state self.last_connected_reuse_stats = { + "num_requests": int(connections.configs.shape[0]), "num_diagonal": int(diag.sum().item()), "num_reused": 0, "num_batched": int(offdiag.numel()), + "num_parallel": 0, "num_groups": 0, "num_grouped_connections": 0, "num_strip_cache_hits": 0, @@ -1598,9 +1621,11 @@ def connected_amplitudes( tn = self._unpack_tn() reference = self._reference_tensor() stats = { + "num_requests": int(connections.configs.shape[0]), "num_diagonal": int(diag.sum().item()), "num_reused": 0, "num_batched": 0, + "num_parallel": 0, "num_groups": 0, "num_grouped_connections": 0, "num_environment_cache_hits": 0, @@ -1616,6 +1641,7 @@ def connected_amplitudes( # PBC edges. They can then reuse one environment pair and one selected # parent-strip template while only their changed projectors differ. groups = {} + fallback_indices = [] for conn_idx_tensor in offdiag: conn_idx = int(conn_idx_tensor) parent_idx = int(connections.batch_ids[conn_idx].item()) @@ -1623,10 +1649,7 @@ def connected_amplitudes( target_config = connections.configs[conn_idx] windows = self._changed_axis_windows(parent_config, target_config) if not windows: - out[conn_idx] = self._contract_value( - self._select_config(tn, target_config), - reference, - ) + fallback_indices.append(conn_idx) stats["num_fallback"] += 1 continue axis, indices = windows[0] @@ -1686,6 +1709,11 @@ def get_context(parent_idx, parent_config, axis, indices): contexts[cache_key] = (envs, strip_tn) return contexts[cache_key] + # Build all primary contexts before dispatching work. This keeps cache + # mutation and boundary-environment construction on the caller thread; + # only independent scalar closures are allowed to run concurrently. + primary_jobs = [] + alternative_jobs = [] for (parent_idx, axis, indices), entries in groups.items(): parent_config = configs[parent_idx] primary_context = get_context( @@ -1695,67 +1723,144 @@ def get_context(parent_idx, parent_config, axis, indices): indices, ) for conn_idx, windows in entries: - target_config = connections.configs[conn_idx] - value_found = False - if primary_context is not None: - envs, strip_tn = primary_context - try: - out[conn_idx] = self._contract_cached_axis_window( - tn, - parent_config, - target_config, - axis, - indices, - envs, - strip_tn, - reference, - ) - except Exception: # pragma: no cover - upstream exceptions vary - pass - else: - value_found = True - stats["num_reused"] += 1 + job = ( + conn_idx, + parent_idx, + axis, + indices, + windows, + primary_context, + ) + if primary_context is None: + alternative_jobs.append(job) + else: + primary_jobs.append(job) + def contract_primary(job): + conn_idx, parent_idx, axis, indices, windows, context = job + parent_config = configs[parent_idx] + target_config = connections.configs[conn_idx] + envs, strip_tn = context + try: + value = self._contract_cached_axis_window( + tn, + parent_config, + target_config, + axis, + indices, + envs, + strip_tn, + reference, + ) + except Exception: # pragma: no cover - upstream exceptions vary + return job, None, False + return job, value, True + + def contract_primary_no_grad(job): + # Torch's grad mode is thread-local; explicitly carry the + # measurement mode into worker threads. + with torch.no_grad(): + return contract_primary(job) + + reference_device = getattr( + getattr(reference, "device", None), + "type", + None, + ) + boundary_workers = _check_positive_int( + "boundary_workers", + getattr(self, "boundary_workers", 1), + ) + use_parallel = ( + boundary_workers > 1 + and len(primary_jobs) > 1 + and not torch.is_grad_enabled() + and reference_device in (None, "cpu") + ) + if use_parallel: + # Threading is deliberately restricted to no-grad CPU inference. + # The PEPS parameters and cached environments remain shared, while + # every worker contracts into its own temporary TensorNetwork. + from concurrent.futures import ThreadPoolExecutor + + with ThreadPoolExecutor(max_workers=boundary_workers) as executor: + primary_results = executor.map( + contract_primary_no_grad, + primary_jobs, + ) + for job, value, value_found in primary_results: + if value_found: + conn_idx = job[0] + out[conn_idx] = value + stats["num_reused"] += 1 + else: + alternative_jobs.append(job) + stats["num_parallel"] = len(primary_jobs) + else: + for job in primary_jobs: + job, value, value_found = contract_primary(job) if value_found: - continue + conn_idx = job[0] + out[conn_idx] = value + stats["num_reused"] += 1 + else: + alternative_jobs.append(job) - for window_index, (alt_axis, alt_indices) in enumerate( - windows[1:], - start=1, - ): - context = get_context( - parent_idx, + # Alternative-axis retries stay serial because they may create new + # cached contexts. They are uncommon for ordinary nearest-neighbor + # updates and retain the existing robustness path. + for conn_idx, parent_idx, _axis, _indices, windows, _context in alternative_jobs: + parent_config = configs[parent_idx] + target_config = connections.configs[conn_idx] + value_found = False + for alt_axis, alt_indices in windows[1:]: + context = get_context( + parent_idx, + parent_config, + alt_axis, + alt_indices, + ) + if context is None: + continue + envs, strip_tn = context + try: + out[conn_idx] = self._contract_cached_axis_window( + tn, parent_config, + target_config, alt_axis, alt_indices, - ) - if context is None: - continue - envs, strip_tn = context - try: - out[conn_idx] = self._contract_cached_axis_window( - tn, - parent_config, - target_config, - alt_axis, - alt_indices, - envs, - strip_tn, - reference, - ) - except Exception: # pragma: no cover - upstream exceptions vary - continue - value_found = True - stats["num_reused"] += 1 - stats["num_alternative_axis_reused"] += 1 - break - - if not value_found: - out[conn_idx] = self._contract_value( - self._select_config(tn, target_config), + envs, + strip_tn, reference, ) - stats["num_fallback"] += 1 + except Exception: # pragma: no cover - upstream exceptions vary + continue + value_found = True + stats["num_reused"] += 1 + stats["num_alternative_axis_reused"] += 1 + break + + if not value_found: + fallback_indices.append(conn_idx) + stats["num_fallback"] += 1 + + if fallback_indices: + fallback_indices = torch.as_tensor( + fallback_indices, + dtype=torch.long, + device=configs.device, + ) + # Re-enter ``forward`` rather than directly contracting each + # target. In inference mode this deduplicates and consults the + # persistent boundary-amplitude cache; when vmap is available it + # can also evaluate unresolved targets as one fixed batch. + # Gradient-enabled calls deliberately retain ``forward``'s normal + # uncached differentiable path. + out[fallback_indices] = self.forward( + connections.configs[fallback_indices], + chunk_size=chunk_size, + ).to(dtype=out.dtype, device=out.device) self.last_connected_reuse_stats = stats return out diff --git a/src/pepsy/vmc/torch/benchmark.py b/src/pepsy/vmc/torch/benchmark.py new file mode 100644 index 0000000..6184493 --- /dev/null +++ b/src/pepsy/vmc/torch/benchmark.py @@ -0,0 +1,254 @@ +"""Reproducible amplitude throughput benchmarks for native Torch VMC.""" + +from __future__ import annotations + +import copy +from dataclasses import dataclass +import time + +from ..torch_types import _check_positive_int, _require_torch +from ._common import _model_device +from .amplitude import ( + _as_long_matrix, + _call_amplitude_fn, + _normalize_amplitude_batching, +) + + +@dataclass(frozen=True) +class TorchAmplitudeBenchmark: + """Timing for one amplitude batching/chunk-size configuration.""" + + amplitude_batching: str | None + executed_batching: str | None + chunk_size: int | None + n_configs: int + repeats: int + elapsed_seconds: float + configurations_per_second: float + + +@dataclass(frozen=True) +class TorchAmplitudeBenchmarkRun: + """Comparable timings for a fixed configuration batch. + + ``executed_batching`` records the path actually used by the amplitude + model. In particular, an unsupported ``"vmap"`` request is reported as + ``"serial"`` rather than being mistaken for a vectorized result. + """ + + entries: tuple[TorchAmplitudeBenchmark, ...] + device: str + n_configs: int + + @property + def best(self): + """Return the fastest tested entry, or ``None`` for an empty run.""" + if not self.entries: + return None + return min(self.entries, key=lambda entry: entry.elapsed_seconds) + + +def _synchronize_for_benchmark(device): + """Synchronize CUDA only when it affects wall-clock timing.""" + torch = _require_torch() + if device.type == "cuda": + torch.cuda.synchronize(device) + + +def _normalize_chunk_sizes(chunk_sizes): + try: + chunk_sizes = tuple(chunk_sizes) + except TypeError as exc: + raise TypeError("chunk_sizes must be an iterable of positive integers or None.") from exc + if not chunk_sizes: + raise ValueError("chunk_sizes must contain at least one entry.") + normalized = [] + for chunk_size in chunk_sizes: + if chunk_size is not None: + chunk_size = _check_positive_int("chunk_size", chunk_size) + if chunk_size not in normalized: + normalized.append(chunk_size) + return tuple(normalized) + + +def _batching_candidates(amplitude_fn, amplitude_batchings): + if amplitude_batchings is None: + if hasattr(amplitude_fn, "amplitude_batching"): + return ("serial", "auto", "vmap") + return (None,) + try: + candidates = tuple(amplitude_batchings) + except TypeError as exc: + raise TypeError("amplitude_batchings must be an iterable or None.") from exc + if not candidates: + raise ValueError("amplitude_batchings must contain at least one entry.") + if not hasattr(amplitude_fn, "amplitude_batching") and any( + candidate is not None for candidate in candidates + ): + raise TypeError( + "amplitude_batchings requires an amplitude model with an " + "amplitude_batching attribute." + ) + normalized = [] + for candidate in candidates: + candidate = ( + None + if candidate is None + else _normalize_amplitude_batching(candidate) + ) + if candidate not in normalized: + normalized.append(candidate) + return tuple(normalized) + + +def _benchmark_model_state(amplitude_fn): + """Capture mutable fast-path flags which a probe is allowed to change.""" + names = ( + "amplitude_batching", + "last_amplitude_batching", + "_vmap_forward_enabled", + "_vmap_log_enabled", + "_proposal_vmap_enabled", + "boundary_cache_size", + "last_amplitude_cache_stats", + ) + state = { + name: getattr(amplitude_fn, name) + for name in names + if hasattr(amplitude_fn, name) + } + if hasattr(amplitude_fn, "_boundary_amplitude_cache"): + state["_boundary_amplitude_cache"] = copy.copy( + amplitude_fn._boundary_amplitude_cache + ) + return state + + +def _restore_benchmark_model_state(amplitude_fn, state): + for name, value in state.items(): + setattr(amplitude_fn, name, value) + + +def _disable_boundary_amplitude_cache(amplitude_fn): + """Force the scalar contraction path while preserving cache state later.""" + if not hasattr(amplitude_fn, "boundary_cache_size"): + return + amplitude_fn.boundary_cache_size = 0 + cache = getattr(amplitude_fn, "_boundary_amplitude_cache", None) + if cache is not None: + amplitude_fn._boundary_amplitude_cache = type(cache)() + + +def benchmark_torch_amplitudes( + amplitude_fn, + configs, + *, + chunk_sizes=(None,), + amplitude_batchings=None, + warmup=1, + repeats=3, + verify=True, + include_cache=False, +): + """Benchmark native amplitude batching and chunk sizes on one batch. + + The model is evaluated under ``torch.no_grad()``. CUDA measurements are + synchronized around every timed region, and every candidate is checked + against the first result by default. By default, boundary-amplitude cache + hits are bypassed so the timing reflects contraction throughput rather than + previously retained samples; set ``include_cache=True`` to measure the + cache-aware serving path. Temporary vectorization probes and cache state + are restored before returning, so a failed benchmarked ``vmap`` attempt + cannot disable the normal path. + + This intentionally measures only the amplitude side of VMC. Use samples + retained by the target calculation (for example ``samples.configs``) to + benchmark representative PEPS configurations without another Markov pass. + """ + torch = _require_torch() + configs = _as_long_matrix(configs).to(device=_model_device(amplitude_fn)) + if configs.shape[0] == 0: + raise ValueError("configs must contain at least one configuration.") + if isinstance(warmup, bool) or not isinstance(warmup, int) or warmup < 0: + raise ValueError("warmup must be a non-negative integer.") + repeats = _check_positive_int("repeats", repeats) + if not isinstance(include_cache, bool): + raise TypeError("include_cache must be a bool.") + chunk_sizes = _normalize_chunk_sizes(chunk_sizes) + amplitude_batchings = _batching_candidates(amplitude_fn, amplitude_batchings) + original_state = _benchmark_model_state(amplitude_fn) + reference = None + entries = [] + try: + for amplitude_batching in amplitude_batchings: + for chunk_size in chunk_sizes: + _restore_benchmark_model_state(amplitude_fn, original_state) + if amplitude_batching is not None: + amplitude_fn.amplitude_batching = amplitude_batching + if not include_cache: + _disable_boundary_amplitude_cache(amplitude_fn) + with torch.no_grad(): + for _ in range(warmup): + _call_amplitude_fn( + amplitude_fn, + configs, + chunk_size=chunk_size, + ) + _synchronize_for_benchmark(configs.device) + started = time.perf_counter() + value = None + for _ in range(repeats): + value = _call_amplitude_fn( + amplitude_fn, + configs, + chunk_size=chunk_size, + ) + _synchronize_for_benchmark(configs.device) + elapsed = time.perf_counter() - started + value = torch.as_tensor(value, device=configs.device) + if tuple(value.shape) != (int(configs.shape[0]),): + raise ValueError( + "amplitude_fn must return one scalar amplitude per " + "configuration." + ) + if reference is None: + reference = value.detach().clone() + elif verify and not torch.allclose(value, reference): + raise RuntimeError( + "Amplitude benchmark candidates returned different " + "values; do not compare their throughput." + ) + entries.append( + TorchAmplitudeBenchmark( + amplitude_batching=amplitude_batching, + executed_batching=getattr( + amplitude_fn, + "last_amplitude_batching", + None, + ), + chunk_size=chunk_size, + n_configs=int(configs.shape[0]), + repeats=repeats, + elapsed_seconds=elapsed, + configurations_per_second=( + int(configs.shape[0]) * repeats / elapsed + if elapsed > 0 + else float("inf") + ), + ) + ) + finally: + _restore_benchmark_model_state(amplitude_fn, original_state) + return TorchAmplitudeBenchmarkRun( + entries=tuple(entries), + device=str(configs.device), + n_configs=int(configs.shape[0]), + ) + + +__all__ = [ + "TorchAmplitudeBenchmark", + "TorchAmplitudeBenchmarkRun", + "benchmark_torch_amplitudes", +] diff --git a/src/pepsy/vmc/torch/distributed.py b/src/pepsy/vmc/torch/distributed.py new file mode 100644 index 0000000..6f8616d --- /dev/null +++ b/src/pepsy/vmc/torch/distributed.py @@ -0,0 +1,199 @@ +"""Optional rank-sharded sampling helpers built on ``torch.distributed``. + +The native VMC sampler remains single-process by default. This module keeps +the distributed layer deliberately small: ranks own independent chains and +only compact scalar statistics are reduced after measurement. PEPS tensors and +sample configurations are never gathered. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from ..torch_types import _require_torch +from .results import TorchDistributedMetadata + + +@dataclass(frozen=True) +class _TorchDistributedRuntime: + """Live process-group handles, intentionally kept out of public results.""" + + module: object + group: object + rank: int + world_size: int + backend: str + + +def resolve_torch_distributed(distributed): + """Return the initialized Torch process group requested by ``distributed``.""" + if distributed is None or distributed is False: + return None + try: + import torch.distributed as dist + except ImportError as exc: # pragma: no cover - torch build dependent + raise RuntimeError( + "distributed=True requires a PyTorch build with torch.distributed." + ) from exc + if not dist.is_available(): + raise RuntimeError( + "distributed=True requires a PyTorch build with torch.distributed." + ) + if not dist.is_initialized(): + raise RuntimeError( + "distributed=True requires an initialized torch.distributed process " + "group. Initialize it before constructing or sampling the VMC run." + ) + group = None if distributed is True else distributed + rank = int(dist.get_rank(group=group)) + world_size = int(dist.get_world_size(group=group)) + if world_size < 1: # pragma: no cover - defensive against broken backends + raise RuntimeError("torch.distributed returned an invalid world size.") + return _TorchDistributedRuntime( + module=dist, + group=group, + rank=rank, + world_size=world_size, + backend=str(dist.get_backend(group=group)), + ) + + +def shard_chain_count(n_chains, runtime): + """Return this rank's deterministic contiguous share of global chains.""" + n_chains = int(n_chains) + if n_chains < runtime.world_size: + raise ValueError( + "The global n_chains must be at least the distributed world size " + "so every rank owns at least one Markov chain." + ) + base, extra = divmod(n_chains, runtime.world_size) + return base + int(runtime.rank < extra) + + +def rank_seed(seed, runtime): + """Derive reproducible, non-overlapping rank-local sampler streams.""" + if seed is None: + return None + # ``torch.Generator.manual_seed`` accepts a bounded integer. Keep every + # rank-local derivation in its portable positive range even when callers + # supplied a very large Python integer. + return (int(seed) + 104_729 * runtime.rank) % (2**63 - 1) + + +def distributed_metadata( + runtime, + *, + global_n_chains, + local_n_chains, + global_n_samples, + local_n_samples, +): + """Create a serializable distributed-run description for result records.""" + return TorchDistributedMetadata( + rank=runtime.rank, + world_size=runtime.world_size, + backend=runtime.backend, + global_n_chains=int(global_n_chains), + local_n_chains=int(local_n_chains), + global_n_samples=int(global_n_samples), + local_n_samples=int(local_n_samples), + ) + + +def _all_reduce(tensor, runtime, *, op): + runtime.module.all_reduce(tensor, op=op, group=runtime.group) + return tensor + + +def distributed_sum_int(value, runtime, *, device): + """Sum a Python integer across ranks without moving PEPS data.""" + torch = _require_torch() + total = torch.as_tensor(int(value), dtype=torch.int64, device=device) + if runtime.world_size > 1: + _all_reduce(total, runtime, op=runtime.module.ReduceOp.SUM) + return int(total.item()) + + +def distributed_max_float(value, runtime, *, device): + """Return the slowest rank's elapsed time for global throughput reporting.""" + torch = _require_torch() + maximum = torch.as_tensor(float(value), dtype=torch.float64, device=device) + if runtime.world_size > 1: + _all_reduce(maximum, runtime, op=runtime.module.ReduceOp.MAX) + return float(maximum.item()) + + +def distributed_unweighted_statistics( + local_values, + *, + local_effective_sample_size, + runtime, +): + """Reduce mean, variance, and local-chain ESS without gathering samples. + + ``local_effective_sample_size`` is obtained from each rank's independent + chain diagnostics. Its sum preserves local autocorrelation corrections, + while global R-hat is intentionally unavailable because configurations are + never all-gathered. + """ + torch = _require_torch() + local_values = torch.as_tensor(local_values).reshape(-1) + if local_values.numel() == 0: + raise ValueError("distributed measurement requires non-empty local samples.") + real_dtype = torch.float64 + real = local_values.real.to(dtype=real_dtype) + imag = ( + local_values.imag.to(dtype=real_dtype) + if local_values.is_complex() + else torch.zeros_like(real) + ) + moments = torch.stack( + ( + real.sum(), + imag.sum(), + local_values.abs().square().to(dtype=real_dtype).sum(), + torch.as_tensor( + float(local_values.numel()), + dtype=real_dtype, + device=local_values.device, + ), + torch.as_tensor( + local_effective_sample_size, + dtype=real_dtype, + device=local_values.device, + ), + ) + ) + if runtime.world_size > 1: + _all_reduce(moments, runtime, op=runtime.module.ReduceOp.SUM) + total = moments[3] + mean_real = moments[0] / total + mean_imag = moments[1] / total + if local_values.is_complex(): + mean = torch.complex(mean_real, mean_imag).to(dtype=local_values.dtype) + else: + mean = mean_real.to(dtype=local_values.dtype) + variance = torch.clamp( + moments[2] / total - mean_real.square() - mean_imag.square(), + min=0.0, + ) + effective_sample_size = torch.clamp(moments[4], min=1.0) + return ( + mean, + variance, + torch.sqrt(variance / effective_sample_size), + torch.sqrt(variance / total), + effective_sample_size, + int(total.item()), + ) + + +__all__ = [ + "distributed_max_float", + "distributed_metadata", + "distributed_sum_int", + "distributed_unweighted_statistics", + "rank_seed", + "resolve_torch_distributed", + "shard_chain_count", +] diff --git a/src/pepsy/vmc/torch/driver.py b/src/pepsy/vmc/torch/driver.py index 3e2bf50..13d9134 100644 --- a/src/pepsy/vmc/torch/driver.py +++ b/src/pepsy/vmc/torch/driver.py @@ -16,7 +16,17 @@ _resolve_log_amplitude_fn, _unique_config_rows, ) +from .benchmark import benchmark_torch_amplitudes from .connections import _driver_terms_connections, compile_operator_sum_torch +from .distributed import ( + distributed_max_float, + distributed_metadata, + distributed_sum_int, + distributed_unweighted_statistics, + rank_seed, + resolve_torch_distributed, + shard_chain_count, +) from .local_energy import ( _adaptive_measurement_options, _adaptive_thinning_interval, @@ -226,6 +236,42 @@ def refresh_amplitudes(self): self._refresh_log_amplitudes() return self.amplitudes + def benchmark_amplitudes( + self, + configs=None, + *, + chunk_sizes=(None,), + amplitude_batchings=None, + warmup=1, + repeats=3, + verify=True, + include_cache=False, + ): + """Benchmark chunked and vectorized amplitude evaluation. + + ``configs`` defaults to the driver's active walkers. Pass retained + chain samples with shape ``(n_steps, n_chains, n_sites)`` to time a + representative production batch without drawing more samples. Boundary + amplitude-cache hits are bypassed by default; pass + ``include_cache=True`` to time the cache-aware serving path instead. + """ + if configs is None: + configs = self.configs + else: + configs = _require_torch().as_tensor(configs, dtype=_require_torch().long) + if configs.ndim == 3: + configs = configs.reshape(-1, configs.shape[-1]) + return benchmark_torch_amplitudes( + self.model, + configs, + chunk_sizes=chunk_sizes, + amplitude_batchings=amplitude_batchings, + warmup=warmup, + repeats=repeats, + verify=verify, + include_cache=include_cache, + ) + def _refresh_log_amplitudes(self): if self.log_amplitude_fn is None: self.log_abs_amplitudes = None @@ -398,15 +444,20 @@ def sample_bp( n_samples = config_kwargs.pop("n_samples") n_chains = config_kwargs.pop("n_chains") if n_discard_per_chain is not None and n_discard_per_chain != config_kwargs["n_discard_per_chain"]: - raise ValueError("n_discard_per_chain conflicts with sampling.burn_in.") + raise ValueError( + "n_discard_per_chain conflicts with sampling.n_discard_per_chain." + ) if n_discard is not None and n_discard != config_kwargs["n_discard_per_chain"]: - raise ValueError("n_discard conflicts with sampling.burn_in.") - if sweep_size is not None and sweep_size != config_kwargs["n_thin"]: - raise ValueError("sweep_size conflicts with sampling.thin.") - if n_thin is not None and n_thin != config_kwargs["n_thin"]: - raise ValueError("n_thin conflicts with sampling.thin.") + raise ValueError( + "n_discard conflicts with sampling.n_discard_per_chain." + ) + if sweep_size is not None and sweep_size != config_kwargs["sweep_size"]: + raise ValueError("sweep_size conflicts with sampling.sweep_size.") + if n_thin is not None and n_thin != config_kwargs["sweep_size"]: + raise ValueError("n_thin conflicts with sampling.sweep_size.") n_discard_per_chain = config_kwargs["n_discard_per_chain"] - n_thin = config_kwargs["n_thin"] + sweep_size = config_kwargs["sweep_size"] + n_thin = None seed = config_kwargs["seed"] sampler_seed = config_kwargs["sampler_seed"] sampling_chunk_size = sampling.chunk_size @@ -452,11 +503,23 @@ def sample( sweep_size=None, n_thin=None, progress=False, + progress_postfix="full", seed=None, sampler_seed=None, track_proposal_stats=False, + distributed=False, ): - """Collect chain-preserving samples and update the driver state.""" + """Collect chain-preserving samples and update the driver state. + + Set ``distributed=True`` after initializing ``torch.distributed`` to + shard the *global* ``SamplingConfig.n_chains`` across ranks. Each rank + owns independent chains and returns only its local configurations; + metadata records the global chain/sample counts. Distributed sampling + requires ``sampling=...`` so the global chain semantics are explicit. + ``progress_postfix=\"acceptance\"`` keeps the live progress display + focused on the running sampling acceptance rate. + """ + distributed_runtime = resolve_torch_distributed(distributed) sampling_chunk_size = None sampling_proposal = None if sampling is not None: @@ -467,19 +530,58 @@ def sample( n_samples = config_kwargs.pop("n_samples") n_chains = config_kwargs.pop("n_chains") if n_discard_per_chain is not None and n_discard_per_chain != config_kwargs["n_discard_per_chain"]: - raise ValueError("n_discard_per_chain conflicts with sampling.burn_in.") + raise ValueError( + "n_discard_per_chain conflicts with sampling.n_discard_per_chain." + ) if n_discard is not None and n_discard != config_kwargs["n_discard_per_chain"]: - raise ValueError("n_discard conflicts with sampling.burn_in.") - if sweep_size is not None and sweep_size != config_kwargs["n_thin"]: - raise ValueError("sweep_size conflicts with sampling.thin.") - if n_thin is not None and n_thin != config_kwargs["n_thin"]: - raise ValueError("n_thin conflicts with sampling.thin.") + raise ValueError( + "n_discard conflicts with sampling.n_discard_per_chain." + ) + if sweep_size is not None and sweep_size != config_kwargs["sweep_size"]: + raise ValueError("sweep_size conflicts with sampling.sweep_size.") + if n_thin is not None and n_thin != config_kwargs["sweep_size"]: + raise ValueError("n_thin conflicts with sampling.sweep_size.") n_discard_per_chain = config_kwargs["n_discard_per_chain"] - n_thin = config_kwargs["n_thin"] + sweep_size = config_kwargs["sweep_size"] + n_thin = None seed = config_kwargs["seed"] sampler_seed = config_kwargs["sampler_seed"] sampling_chunk_size = sampling.chunk_size sampling_proposal = sampling.proposal + elif distributed_runtime is not None: + raise ValueError( + "distributed sampling requires sampling=SamplingConfig(...) so " + "n_chains has an unambiguous global meaning." + ) + + distributed_info = None + if distributed_runtime is not None: + global_n_chains = int(n_chains) + local_n_chains = shard_chain_count( + global_n_chains, + distributed_runtime, + ) + if self.n_walkers != local_n_chains: + raise ValueError( + "This rank's TorchVMCDriver has " + f"{self.n_walkers} walkers, but rank " + f"{distributed_runtime.rank} requires {local_n_chains} of " + f"the global {global_n_chains} chains. Initialize each rank " + "with its local shard before calling sample(..., " + "distributed=True)." + ) + if seed is not None: + seed = rank_seed(seed, distributed_runtime) + if sampler_seed is not None: + sampler_seed = rank_seed(sampler_seed, distributed_runtime) + if seed is None and sampler_seed is None: + raise ValueError( + "distributed sampling requires SamplingConfig.seed or " + "SamplingConfig.sampler_seed so ranks do not duplicate " + "their random streams." + ) + n_chains = local_n_chains + n_samples = int(sampling.n_samples_per_chain) * local_n_chains sampler = self.make_sampler( n_chains=n_chains, seed=seed, @@ -494,6 +596,7 @@ def sample( sweep_size=sweep_size, n_thin=n_thin, progress=progress, + progress_postfix=progress_postfix, track_proposal_stats=track_proposal_stats, ) self.configs = sampler.configs @@ -501,6 +604,20 @@ def sample( self.generator = sampler.generator if track_proposal_stats: self.last_proposal_stats = result.proposal_stats + if distributed_runtime is not None: + global_n_samples = distributed_sum_int( + result.n_samples, + distributed_runtime, + device=self.configs.device, + ) + distributed_info = distributed_metadata( + distributed_runtime, + global_n_chains=global_n_chains, + local_n_chains=local_n_chains, + global_n_samples=global_n_samples, + local_n_samples=result.n_samples, + ) + result = replace(result, distributed=distributed_info) return result def make_connections(self, configs=None, *, terms=None): @@ -1091,6 +1208,7 @@ def measure_samples( profile=False, deduplicate=True, progress=False, + distributed=None, ): """Measure saved chain samples without running another sampler. @@ -1134,6 +1252,21 @@ def measure_samples( model_device = _model_device(self.model) sample_object = samples if hasattr(samples, "configs") else None + sample_distributed = getattr(sample_object, "distributed", None) + if distributed is None: + distributed = sample_distributed is not None + distributed_runtime = resolve_torch_distributed(distributed) + if distributed_runtime is not None and sample_distributed is not None: + if ( + sample_distributed.rank != distributed_runtime.rank + or sample_distributed.world_size != distributed_runtime.world_size + ): + raise RuntimeError( + "The supplied samples belong to a different distributed " + "rank layout than the active torch.distributed process group." + ) + if distributed_runtime is not None: + progress = bool(progress) and distributed_runtime.rank == 0 provenance = getattr(sample_object, "provenance", None) if provenance is not None and provenance != _torch_sample_provenance(self.model): raise RuntimeError( @@ -1171,6 +1304,15 @@ def measure_samples( if n_steps <= 0 or n_chains <= 0 or n_sites <= 0: raise ValueError("samples must contain at least one configuration.") flat_configs = chain_configs.reshape(-1, n_sites) + global_n_chains = ( + distributed_sum_int( + n_chains, + distributed_runtime, + device=model_device, + ) + if distributed_runtime is not None + else n_chains + ) unique_parent_count = ( int(_unique_config_rows(flat_configs)[0].shape[0]) if deduplicate @@ -1252,6 +1394,12 @@ def measure_samples( ) else: importance_weights = None + if distributed_runtime is not None and importance_weights is not None: + raise NotImplementedError( + "Distributed measurement currently supports only unweighted " + "rank-sharded Markov samples. Importance weights require a " + "global normalization and are not reduced implicitly." + ) if observables is None: observable_items = (("observable", None),) @@ -1327,6 +1475,23 @@ def set_phase(stage, *, n_connections=None): ) if sample_object is not None else 0.0 n_proposed = int(getattr(sample_object, "n_proposed", 0)) if sample_object is not None else 0 n_accepted = int(getattr(sample_object, "n_accepted", 0)) if sample_object is not None else 0 + if distributed_runtime is not None: + elapsed = distributed_max_float( + elapsed, + distributed_runtime, + device=model_device, + ) + n_proposed = distributed_sum_int( + n_proposed, + distributed_runtime, + device=model_device, + ) + n_accepted = distributed_sum_int( + n_accepted, + distributed_runtime, + device=model_device, + ) + acceptance_rate = n_accepted / n_proposed if n_proposed else 0.0 profile_data = None if profile: profile_data = { @@ -1362,10 +1527,34 @@ def set_phase(stage, *, n_connections=None): if importance_weights is None else (*_weighted_energy_statistics(flat_values[name], importance_weights), None) ) + distributed_info = None + if distributed_runtime is not None: + ( + energy_mean, + energy_variance, + energy_stderr, + energy_stderr_naive, + effective_sample_size, + global_n_samples, + ) = distributed_unweighted_statistics( + local_values, + local_effective_sample_size=effective_sample_size, + runtime=distributed_runtime, + ) + chain_diagnostics = None + distributed_info = distributed_metadata( + distributed_runtime, + global_n_chains=global_n_chains, + local_n_chains=n_chains, + global_n_samples=global_n_samples, + local_n_samples=int(local_values.numel()), + ) result_profile = None if profile_data is not None: result_profile = dict(profile_data) result_profile["observable"] = name + if distributed_info is not None: + result_profile["distributed"] = distributed_info results[name] = TorchVMCEnergyEstimate( configs=chain_configs, amplitudes=chain_amplitudes, @@ -1376,11 +1565,19 @@ def set_phase(stage, *, n_connections=None): acceptance_rate=acceptance_rate, n_proposed=n_proposed, n_accepted=n_accepted, - n_samples=int(local_values.numel()), + n_samples=( + int(local_values.numel()) + if distributed_info is None + else distributed_info.global_n_samples + ), n_measurements=n_steps, elapsed_seconds=elapsed, samples_per_second=( - int(local_values.numel()) / elapsed + ( + int(local_values.numel()) + if distributed_info is None + else distributed_info.global_n_samples + ) / elapsed if elapsed > 0 else float("inf") ), @@ -1404,6 +1601,7 @@ def set_phase(stage, *, n_connections=None): name="proposal_log_probs", ).reshape(n_steps, n_chains) ), + distributed=distributed_info, ) if phase_bar is not None: diff --git a/src/pepsy/vmc/torch/fermion.py b/src/pepsy/vmc/torch/fermion.py index 450a3a0..c24d677 100644 --- a/src/pepsy/vmc/torch/fermion.py +++ b/src/pepsy/vmc/torch/fermion.py @@ -19,6 +19,7 @@ make_torch_peps_amplitude_model, ) from .connections import compile_operator_sum_torch, _normalize_terms_site_labels +from .distributed import rank_seed, resolve_torch_distributed, shard_chain_count from .driver import TorchVMCDriver from .metadata import _infer_torch_fermion_metadata from .results import TorchVMCMeasurementRun, TorchVMCWarmupResult @@ -353,6 +354,7 @@ def __init__( log_amplitude_fn=None, proposal_batching="auto", proposal_vmap_min_batch=8, + boundary_workers=1, generator=None, seed=None, amplitude_floor=0.0, @@ -392,13 +394,10 @@ def __init__( from ..api import OperatorSum if terms is None: - if fermion is None: - raise ValueError( - "Pass fermion=... when terms are omitted so the default " - "Hamiltonian can be constructed." - ) - hamiltonian = fermion.hamiltonian(metadata.edges) - terms = hamiltonian.terms + raise ValueError( + "Pass explicit hamiltonian=... or terms=.... Fermion stores " + "local symmetry conventions, not t/U/V/mu couplings." + ) elif isinstance(terms, OperatorSum): hamiltonian = terms terms = compile_operator_sum_torch( @@ -444,6 +443,10 @@ def __init__( "amplitude_batching": amplitude_batching, "proposal_batching": proposal_batching, "proposal_vmap_min_batch": proposal_vmap_min_batch, + "boundary_workers": _check_positive_int( + "boundary_workers", + boundary_workers, + ), } self._driver_options = { "proposal": proposal, @@ -467,6 +470,7 @@ def _ensure_initialized( contraction=None, contraction_opts=None, n_walkers=None, + initialization_seed=None, ): """Initialize the native driver once, from the measurement recipe. @@ -524,12 +528,12 @@ def _ensure_initialized( ) return - self._initialize_driver( - requested_contraction, - n_walkers=requested_n_walkers, - ) + initialization_kwargs = {"n_walkers": requested_n_walkers} + if initialization_seed is not None: + initialization_kwargs["initialization_seed"] = initialization_seed + self._initialize_driver(requested_contraction, **initialization_kwargs) - def _initialize_driver(self, contraction, *, n_walkers): + def _initialize_driver(self, contraction, *, n_walkers, initialization_seed=None): """Build the amplitude model and initial walkers for a first run.""" torch = _require_torch() model_kwargs = { @@ -546,6 +550,7 @@ def _initialize_driver(self, contraction, *, n_walkers): proposal_vmap_min_batch=self._model_options[ "proposal_vmap_min_batch" ], + boundary_workers=self._model_options["boundary_workers"], ) model = make_torch_peps_amplitude_model(self.peps, **model_kwargs) model_device = _model_device( @@ -554,12 +559,14 @@ def _initialize_driver(self, contraction, *, n_walkers): ) generator = self._initial_generator - if self._initial_seed is not None: + if initialization_seed is None: + initialization_seed = self._initial_seed + if initialization_seed is not None: try: generator = torch.Generator(device=model_device) except (RuntimeError, TypeError, ValueError): generator = torch.Generator() - generator.manual_seed(int(self._initial_seed)) + generator.manual_seed(int(initialization_seed)) metadata = self.metadata configs = self._initial_configs @@ -719,8 +726,8 @@ def _sampling_estimator_kwargs(sampling, kwargs): "n_chains": configured["n_chains"], "n_discard_per_chain": configured["n_discard_per_chain"], "n_discard": configured["n_discard_per_chain"], - "sweep_size": configured["n_thin"], - "n_thin": configured["n_thin"], + "sweep_size": configured["sweep_size"], + "n_thin": configured["sweep_size"], "seed": configured["seed"], "sampler_seed": configured["sampler_seed"], } @@ -768,6 +775,7 @@ def sample( contraction=None, contraction_opts=None, proposal=None, + distributed=False, **kwargs, ): """Collect reusable Markov or external-proposal samples. @@ -785,6 +793,11 @@ def sample( ``n_samples=...`` for that path. """ if proposal is not None: + if distributed: + raise NotImplementedError( + "Distributed sampling currently supports native Markov " + "chains, not external proposal/importance samples." + ) if sampling is not None: raise ValueError( "sampling= describes target-Metropolis burn-in and " @@ -819,13 +832,27 @@ def sample( progress=progress, amplitude_floor=amplitude_floor, ) + distributed_runtime, initialization_sampling = self._rank_sharded_sampling_config( + sampling, + distributed, + ) + initialization_seed = ( + self._sampling_seed(initialization_sampling) + if distributed_runtime is not None + else None + ) self._ensure_initialized( - sampling=sampling, + sampling=initialization_sampling, contraction=contraction, contraction_opts=contraction_opts, n_walkers=kwargs.get("n_chains"), + initialization_seed=initialization_seed, + ) + return super().sample( + sampling=sampling, + distributed=distributed, + **kwargs, ) - return super().sample(sampling=sampling, **kwargs) def check_mc_convergence( self, @@ -851,6 +878,51 @@ def check_mc_convergence( ) return super().check_mc_convergence(compiled, **kwargs) + @staticmethod + def _rank_sharded_sampling_config(sampling, distributed): + """Return the rank-local recipe required for lazy driver setup.""" + distributed_runtime = resolve_torch_distributed(distributed) + if distributed_runtime is None: + return None, sampling + if sampling is None: + raise ValueError( + "distributed sampling requires sampling=SamplingConfig(...) so " + "n_chains has an unambiguous global meaning." + ) + from ..api import SamplingConfig + if not isinstance(sampling, SamplingConfig): + raise TypeError("sampling must be a SamplingConfig or None.") + rank_local_seed = ( + rank_seed(sampling.seed, distributed_runtime) + if sampling.seed is not None + else None + ) + rank_local_sampler_seed = ( + rank_seed(sampling.sampler_seed, distributed_runtime) + if sampling.sampler_seed is not None + else None + ) + return distributed_runtime, replace( + sampling, + n_chains=shard_chain_count( + sampling.n_chains, + distributed_runtime, + ), + seed=rank_local_seed, + sampler_seed=rank_local_sampler_seed, + ) + + @staticmethod + def _sampling_seed(sampling): + """Return the configured sampler seed, if the recipe has one.""" + if sampling is None: + return None + return ( + sampling.seed + if sampling.seed is not None + else sampling.sampler_seed + ) + def measure( self, samples, @@ -862,6 +934,7 @@ def measure( profile=False, deduplicate=True, progress=False, + distributed=None, _include_energy=False, ): """Measure observables from retained samples without resampling. @@ -876,15 +949,20 @@ def measure( observables, include_energy=_include_energy or observables is None, ) + measure_kwargs = { + "observables": compiled, + "amplitudes": amplitudes, + "weights": weights, + "proposal_log_probs": proposal_log_probs, + "profile": profile, + "deduplicate": deduplicate, + "progress": progress, + } + if distributed is not None: + measure_kwargs["distributed"] = distributed return self.measure_samples( samples, - observables=compiled, - amplitudes=amplitudes, - weights=weights, - proposal_log_probs=proposal_log_probs, - profile=profile, - deduplicate=deduplicate, - progress=progress, + **measure_kwargs, ) def warmup( @@ -941,6 +1019,7 @@ def run_measurement( warmup_sweeps=0, progress=False, profile=False, + distributed=False, ): """Warm up, sample, and estimate PEPS Fermion observables once. @@ -949,23 +1028,39 @@ def run_measurement( observable estimates. ``progress=True`` reports optional burn-in, MCMC sampling, then the connection/contraction/statistics phases. """ + distributed_runtime, initialization_sampling = self._rank_sharded_sampling_config( + sampling, + distributed, + ) + initialization_seed = ( + self._sampling_seed(initialization_sampling) + if distributed_runtime is not None + else None + ) self._ensure_initialized( - sampling=sampling, + sampling=initialization_sampling, contraction=contraction, contraction_opts=contraction_opts, + initialization_seed=initialization_seed, ) start = time.perf_counter() + rank_progress = bool(progress) and ( + distributed_runtime is None or distributed_runtime.rank == 0 + ) warmup_result = ( - self.warmup(n_sweeps=warmup_sweeps, progress=progress) + self.warmup(n_sweeps=warmup_sweeps, progress=rank_progress) if warmup else None ) - samples = self.sample(sampling=sampling, progress=progress) + sample_kwargs = {"sampling": sampling, "progress": rank_progress} + if distributed_runtime is not None: + sample_kwargs["distributed"] = distributed + samples = self.sample(**sample_kwargs) estimates = self.measure( samples, observables=observables, profile=profile, - progress=progress, + progress=rank_progress, _include_energy=True, ) return TorchVMCMeasurementRun( @@ -986,6 +1081,7 @@ def run( warmup=None, warmup_sweeps=0, progress=False, + distributed=False, **kwargs, ): """Run either a PEPS measurement workflow or optimization updates. @@ -1014,11 +1110,12 @@ def run( or contraction is not None or contraction_opts is not None or warmup_sweeps != 0 + or not (distributed is False or distributed is None) ): raise ValueError( - "observables, sampling, contraction settings, and " - "warmup_sweeps apply only to measurement runs; omit " - "n_steps to use them." + "observables, sampling, contraction settings, warmup_sweeps, " + "and distributed apply only to measurement runs; omit n_steps " + "to use them." ) if warmup is not None: raise ValueError("warmup applies only to a measurement run.") @@ -1037,6 +1134,7 @@ def run( warmup_sweeps=warmup_sweeps, progress=progress, profile=profile, + distributed=distributed, ) def make_bp_sampler( diff --git a/src/pepsy/vmc/torch/local_energy.py b/src/pepsy/vmc/torch/local_energy.py index 79188b8..e82cc27 100644 --- a/src/pepsy/vmc/torch/local_energy.py +++ b/src/pepsy/vmc/torch/local_energy.py @@ -371,7 +371,14 @@ def _connected_amplitudes_with_target_dedup( reuse_diagonal=True, deduplicate_targets=False, ): - """Evaluate connected amplitudes, optionally sharing target rows globally.""" + """Evaluate connected amplitudes, optionally sharing target rows globally. + + A connected target can also be one of the parent configurations in the + retained batch. In that case its amplitude is already known, even when + the corresponding connection belongs to a different walker. Reusing that + value avoids an otherwise unnecessary PEPS contraction before dispatching + the remaining unique targets to the amplitude model. + """ torch = _require_torch() if not deduplicate_targets or connections.configs.shape[0] <= 1: return _connected_amplitudes_for_connections( @@ -383,16 +390,58 @@ def _connected_amplitudes_with_target_dedup( reuse_diagonal=reuse_diagonal, ) - target_configs, target_inverse = _unique_config_rows(connections.configs) - if target_inverse is None: # pragma: no cover - guarded by shape - target_inverse = torch.zeros( + # Deduplicate parent and target configurations together. Besides avoiding + # a second expensive row comparison, this supplies a global parent lookup + # for target configurations that another retained walker has already + # evaluated. + all_configs = torch.cat((configs, connections.configs), dim=0) + unique_configs, all_inverse = _unique_config_rows(all_configs) + if all_inverse is None: # pragma: no cover - guarded by shape + all_inverse = torch.zeros( 1, dtype=torch.long, device=configs.device, ) - # Pick one parent for each unique target. The target amplitude is - # independent of its parent, while the representative parent still lets - # boundary backends reuse the appropriate environment. + n_parents = configs.shape[0] + parent_keys = all_inverse[:n_parents] + connection_keys = all_inverse[n_parents:] + target_keys, target_inverse = torch.unique( + connection_keys, + sorted=False, + return_inverse=True, + ) + target_configs = unique_configs[target_keys] + + # Associate every globally unique configuration with any retained parent + # carrying it. The particular parent is immaterial because equal + # configurations have equal amplitudes at the current model state. + parent_order = torch.argsort(parent_keys) + sorted_parent_keys = parent_keys[parent_order] + first_parent_key = torch.ones( + sorted_parent_keys.shape[0], + dtype=torch.bool, + device=configs.device, + ) + if first_parent_key.numel() > 1: + first_parent_key[1:] = ( + sorted_parent_keys[1:] != sorted_parent_keys[:-1] + ) + parent_for_key = torch.full( + (unique_configs.shape[0],), + -1, + dtype=torch.long, + device=configs.device, + ) + parent_for_key[sorted_parent_keys[first_parent_key]] = parent_order[ + first_parent_key + ] + target_parent_indices = parent_for_key[target_keys] + reusable_targets = target_parent_indices >= 0 + + # Pick one parent for each unique target that is not already present in + # the parent batch. The target amplitude is independent of its parent, + # while the representative parent still lets boundary backends reuse the + # appropriate environment. order = torch.argsort(target_inverse) sorted_inverse = target_inverse[order] first = torch.ones( @@ -403,23 +452,41 @@ def _connected_amplitudes_with_target_dedup( if first.numel() > 1: first[1:] = sorted_inverse[1:] != sorted_inverse[:-1] representative = order[first] - unique_connections = TorchConnections( - configs=target_configs, - coeffs=torch.ones( - target_configs.shape[0], - dtype=amplitudes.dtype, - device=configs.device, - ), - batch_ids=connections.batch_ids[representative], + unique_amplitudes = torch.empty( + target_configs.shape[0], + dtype=amplitudes.dtype, + device=configs.device, ) - unique_amplitudes = _connected_amplitudes_for_connections( - configs, - amplitudes, - unique_connections, - amplitude_fn, - chunk_size=chunk_size, - reuse_diagonal=reuse_diagonal, + if reuse_diagonal and bool(torch.any(reusable_targets)): + unique_amplitudes[reusable_targets] = amplitudes[ + target_parent_indices[reusable_targets] + ] + + unreused_targets = ~reusable_targets if reuse_diagonal else torch.ones( + target_configs.shape[0], + dtype=torch.bool, + device=configs.device, ) + if bool(torch.any(unreused_targets)): + unique_connections = TorchConnections( + configs=target_configs[unreused_targets], + coeffs=torch.ones( + int(unreused_targets.sum().item()), + dtype=amplitudes.dtype, + device=configs.device, + ), + batch_ids=connections.batch_ids[ + representative[unreused_targets] + ], + ) + unique_amplitudes[unreused_targets] = _connected_amplitudes_for_connections( + configs, + amplitudes, + unique_connections, + amplitude_fn, + chunk_size=chunk_size, + reuse_diagonal=reuse_diagonal, + ) return unique_amplitudes[target_inverse] diff --git a/src/pepsy/vmc/torch/results.py b/src/pepsy/vmc/torch/results.py index b5ccb00..f4f1701 100644 --- a/src/pepsy/vmc/torch/results.py +++ b/src/pepsy/vmc/torch/results.py @@ -64,6 +64,25 @@ class TorchSampleProvenance: contraction_signature: tuple[Any, Any, Any, Any] +@dataclass(frozen=True) +class TorchDistributedMetadata: + """Rank-sharding information attached to local distributed VMC results. + + PEPS tensors and configurations remain rank-local. Unweighted measurement + estimates reduce moments and rank-local effective sample sizes, but global + chain diagnostics (notably R-hat) are intentionally unavailable without an + all-gather of Markov histories. + """ + + rank: int + world_size: int + backend: str + global_n_chains: int + local_n_chains: int + global_n_samples: int + local_n_samples: int + + def _torch_sample_provenance(model): """Capture the mutable model state relevant to stored MCMC amplitudes.""" parameters = getattr(model, "parameters", None) @@ -117,6 +136,7 @@ class TorchMCMCSamples: log_abs_amplitudes: Any = None proposal_stats: Any = None provenance: TorchSampleProvenance | None = None + distributed: TorchDistributedMetadata | None = None def diagnostics(self, values=None, *, max_lag=None, split_rhat=False): """Compute chain diagnostics for a scalar observable. @@ -139,6 +159,25 @@ def to_common(self): """Convert to the backend-neutral :class:`pepsy.vmc.VMCSamples`.""" from ..api import VMCSamples + diagnostics = { + "n_samples": self.n_samples, + "n_discard_per_chain": self.n_discard_per_chain, + "sweep_size": self.sweep_size, + "n_proposed": self.n_proposed, + "n_accepted": self.n_accepted, + "elapsed_seconds": self.elapsed_seconds, + "samples_per_second": self.samples_per_second, + } + if self.distributed is not None: + diagnostics["distributed"] = { + "rank": self.distributed.rank, + "world_size": self.distributed.world_size, + "backend": self.distributed.backend, + "global_n_chains": self.distributed.global_n_chains, + "local_n_chains": self.distributed.local_n_chains, + "global_n_samples": self.distributed.global_n_samples, + "local_n_samples": self.distributed.local_n_samples, + } return VMCSamples( configs=self.configs, amplitudes=self.amplitudes, @@ -146,15 +185,7 @@ def to_common(self): n_samples_per_chain=self.n_samples_per_chain, n_chains=self.n_chains, acceptance_rate=self.acceptance_rate, - diagnostics={ - "n_samples": self.n_samples, - "n_discard_per_chain": self.n_discard_per_chain, - "sweep_size": self.sweep_size, - "n_proposed": self.n_proposed, - "n_accepted": self.n_accepted, - "elapsed_seconds": self.elapsed_seconds, - "samples_per_second": self.samples_per_second, - }, + diagnostics=diagnostics, native=self, ) @@ -364,6 +395,7 @@ class TorchVMCEnergyEstimate: effective_sample_size: Any = None importance_weights: Any = None proposal_log_probs: Any = None + distributed: TorchDistributedMetadata | None = None @dataclass(frozen=True) @@ -504,15 +536,20 @@ def _connected_target_progress(model): if not stats: return {} fields = {} + requests = int(stats.get("num_requests", 0)) diagonal = int(stats.get("num_diagonal", 0)) reused = int(stats.get("num_reused", 0)) batched = int(stats.get("num_batched", 0)) + parallel = int(stats.get("num_parallel", 0)) fallback = int(stats.get("num_fallback", 0)) if diagonal or reused or batched or fallback: fields["targets"] = ( - f"diag={diagonal}, env={reused}, " + (f"req={requests}, " if requests else "") + + f"diag={diagonal}, env={reused}, " f"batch={batched}, direct={fallback}" ) + if parallel: + fields["target-workers"] = parallel environment_hits = int(stats.get("num_environment_cache_hits", 0)) environment_builds = int(stats.get("num_environment_builds", 0)) if environment_hits or environment_builds: @@ -541,10 +578,23 @@ def _set_vmc_progress_postfix( burn_in=None, thin=None, phase=None, + progress_postfix="full", ): """Update a Metropolis/VMC bar without affecting numerical work.""" if bar is None: return + if progress_postfix not in {"full", "acceptance"}: + raise ValueError( + "progress_postfix must be 'full' or 'acceptance'." + ) + if progress_postfix == "acceptance": + postfix = {} + if result is not None: + postfix["accept"] = f"{result.acceptance_rate:.3f}" + set_postfix = getattr(bar, "set_postfix", None) + if callable(set_postfix): + set_postfix(postfix) + return postfix = _model_progress_fields(model) if n_chains is not None: postfix["walkers"] = int(n_chains) @@ -640,6 +690,7 @@ def _accumulate_cache_profile(total, snapshot): __all__ = [ "TorchChainDiagnostics", + "TorchDistributedMetadata", "TorchImportanceSamples", "TorchMCMCSamples", "TorchMetropolisResult", diff --git a/src/pepsy/vmc/torch/sampler.py b/src/pepsy/vmc/torch/sampler.py index b7cf8ee..0bd5712 100644 --- a/src/pepsy/vmc/torch/sampler.py +++ b/src/pepsy/vmc/torch/sampler.py @@ -397,6 +397,7 @@ def sample( sweep_size=None, n_thin=None, progress=False, + progress_postfix="full", track_proposal_stats=False, ): """Discard and collect chain-preserving Metropolis samples. @@ -407,7 +408,10 @@ def sample( ``n_discard_per_chain`` and ``sweep_size`` respectively. Both the discard and retained portions use that sweep interval, so the progress bar totals ``(n_discard_per_chain + n_samples_per_chain) * - sweep_size`` batched Metropolis sweeps. + sweep_size`` batched Metropolis sweeps. Set + ``progress_postfix=\"acceptance\"`` to keep the live postfix focused on + the running acceptance rate; the default ``\"full\"`` shows sampler and + cache details as well. """ torch = _require_torch() n_samples = _check_positive_int("n_samples", n_samples) @@ -426,6 +430,10 @@ def sample( n_discard_per_chain, ) sweep_size = _check_positive_int("sweep_size", sweep_size) + if progress_postfix not in {"full", "acceptance"}: + raise ValueError( + "progress_postfix must be 'full' or 'acceptance'." + ) n_samples_per_chain = ( n_samples + self.n_chains - 1 ) // self.n_chains @@ -450,6 +458,7 @@ def sample( thin=sweep_size, phase=("equilibrate" if n_discard_per_chain else "retain 0/" f"{n_samples_per_chain}"), + progress_postfix=progress_postfix, ) start = time.perf_counter() n_proposed = 0 @@ -488,6 +497,7 @@ def advance_one_sweep(): burn_in=n_discard_per_chain, thin=sweep_size, phase=sampling_phase(), + progress_postfix=progress_postfix, ) for _ in range(n_discard_per_chain * sweep_size): diff --git a/tests/test_backends.py b/tests/test_backends.py index 095588c..fbd28cd 100644 --- a/tests/test_backends.py +++ b/tests/test_backends.py @@ -72,3 +72,41 @@ def test_to_float_handles_torch_scalar_if_available(): value = torch.tensor(3.5 + 1.25j, dtype=torch.complex128, requires_grad=True) assert pepsy.to_float(value) == pytest.approx(3.5) + + +@pytest.mark.parametrize("complex_input", (False, True)) +def test_jax_svd_vjp_restores_quimb_truncated_cotangents(complex_input): + """Fixed-rank contractions must not mix full and truncated SVD axes.""" + jax = pytest.importorskip("jax") + import jax.numpy as jnp + + from pepsy.backends.linalg_jax import jaxsvd_bwd, jaxsvd_fwd + + matrix = jnp.arange(64, dtype=jnp.float32).reshape(8, 8) + jnp.eye(8) + if complex_input: + matrix = matrix * (1.0 + 0.1j) + + outputs, residual = jaxsvd_fwd(matrix) + U, S, Vh = outputs + retained = 2 + u_tangent = 0.25 - 0.1j if complex_input else 0.25 + vh_tangent = -0.2 + 0.3j if complex_input else -0.2 + truncated_tangents = ( + jnp.full_like(U[:, :retained], u_tangent), + jnp.full_like(S[:retained], 0.5), + jnp.full_like(Vh[:retained, :], vh_tangent), + ) + actual = jaxsvd_bwd(residual, truncated_tangents)[0] + + full_tangents = type(outputs)( + jnp.zeros_like(U).at[:, :retained].set(truncated_tangents[0]), + jnp.zeros_like(S).at[:retained].set(truncated_tangents[1]), + jnp.zeros_like(Vh).at[:retained, :].set(truncated_tangents[2]), + ) + _, pullback = jax.vjp( + lambda value: jnp.linalg.svd(value, full_matrices=False), + matrix, + ) + expected = pullback(full_tangents)[0] + + np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-5) diff --git a/tests/test_fermion_gate_cache.py b/tests/test_fermion_gate_cache.py index 3d6e143..6713abd 100644 --- a/tests/test_fermion_gate_cache.py +++ b/tests/test_fermion_gate_cache.py @@ -27,7 +27,7 @@ def _dense(gate): def test_fingerprint_is_content_addressed_not_identity(): """Equal contents share a fingerprint; different contents do not.""" - fermion = Fermion(spinful=True, symmetry="U1", t=1.0, U=8.0) + fermion = Fermion(spinful=True, symmetry="U1") up_a = fermion.hopping_operator(spin="up") up_b = fermion.hopping_operator(spin="up") @@ -53,7 +53,7 @@ def test_operator_gate_does_not_alias_distinct_operators_under_id_reuse(): of different kinds at the same ``theta`` used to let a recycled ``id`` return a stale cached gate. Every gate must match an independent exponential. """ - fermion = Fermion(spinful=True, symmetry="U1", t=1.0, U=8.0) + fermion = Fermion(spinful=True, symmetry="U1") theta = 0.10667747 references = { @@ -78,7 +78,7 @@ def test_operator_gate_does_not_alias_distinct_operators_under_id_reuse(): def test_operator_gate_cache_hits_return_equivalent_gate(): """Repeated calls with equal contents reuse a single correct gate.""" - fermion = Fermion(spinful=True, symmetry="U1", t=1.0, U=8.0) + fermion = Fermion(spinful=True, symmetry="U1") theta = 0.37 reference = _dense(_gate_from_term(fermion.hopping_operator(spin="up"), theta)) diff --git a/tests/test_netket_flat_z2.py b/tests/test_netket_flat_z2.py new file mode 100644 index 0000000..3fb71b7 --- /dev/null +++ b/tests/test_netket_flat_z2.py @@ -0,0 +1,437 @@ +"""Regression tests for NetKet's flat-Z2 PEPS and warmup reporting.""" + +from types import SimpleNamespace + +import numpy as np +import pytest + + +@pytest.mark.smoke +def test_prepare_fermionic_peps_flattens_odd_z2_bond_dimension(): + """Odd D creates unequal Z2 blocks that must be padded before JAX VMC.""" + jax = pytest.importorskip("jax") + jnp = pytest.importorskip("jax.numpy") + + import pepsy as py + from pepsy.vmc.netket import prepare_fermionic_peps_for_netket + + Lx = Ly = 2 + sites = tuple((x, y) for x in range(Lx) for y in range(Ly)) + occupations = { + (0, 0): 1, + (0, 1): 1, + (1, 0): 1, + (1, 1): 0, + } + fermion = py.Fermion( + spinful=True, + symmetry="Z2", + to_backend=py.backend_jax( + device=jax.devices()[0], dtype=jnp.complex64 + ), + ) + peps = py.hrs_to_peps( + (Lx, Ly), + fermion=fermion, + occupations=occupations, + chi=3, + seed=31, + dtype="float32", + ) + + with pytest.warns(RuntimeWarning, match="Zero-padded unequal Z2"): + prepared = prepare_fermionic_peps_for_netket(peps) + + assert tuple(prepared.sites) == sites + assert all("Flat" in type(prepared[site].data).__name__ for site in sites) + + +def test_explicit_native_hubbard_terms_compile_to_matching_netket_terms(): + """NetKet uses the supplied native term mapping, not Fermion attributes.""" + import pepsy as py + from pepsy.vmc.netket import ( + _native_fermi_hubbard_terms_to_netket, + fermion_model_terms, + ) + + fermion = py.Fermion(spinful=True, symmetry="U1U1") + t, U, mu = 1.25, 3.5, 0.2 + terms = { + (0, 1): -t * fermion.hopping_operator(), + 0: fermion.onsite_term(0, U=U, mu=mu), + 1: fermion.onsite_term(1, U=U, mu=mu), + } + hamiltonian = fermion.hamiltonian(terms) + actual, constant = _native_fermi_hubbard_terms_to_netket( + fermion, + hamiltonian, + site_order=(0, 1), + ) + expected = fermion_model_terms( + fermion, + ((0, 1),), + t=t, + U=U, + mu=mu, + n_sites=2, + ) + + def collect(entries): + result = {} + for coefficient, operators in entries: + result[tuple(operators)] = result.get(tuple(operators), 0.0) + coefficient + return result + + assert constant == pytest.approx(0.0) + assert collect(actual) == pytest.approx(collect(expected)) + + +@pytest.mark.smoke +def test_warmup_summary_reports_stage_times_and_amplitude_rows(capsys): + """Warmup uses NetKet's JIT forward and gradient routes at chunk size.""" + from pepsy.vmc.netket import warmup_netket_vmc + + class State: + sampler = SimpleNamespace(n_chains=2, sweep_size=4) + sampler_state = None + samples = np.zeros((2, 3, 4), dtype=np.int8) + n_samples = 6 + n_discard_per_chain = 1 + chunk_size = 2 + + def __init__(self): + self.log_value_batches = [] + self.expect_and_grad_calls = 0 + + def reset(self): + pass + + def log_value(self, configs): + self.log_value_batches.append(tuple(configs.shape)) + return configs.sum(axis=-1) + + def expect_and_grad(self, hamiltonian): + self.expect_and_grad_calls += 1 + return SimpleNamespace(mean=-1.25), {"tensor": np.ones(1)} + + state = State() + elapsed = warmup_netket_vmc( + SimpleNamespace(vstate=state, hamiltonian=object()), + progress=False, + verbose=True, + ) + output = capsys.readouterr().out + + assert elapsed >= 0.0 + assert "sampler" in output + assert "JIT log amplitudes" in output + assert "energy + gradient" in output + assert "6 retained = 2 chains x 3/chain" in output + assert "representative 2-row forward chunk" in output + assert state.log_value_batches == [(2, 4)] + assert state.expect_and_grad_calls == 1 + + +@pytest.mark.smoke +def test_vmc_progress_starts_before_the_first_netket_update(monkeypatch): + """The user sees a status while the first JIT-heavy update is running.""" + from pepsy.vmc import netket as netket_module + + class Bar: + def __init__(self): + self.description = None + self.postfix = None + self.updates = 0 + self.closed = False + + def set_description_str(self, value): + self.description = value + + def set_postfix_str(self, value): + self.postfix = value + + def update(self, value): + self.updates += value + + def close(self): + self.closed = True + + bar = Bar() + monkeypatch.setattr( + netket_module, "_make_progress_bar", lambda **kwargs: bar + ) + callback = netket_module._VMCProgressCallback(3, enabled=True) + + callback.start("first update: sampling and compiling gradients") + assert bar.description == "VMC 0/3: preparing" + assert "compiling gradients" in bar.postfix + + class PendingParameter: + def __init__(self): + self.synchronized = False + + def block_until_ready(self): + self.synchronized = True + + pending_parameter = PendingParameter() + stats = SimpleNamespace(mean=-1.0, error_of_mean=0.1, variance=0.2) + driver = SimpleNamespace( + _loss_name="Energy", + _loss_stats=stats, + variational_state=SimpleNamespace(parameters={"t0": pending_parameter}), + ) + callback(1, {"Energy": stats}, driver) + + assert pending_parameter.synchronized + assert bar.description == "VMC energy" + assert "first update" in bar.postfix + assert bar.updates == 1 + callback.close() + assert bar.closed + + +@pytest.mark.smoke +def test_vmc_progress_fails_fast_after_a_nonfinite_parameter_update(): + """A diverged PEPS update must not spend further VMC steps sampling NaNs.""" + from pepsy.vmc import netket as netket_module + + stats = SimpleNamespace(mean=-1.0, error_of_mean=0.1, variance=0.2) + driver = SimpleNamespace( + _loss_name="Energy", + _loss_stats=stats, + variational_state=SimpleNamespace(parameters={"t0": np.array([np.nan])}), + ) + callback = netket_module._VMCProgressCallback(2, enabled=False) + + with pytest.raises(FloatingPointError, match="non-finite PEPS parameters"): + callback(1, {"Energy": stats}, driver) + + +@pytest.mark.smoke +def test_sample_resource_monitor_reports_and_retains_metrics(monkeypatch, capsys): + """Sampling forwards the opt-in GPU/RSS report into its diagnostics.""" + from pepsy.vmc import netket as netket_module + + report = netket_module.NetKetResourceUsage( + elapsed_seconds=0.2, + host_rss_before_mib=128.0, + host_rss_after_mib=144.0, + host_rss_peak_mib=160.0, + gpu_after=( + netket_module.NetKetGPUUsage( + index=0, + name="test GPU", + memory_used_mib=1024, + memory_total_mib=2048, + utilization_percent=25, + memory_utilization_percent=10, + process_memory_mib=512, + ), + ), + gpu_peak=( + netket_module.NetKetGPUUsage( + index=0, + name="test GPU", + memory_used_mib=1536, + memory_total_mib=2048, + utilization_percent=90, + memory_utilization_percent=20, + process_memory_mib=1024, + ), + ), + ) + + class Monitor: + def __init__(self, *, interval): + self.interval = interval + self.started = False + + def start(self): + self.started = True + return self + + def stop(self): + assert self.started + return report + + class State: + sampler_state = None + samples = np.zeros((2, 3, 4), dtype=np.int8) + n_samples = 6 + n_discard_per_chain = 1 + chunk_size = 2 + + sampler = SimpleNamespace(n_chains=2, sweep_size=4) + setup = SimpleNamespace(vstate=State(), sampler=sampler) + monkeypatch.setattr(netket_module, "_NetKetResourceMonitor", Monitor) + + sampled = netket_module.NetKetPEPSVMC.sample( + setup, + resource_monitor=True, + resource_interval=0.5, + ) + output = capsys.readouterr().out + + assert sampled.diagnostics["resources"]["host_rss_peak_mib"] == 160.0 + assert sampled.diagnostics["resources"]["gpu_peak"][0]["utilization_percent"] == 90 + assert "NetKet sampling resources" in output + + +@pytest.mark.smoke +def test_sample_fresh_resets_the_retained_netket_cache(): + """``fresh=True`` guarantees a new batch without resetting chain state.""" + from pepsy.vmc.netket import NetKetPEPSVMC + + class State: + sampler_state = None + n_samples = 2 + n_discard_per_chain = 0 + chunk_size = 2 + + def __init__(self): + self._samples = np.zeros((1, 2, 3), dtype=np.int8) + self.resets = 0 + + def reset(self): + self.resets += 1 + self._samples = None + + @property + def samples(self): + if self._samples is None: + self._samples = np.full( + (1, 2, 3), self.resets, dtype=np.int8 + ) + return self._samples + + state = State() + setup = SimpleNamespace( + vstate=state, + sampler=SimpleNamespace(n_chains=1, sweep_size=1), + ) + sampled = NetKetPEPSVMC.sample(setup, fresh=True) + + assert state.resets == 1 + assert sampled.native is state._samples + assert np.all(sampled.native == 1) + + +@pytest.mark.smoke +def test_warmup_stops_the_resource_monitor_when_jit_fails(monkeypatch): + """Telemetry must not leave a polling thread running after a JIT error.""" + from pepsy.vmc import netket as netket_module + + monitor = SimpleNamespace(started=False, stopped=False) + monitor.start = lambda: setattr(monitor, "started", True) or monitor + monitor.stop = lambda: setattr(monitor, "stopped", True) or None + monkeypatch.setattr( + netket_module, + "_NetKetResourceMonitor", + lambda **kwargs: monitor, + ) + + class State: + samples = np.zeros((1, 2, 3), dtype=np.int8) + chunk_size = 2 + + def reset(self): + pass + + def log_value(self, configs): + raise RuntimeError("synthetic JIT failure") + + with pytest.raises(RuntimeError, match="synthetic JIT failure"): + netket_module.warmup_netket_vmc( + SimpleNamespace(vstate=State(), hamiltonian=object()), + progress=False, + resource_monitor=True, + ) + assert monitor.started + assert monitor.stopped + + +@pytest.mark.smoke +def test_netket_mc_diagnostic_facade_forwards_and_version_gates(): + """Recent NetKet convergence helpers stay optional-dependency friendly.""" + from pepsy.vmc.netket import NetKetPEPSVMC + + class State: + def check_mc_convergence(self, hamiltonian, **kwargs): + return "check", hamiltonian, kwargs + + def thermalise(self, hamiltonian, **kwargs): + return "thermalise", hamiltonian, kwargs + + def expect_to_precision(self, observable, **kwargs): + return "precision", observable, kwargs + + setup = NetKetPEPSVMC( + hilbert=None, + graph=None, + hamiltonian="energy", + sampler=None, + vstate=State(), + model=None, + ansatz=SimpleNamespace(n_sites=1, n_params=1), + config_map=None, + preconditioner=None, + ) + assert setup.check_mc_convergence(min_chain_length=100) == ( + "check", + "energy", + {"min_chain_length": 100}, + ) + assert setup.thermalise(max_chain_length=200) == ( + "thermalise", + "energy", + {"max_chain_length": 200}, + ) + assert setup.expect_to_precision(rtol=0.01) == ( + "precision", + "energy", + {"rtol": 0.01}, + ) + + missing = NetKetPEPSVMC( + hilbert=None, + graph=None, + hamiltonian="energy", + sampler=None, + vstate=object(), + model=None, + ansatz=SimpleNamespace(n_sites=1, n_params=1), + config_map=None, + preconditioner=None, + ) + with pytest.raises(RuntimeError, match="NetKet >= 3.22"): + missing.check_mc_convergence() + + +@pytest.mark.smoke +def test_final_fermion_observables_use_conserving_netket_operators(): + """Fixed-sector final measurements avoid generic FermionOperator2nd.""" + nk = pytest.importorskip("netket") + from pepsy.vmc.netket import ( + NetKetEtaPairObservable, + _netket_eta_pair_operator, + standard_fermion_observables, + ) + + hilbert = nk.hilbert.SpinOrbitalFermions( + 4, + s=1 / 2, + n_fermions_per_spin=(2, 2), + ) + expected_name = "ParticleNumberAndSpinConservingFermioperator2nd" + observables = standard_fermion_observables(hilbert) + assert all(type(operator).__name__ == expected_name for operator in observables.values()) + + ansatz = SimpleNamespace( + orbital_sites=((0, 0), (0, 1), (1, 0), (1, 1)) + ) + eta_pair = _netket_eta_pair_operator( + hilbert, + ansatz, + NetKetEtaPairObservable(1, 0), + ) + assert type(eta_pair).__name__ == expected_name diff --git a/tests/test_optimize_mera.py b/tests/test_optimize_mera.py index 8be752c..ef9ba40 100644 --- a/tests/test_optimize_mera.py +++ b/tests/test_optimize_mera.py @@ -666,19 +666,13 @@ def test_unified_fermion_helper_adapts_to_qmera_mode_terms(): """One Fermion model should feed both site and qMERA energy layouts.""" pytest.importorskip("symmray") geometry = QMeraGeometry(shape=3, site_modes=("up", "down")) - fermion = Fermion( - spinful=True, - symmetry="U1U1", - t=0.5, - U=4.0, - mu=0.1, - ) + fermion = Fermion(spinful=True, symmetry="U1U1") - site_terms = fermion.local_terms(((0, 1), (1, 2))) - mode_terms = fermion.local_terms(geometry, layout="qmera") + site_terms = fermion.local_terms(((0, 1), (1, 2)), t=0.5, U=4.0, mu=0.1) + mode_terms = fermion.local_terms(geometry, layout="qmera", t=0.5, U=4.0, mu=0.1) direct_terms = qmera_symmray_fermi_hubbard_terms( geometry, - fermion=fermion, + fermion=fermion, t=0.5, U=4.0, mu=0.1, ) assert set(site_terms) == {(0, 1), (1, 2)} @@ -700,13 +694,7 @@ def test_unified_fermion_qmera_optimizer_runs_torch_autodiff(): array_backend = backend_torch(dtype=torch.complex128) backend = QMeraSymmrayFermionBackend(to_backend=array_backend) registry = symmray_fermion_gate_registry(backend=backend) - fermion = Fermion( - spinful=True, - symmetry="U1U1", - t=0.2, - U=0.5, - mu=0.1, - ) + fermion = Fermion(spinful=True, symmetry="U1U1") def product_state_factory(schedule, sites, **kwargs): return backend.product_state( @@ -732,9 +720,10 @@ def product_state_factory(schedule, sites, **kwargs): param_scale=0.01, product_state_factory=product_state_factory, ) - terms = builder.fermion_terms() + terms = builder.fermion_terms(t=0.2, U=0.5, mu=0.1) optimizer = builder.fermion_parametric_optimizer( energy_per_site=False, + term_params={"t": 0.2, "U": 0.5, "mu": 0.1}, ) initial = optimizer.loss(energy_per_site=False) result = optimizer.run( @@ -1052,16 +1041,10 @@ def product_state_factory(schedule, sites, **kwargs): ) hubbard_backend = QMeraSymmrayFermionBackend() hubbard_registry = symmray_fermion_gate_registry(backend=hubbard_backend) - hubbard = Fermion( - spinful=True, - symmetry="U1U1", - t=0.2, - U=0.5, - mu=0.1, - ) + hubbard = Fermion(spinful=True, symmetry="U1U1") hubbard_terms = qmera_symmray_fermi_hubbard_terms( hubbard_geometry, - fermion=hubbard, + fermion=hubbard, t=0.2, U=0.5, mu=0.1, ) hubbard_lightcone, hubbard_direct = run_case( hubbard_geometry, diff --git a/tests/test_optimize_mps.py b/tests/test_optimize_mps.py index d198f68..256a00b 100644 --- a/tests/test_optimize_mps.py +++ b/tests/test_optimize_mps.py @@ -153,8 +153,6 @@ def test_mps_optimizer_simple_update_routes_torch_u1u1_long_range_gate(): fermion = py.Fermion( spinful=True, symmetry="U1U1", - t=1.0, - U=8.0, dtype="float64", ) state = py.hrs_to_mps( @@ -168,7 +166,7 @@ def test_mps_optimizer_simple_update_routes_torch_u1u1_long_range_gate(): ) state.apply_to_arrays(backend) fermion.to_backend = backend - hopping = fermion.hopping_gate(0.001, imaginary=True) + hopping = fermion.hopping_gate(0.001, t=1.0, imaginary=True) optimizer = py.MpsOptimizer( state, diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index f0e844e..43c80d4 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -3267,9 +3267,7 @@ def test_tree_native_fermionic_gate_stream_matches_mps(): t, U, dt = 1.0, 8.0, 0.05 state_dtype = "complex128" - fermion = pepsy.Fermion( - spinful=True, symmetry="U1U1", t=t, U=U, mu=0.0, dtype=state_dtype - ) + fermion = pepsy.Fermion(spinful=True, symmetry="U1U1", dtype=state_dtype) setup = fermion.lattice_half_filling(Lx, Ly, pattern="checkerboard", cyclic=True) mapper = tensors.OneDMap(Lx, Ly, mode="snake") _, coo2idx = mapper.build() @@ -3356,9 +3354,6 @@ def test_native_fermionic_submpo_keeps_graded_hub_recovery(monkeypatch): fermion = pepsy.Fermion( spinful=True, symmetry="U1U1", - t=1.0, - U=8.0, - mu=0.0, dtype="complex128", ) occupations = ((1, 0), (0, 1), (1, 0), (0, 1)) diff --git a/tests/test_package_layout.py b/tests/test_package_layout.py index dfb2774..5e9cfbc 100644 --- a/tests/test_package_layout.py +++ b/tests/test_package_layout.py @@ -65,6 +65,7 @@ ContractionConfig, FermionSiteEncoding, MCState, + NetKetEtaPairObservable, NetKetVMCSetup, OptimizationConfig, SamplingConfig, @@ -165,6 +166,7 @@ def test_new_namespace_imports_resolve(): assert TorchPEPSBoundaryAmplitude is not None assert TorchVMCDriver is not None assert TorchVMCSetup is not None + assert NetKetEtaPairObservable is not None assert NetKetVMCSetup is not None assert TorchVMCStepResult is not None assert TorchSquareLattice is not None diff --git a/tests/test_symmetric_tensors.py b/tests/test_symmetric_tensors.py index e36d2a5..ab902fe 100644 --- a/tests/test_symmetric_tensors.py +++ b/tests/test_symmetric_tensors.py @@ -63,9 +63,15 @@ def test_spinful_fermion_helper_bundles_symmetry_aware_building_blocks( symmetry, expected_model, expected_charge ): """The convenience helper should cover the native U1 and U1U1 workflows.""" - fermions = SpinfulFermion(symmetry=symmetry, t=1.0, U=3.0) + fermions = SpinfulFermion(symmetry=symmetry) assert SpinfulFermionHubbard is SpinfulFermion + assert not hasattr(fermions, "t") + assert not hasattr(fermions, "U") + assert not hasattr(fermions, "V") + assert not hasattr(fermions, "mu") + with pytest.raises(TypeError, match="unexpected keyword argument 't'"): + SpinfulFermion(symmetry=symmetry, t=1.0) assert isinstance(SymmFermions.spinful(symmetry=symmetry), SpinfulFermion) assert pepsy.SpinfulFermion is SpinfulFermion assert pepsy.SymmFermions is SymmFermions @@ -89,13 +95,17 @@ def test_spinful_fermion_helper_bundles_symmetry_aware_building_blocks( ((0, 1), (2, 3)), ((1, 2), (3, 0)), ) - stream = fermions.strang_gate_stream(edges, 0.02, sites=range(4)) + stream = fermions.strang_gate_stream( + edges, 0.02, sites=range(4), t=1.0, U=3.0 + ) assert isinstance(stream, SymGateStream) assert stream.order == 2 assert len(stream) == 16 assert all(len(where) == 2 for _, where in stream[4:12]) - ham = fermions.hamiltonian(((0, 1), (1, 2)), mu=0.1) + ham = fermions.hamiltonian( + ((0, 1), (1, 2)), t=1.0, U=3.0, mu=0.1 + ) assert isinstance(ham, SymHamiltonian) assert ham.model == expected_model assert ham.symmetry == symmetry @@ -104,13 +114,7 @@ def test_spinful_fermion_helper_bundles_symmetry_aware_building_blocks( @pytest.mark.parametrize("symmetry", ["U1", "Z2"]) def test_unified_fermion_helper_supports_spinless_native_workflow(symmetry): """The unified helper should expose the spinless native t-V workflow.""" - fermions = Fermion( - spinful=False, - symmetry=symmetry, - t=1.0, - V=2.0, - mu=0.3, - ) + fermions = Fermion(spinful=False, symmetry=symmetry) assert fermions.model == "fermi_hubbard_spinless" assert fermions.physical_sectors == default_physical_sectors(symmetry, 2) @@ -121,7 +125,7 @@ def test_unified_fermion_helper_supports_spinless_native_workflow(symmetry): assert fermions.operator_charge("create") == (1 if symmetry == "Z2" else 1) assert type(fermions.operator("number")).__name__.endswith("FermionicArray") - density = fermions.gate("density", 0.1) + density = fermions.gate("density", 0.1, V=2.0) assert type(density).__name__.endswith("FermionicArray") assert density.to_dense().shape == (2, 2, 2, 2) @@ -130,25 +134,25 @@ def test_unified_fermion_helper_supports_spinless_native_workflow(symmetry): 0.01, sites=range(3), order=2, + t=1.0, + V=2.0, + mu=0.3, ) assert isinstance(stream, SymGateStream) assert stream.order == 2 assert all(len(where) == 2 for _, where in stream if isinstance(where, tuple)) - ham = fermions.hamiltonian(((0, 1), (1, 2))) + ham = fermions.hamiltonian(((0, 1), (1, 2)), t=1.0, V=2.0, mu=0.3) assert ham.model == "fermi_hubbard_spinless" assert ham.symmetry == symmetry - assert set(fermions.local_terms(((0, 1), (1, 2)))) == {(0, 1), (1, 2)} + assert set(fermions.local_terms( + ((0, 1), (1, 2)), t=1.0, V=2.0, mu=0.3 + )) == {(0, 1), (1, 2)} def test_unified_spinful_fermion_gate_stream_runs_native_mps(): """The unified spinful model should evolve a charge-conserving MPS.""" - fermion = Fermion( - spinful=True, - symmetry="U1U1", - t=0.5, - U=2.0, - ) + fermion = Fermion(spinful=True, symmetry="U1U1") state = SymMPS.random( 3, symmetry="U1U1", @@ -161,7 +165,9 @@ def test_unified_spinful_fermion_gate_stream_runs_native_mps(): ) state.apply_gates( - fermion.gate_stream(((0, 1), (1, 2)), 0.01, sites=range(3)), + fermion.gate_stream( + ((0, 1), (1, 2)), 0.01, sites=range(3), t=0.5, U=2.0 + ), method="direct", contract="split", max_bond=4, @@ -175,7 +181,7 @@ def test_unified_spinful_fermion_gate_stream_runs_native_mps(): def test_ps_to_mps_accepts_fermion_and_builds_native_charge_sector(): """The public MPS constructor should hide SymMPS for Fermion starts.""" - fermion = Fermion(spinful=True, symmetry="U1U1", t=0.5, U=2.0) + fermion = Fermion(spinful=True, symmetry="U1U1") mps = pepsy.ps_to_mps(4, fermion=fermion, seed=17) assert mps.L == 4 @@ -240,7 +246,7 @@ def test_hrs_to_mps_direct_uses_symmray_random_blocks(): def test_ps_to_peps_accepts_fermion_coordinate_occupations(): """The public PEPS constructor should build a native fixed-charge seed.""" - fermion = Fermion(spinful=True, symmetry="U1U1", U=2.0) + fermion = Fermion(spinful=True, symmetry="U1U1") occupations = { (x, y): (1, 0) if (x + y) % 2 == 0 else (0, 1) for x in range(2) @@ -332,6 +338,7 @@ def test_hrs_to_peps_accepts_fermion_sector_and_chi(): subsizes="maximal", seed=32, dtype="complex128", + normalize=True, ) assert (peps.Lx, peps.Ly) == (2, 2) @@ -344,6 +351,31 @@ def test_hrs_to_peps_accepts_fermion_sector_and_chi(): assert pepsy.hrps_to_peps is pepsy.hrs_to_peps +def test_hrs_to_peps_skips_global_norm_by_default_for_vmc(monkeypatch): + """The VMC-safe default avoids the unrelated CPU boundary-norm contraction.""" + fermion = Fermion(spinful=True, symmetry="Z2") + occupations = { + (x, y): 1 if (x + y) % 2 == 0 else 0 + for x in range(2) + for y in range(2) + } + + def unexpected_normalize(_): + raise AssertionError("global PEPS normalization must be skipped") + + monkeypatch.setattr(SymPEPS, "normalize", unexpected_normalize) + peps = pepsy.hrs_to_peps( + (2, 2), + fermion=fermion, + occupations=occupations, + chi=3, + seed=33, + ) + + assert (peps.Lx, peps.Ly) == (2, 2) + assert 1 < peps.max_bond() <= 3 + + @pytest.mark.parametrize( ("symmetry", "expected_charge", "expected_occupation"), [ @@ -356,7 +388,7 @@ def test_lattice_half_filling_prepares_explicit_peps_metadata( symmetry, expected_charge, expected_occupation ): """Lattice setup normalizes occupations without building terms or gates.""" - fermion = Fermion(spinful=True, symmetry=symmetry, U=2.0) + fermion = Fermion(spinful=True, symmetry=symmetry) setup = fermion.lattice_half_filling(3, 2, pattern="checkerboard") @@ -377,9 +409,9 @@ def test_lattice_half_filling_prepares_explicit_peps_metadata( def test_symdmrg_fermionic_state_accepts_raw_fermion_constructor_output(): """SymDMRG should restore native tensors without a SymMPS wrapper input.""" - fermion = Fermion(spinful=True, symmetry="U1U1", t=0.5, U=2.0) + fermion = Fermion(spinful=True, symmetry="U1U1") mps = pepsy.ps_to_mps(3, fermion=fermion, seed=23) - mpo = fermion.hamiltonian(((0, 1), (1, 2))).to_mpo( + mpo = fermion.hamiltonian(((0, 1), (1, 2)), t=0.5, U=2.0).to_mpo( L=3, max_bond=8, compress=True, @@ -645,7 +677,7 @@ def test_native_fermionic_compiled_plan_preserves_phases_and_dummy_modes(): @pytest.mark.parametrize("symmetry", ["Z2", "Z2Z2"]) def test_unified_spinful_fermion_supports_symmray_parity_symmetries(symmetry): """Spinful parity symmetries should use native charges and gates.""" - fermion = Fermion(spinful=True, symmetry=symmetry, t=0.5, U=2.0, mu=0.1) + fermion = Fermion(spinful=True, symmetry=symmetry) assert fermion.model == ( "fermi_hubbard" if symmetry == "Z2" else "fermi_hubbard_u1u1" @@ -662,14 +694,16 @@ def test_unified_spinful_fermion_supports_symmray_parity_symmetries(symmetry): 0 if symmetry == "Z2" else (1, 1) ) assert type(fermion.observable("number")).__name__.endswith("FermionicArray") - assert type(fermion.hopping_gate(0.01)).__name__.endswith("FermionicArray") - assert type(fermion.interaction_gate(0.01)).__name__.endswith("FermionicArray") + assert type(fermion.hopping_gate(0.01, t=0.5)).__name__.endswith("FermionicArray") + assert type(fermion.interaction_gate(0.01, U=2.0)).__name__.endswith("FermionicArray") -def test_unified_spinful_fermion_hamiltonian_keeps_mu_parameter(): - """Configured spinful chemical potentials must reach Symmray terms.""" - fermion = Fermion(spinful=True, symmetry="U1U1", U=3.0, mu=(0.2, 0.4)) - hamiltonian = fermion.hamiltonian(((0, 1),)) +def test_unified_spinful_fermion_hamiltonian_keeps_explicit_mu_parameter(): + """Explicit spinful chemical potentials must reach Symmray terms.""" + fermion = Fermion(spinful=True, symmetry="U1U1") + hamiltonian = fermion.hamiltonian( + ((0, 1),), t=1.0, U=3.0, mu=(0.2, 0.4) + ) assert hamiltonian.parameters["mu"] == (0.2, 0.4) dense = hamiltonian.terms[(0, 1)].to_dense() @@ -684,14 +718,7 @@ def test_unified_spinful_fermion_hamiltonian_keeps_mu_parameter(): def test_fermion_exposes_explicit_native_operator_terms(): """Named and generic APIs return the unexponentiated fermion terms.""" - fermion = Fermion( - spinful=True, - symmetry="U1U1", - t=1.7, - U=3.0, - mu=0.4, - V=0.6, - ) + fermion = Fermion(spinful=True, symmetry="U1U1") hopping = fermion.hopping_operator() reference_hopping = fermion.hamiltonian( @@ -720,7 +747,7 @@ def test_fermion_exposes_explicit_native_operator_terms(): (0.0, 1.0, 1.0, 2.0), ) np.testing.assert_allclose( - np.diag(fermion.onsite_term(0).to_dense()), + np.diag(fermion.onsite_term(0, U=3.0, mu=0.4).to_dense()), (0.0, -0.4, -0.4, 3.0 - 0.8), ) @@ -883,7 +910,7 @@ def test_fermion_spin_gates_preserve_torch_backend(): def test_mps_energy_uses_explicit_fermion_terms_natively(): """An explicit one- plus two-site SymHamiltonian stays on native MPS terms.""" - fermion = Fermion(spinful=True, symmetry="U1U1", U=2.0, mu=0.0) + fermion = Fermion(spinful=True, symmetry="U1U1") ham = fermion.hamiltonian( { (0, 1): -fermion.hopping_operator(), @@ -913,13 +940,7 @@ def test_mps_energy_uses_explicit_fermion_terms_natively(): def test_fermion_explicit_coordinate_terms_preserve_peps_locations(): """Coordinate-site terms remain usable by PEPS and mapped MPO workflows.""" - fermion = Fermion( - spinful=True, - symmetry="U1U1", - t=1.0, - U=3.0, - mu=0.2, - ) + fermion = Fermion(spinful=True, symmetry="U1U1") edges = ( ((0, 0), (0, 1)), ((0, 0), (1, 0)), @@ -927,7 +948,7 @@ def test_fermion_explicit_coordinate_terms_preserve_peps_locations(): ) sites = tuple(sorted({site for edge in edges for site in edge})) terms = {edge: -fermion.hopping_operator() for edge in edges} - terms |= {site: fermion.onsite_term(site) for site in sites} + terms |= {site: fermion.onsite_term(site, U=3.0, mu=0.2) for site in sites} hamiltonian = fermion.hamiltonian(terms) @@ -960,8 +981,8 @@ def test_unified_fermion_peps_energy_accepts_boundary_chi(): seed=17, dtype="complex128", ) - fermion = Fermion(symmetry="U1U1", U=2.0, mu=(0.1, 0.2)) - ham = fermion.hamiltonian(peps.edges) + fermion = Fermion(symmetry="U1U1") + ham = fermion.hamiltonian(peps.edges, t=1.0, U=2.0, mu=(0.1, 0.2)) exact = peps.energy(ham) boundary = peps.energy(ham, chi=8) @@ -979,10 +1000,10 @@ def test_unified_fermion_gates_and_hamiltonian_preserve_backend(): ) values = ( - fermion.onsite_gate(0.01), - fermion.hopping_gate(0.01), + fermion.onsite_gate(0.01, U=8.0), + fermion.hopping_gate(0.01, t=1.0), fermion.density_gate(0.01, V=0.2), - fermion.hamiltonian(((0, 1),)).terms[(0, 1)], + fermion.hamiltonian(((0, 1),), t=1.0, U=8.0).terms[(0, 1)], ) assert all(value.backend == "torch" for value in values) @@ -990,14 +1011,16 @@ def test_unified_fermion_gates_and_hamiltonian_preserve_backend(): def test_spinful_interaction_gate_has_exact_doublon_phase(): """The onsite gate should only phase the doubly occupied basis state.""" theta = 0.2 * 3.0 - gate = Fermion(spinful=True, symmetry="U1U1", U=3.0).gate( + gate = Fermion(spinful=True, symmetry="U1U1").gate( "interaction", 0.2, + U=3.0, ) expected = np.diag([1.0, 1.0, 1.0, np.exp(-1j * theta)]) np.testing.assert_allclose(gate.to_dense(), expected) - imaginary = Fermion(spinful=True, symmetry="U1U1", U=3.0).interaction_gate( + imaginary = Fermion(spinful=True, symmetry="U1U1").interaction_gate( 0.2, + U=3.0, imaginary=True, ) np.testing.assert_allclose( @@ -1008,7 +1031,7 @@ def test_spinful_interaction_gate_has_exact_doublon_phase(): def test_fermion_generic_exponential_matches_named_interaction_gate(): """Generic neutral monomials should share the native exponentiator.""" - fermion = Fermion(spinful=True, symmetry="U1U1", U=3.0) + fermion = Fermion(spinful=True, symmetry="U1U1") generic = fermion.exponential( [(3.0, ((0, "number_up"), (0, "number_down")))], 0.2, @@ -1017,16 +1040,16 @@ def test_fermion_generic_exponential_matches_named_interaction_gate(): np.testing.assert_allclose( generic.to_dense(), - fermion.interaction_gate(0.2).to_dense(), + fermion.interaction_gate(0.2, U=3.0).to_dense(), ) assert type(generic).__name__ == "U1U1FermionicArray" def test_fermion_hopping_gate_matches_native_hamiltonian_exponential(): """Native hopping imaginary time must lower the matching term energy.""" - fermion = Fermion(spinful=True, symmetry="U1U1", t=1.0, U=0.0) + fermion = Fermion(spinful=True, symmetry="U1U1") ham = fermion.hamiltonian({(0, 1): -fermion.hopping_operator()}) - named = fermion.hopping_gate(0.01, imaginary=True) + named = fermion.hopping_gate(0.01, t=1.0, imaginary=True) reference = ham.trotter_gates(0.01, imaginary=True)[0][0] np.testing.assert_allclose(named.to_dense(), reference.to_dense()) @@ -1097,7 +1120,6 @@ def test_native_hopping_gate_has_correct_long_range_parity_sign( fermion = Fermion( spinful=spinful, symmetry="U1U1" if spinful else "U1", - t=1.3, ) dt = 0.13 state = pepsy.ps_to_mps( @@ -3076,8 +3098,6 @@ def test_symmray_mpo_real_time_fermion_stream_preserves_norm_without_truncation( fermion = Fermion( spinful=True, symmetry="U1U1", - t=1.0, - U=8.0, to_backend=backend, ) lx, ly = 2, 3 diff --git a/tests/test_tree_sampler.py b/tests/test_tree_sampler.py index 72a1ab7..44075b6 100644 --- a/tests/test_tree_sampler.py +++ b/tests/test_tree_sampler.py @@ -284,9 +284,7 @@ def _fermionic_tree(*, chi=64, steps=4, root_qubit=None): t, U, dt = 1.0, 4.0, 0.05 dtype = "complex128" - fermion = pepsy.Fermion( - spinful=True, symmetry="U1U1", t=t, U=U, mu=0.0, dtype=dtype - ) + fermion = pepsy.Fermion(spinful=True, symmetry="U1U1", dtype=dtype) setup = fermion.lattice_half_filling(Lx, Ly, pattern="checkerboard", cyclic=True) mapper = tensors.OneDMap(Lx, Ly, mode="snake") _, coo2idx = mapper.build() diff --git a/tests/test_vmc_api.py b/tests/test_vmc_api.py index e62b742..513c733 100644 --- a/tests/test_vmc_api.py +++ b/tests/test_vmc_api.py @@ -1,5 +1,6 @@ """Tests for the backend-neutral VMC API contracts.""" +import os import numpy as np import pytest from types import SimpleNamespace @@ -89,6 +90,15 @@ def test_mcstate_uses_netket_total_sample_convention_and_bridges_to_problem(): assert state.ansatz is state.peps assert state.to_problem(OperatorSum()).site_order == (0, 1) + canonical_state = MCState( + object(), + n_samples=12, + n_chains=3, + n_discard_per_chain=4, + sweep_size=2, + ) + assert canonical_state.sweep_size == 2 + with pytest.raises(ValueError, match="divisible"): MCState(object(), n_samples=5, n_chains=2) with pytest.raises(ValueError, match="either sampling"): @@ -210,16 +220,172 @@ def test_warning_types_are_backend_neutral(): assert issubclass(VMCBackendCapabilityError, NotImplementedError) -def test_netket_portable_adapter_rejects_external_weighted_batches(): +def test_netket_measure_samples_reuses_only_the_current_cache(): + from pepsy.vmc.netket import NetKetPEPSVMC + + native_samples = object() + seen = [] + + class State: + _samples = native_samples + + def expect(self, observable): + seen.append(observable) + return f"stats:{observable}" + + setup = NetKetPEPSVMC( + hilbert=None, + graph=None, + hamiltonian=None, + sampler=None, + vstate=State(), + model=None, + ansatz=SimpleNamespace(n_sites=1, n_params=1), + config_map=None, + preconditioner=None, + ) + retained = VMCSamples( + configs=np.zeros((1, 1, 1), dtype=np.int64), + n_samples_per_chain=1, + n_chains=1, + native=native_samples, + ) + + assert setup.measure_samples(retained, {"density": "n", "spin": "sz"}) == { + "density": "stats:n", + "spin": "stats:sz", + } + assert seen == ["n", "sz"] + with pytest.raises(VMCBackendCapabilityError, match="current MCState sample cache"): + setup.measure_samples( + VMCSamples( + configs=np.zeros((1, 1, 1), dtype=np.int64), + n_samples_per_chain=1, + n_chains=1, + native=object(), + ), + {"density": "n"}, + ) + + +def test_netket_measure_samples_compiles_eta_pair_on_retained_batch(monkeypatch): + import pepsy.vmc.netket as netket_vmc + from pepsy.vmc.netket import NetKetEtaPairObservable, NetKetPEPSVMC + + native_samples = object() + compiled = [] + + class State: + _samples = native_samples + + def expect(self, observable): + return f"stats:{len(observable)}" + + def fake_fermion_operator(hilbert, terms, *, constant=0.0, conserving=False): + assert hilbert == "fermion-hilbert" + assert constant == 0.0 + assert conserving == "auto" + compiled.append(tuple(terms)) + return compiled[-1] + + monkeypatch.setattr(netket_vmc, "netket_fermion_operator", fake_fermion_operator) + setup = NetKetPEPSVMC( + hilbert="fermion-hilbert", + graph=None, + hamiltonian=None, + sampler=None, + vstate=State(), + model=None, + ansatz=SimpleNamespace( + n_sites=4, + n_params=1, + orbital_sites=((0, 0), (0, 1), (1, 0), (1, 1)), + ), + config_map=None, + preconditioner=None, + ) + retained = VMCSamples( + configs=np.zeros((1, 4, 1), dtype=np.int64), + n_samples_per_chain=1, + n_chains=1, + native=native_samples, + ) + + measured = setup.measure_samples( + retained, + { + "eta": NetKetEtaPairObservable( + 1, + 0, + periodic=True, + staggered=True, + ) + }, + ) + + assert measured == {"eta": "stats:8"} + assert len(compiled) == 1 + assert compiled[0][0] == ( + -0.25, + ((0, 1, True), (0, -1, True), (2, -1, False), (2, 1, False)), + ) + assert compiled[0][1] == ( + -0.25, + ((2, 1, True), (2, -1, True), (0, -1, False), (0, 1, False)), + ) + + +def test_netket_portable_adapter_rejects_weighted_batches(): from pepsy.vmc.netket import NetKetVMCSetup setup = NetKetVMCSetup(setup=object(), problem=object()) - with pytest.raises(VMCBackendCapabilityError, match="externally supplied"): - setup.measure(samples=np.zeros((2, 1), dtype=np.int64)) - with pytest.raises(VMCBackendCapabilityError, match="externally supplied"): + with pytest.raises(VMCBackendCapabilityError, match="weighted or proposal"): + setup.measure(weights=np.ones(2)) + with pytest.raises(VMCBackendCapabilityError, match="weighted sample"): setup.optimize(n_steps=1, weights=np.ones(2)) +def test_netket_portable_adapter_measures_a_retained_batch(): + from pepsy.vmc.netket import NetKetVMCSetup + + retained_native = object() + seen = {} + + class NativeSetup: + hamiltonian = "hamiltonian" + observables = {"density": "density"} + + def measure_samples(self, samples, observables=None): + seen["samples"] = samples + seen["observables"] = observables + return { + name: SimpleNamespace( + mean=float(index), + variance=0.0, + error_of_mean=0.0, + ) + for index, name in enumerate(observables) + } + + facade = NetKetVMCSetup( + setup=NativeSetup(), + problem=VMCProblem(peps=object(), hamiltonian="hamiltonian"), + ) + retained = VMCSamples( + configs=np.zeros((1, 1, 1), dtype=np.int64), + n_samples_per_chain=1, + n_chains=1, + native=retained_native, + ) + measurement = facade.measure(samples=retained) + + assert seen["samples"] is retained + assert seen["observables"] == {"energy": "hamiltonian", "density": "density"} + assert measurement.energy_mean == 0.0 + assert measurement.observables["density"].mean == 1.0 + assert measurement.diagnostics["samples"] is retained + + def test_shared_configuration_objects_normalize_aliases_and_defaults(): contraction = ContractionConfig(method="boundary_mps", chi=4, cutoff=1e-8) sampling = SamplingConfig( @@ -233,11 +399,29 @@ def test_shared_configuration_objects_normalize_aliases_and_defaults(): assert contraction.method == "boundary" assert contraction.chi == 4 assert sampling.thin == 2 + assert sampling.sweep_size == 2 + assert sampling.n_discard_per_chain == 3 assert sampling.n_samples == 16 assert sampling.torch_kwargs()["n_samples"] == 16 assert sampling.netket_kwargs()["n_samples"] == 16 assert optimization.method == "minsr" + canonical = SamplingConfig( + n_samples=16, + n_chains=2, + n_discard_per_chain=3, + sweep_size=2, + ) + assert canonical.n_samples_per_chain == 8 + assert canonical.burn_in == 3 + assert canonical.thin == 2 + assert canonical.torch_kwargs()["sweep_size"] == 2 + + with pytest.raises(ValueError, match="either n_samples"): + SamplingConfig(n_samples=16, n_samples_per_chain=8, n_chains=2) + with pytest.raises(ValueError, match="either sweep_size"): + SamplingConfig(sweep_size=2, thin=2) + with pytest.raises(ValueError, match="chi is required"): ContractionConfig(method="ctmrg") with pytest.raises(ValueError, match="n_samples_per_chain"): @@ -267,10 +451,10 @@ def forward(self, configs): ) samples = driver.sample( sampling=SamplingConfig( - n_samples_per_chain=2, + n_samples=4, n_chains=2, - burn_in=0, - thin=1, + n_discard_per_chain=0, + sweep_size=1, seed=12, ) ) @@ -539,6 +723,14 @@ def set_postfix(self, postfix): "accept": "0.750", } + _set_vmc_progress_postfix( + bar, + SimpleNamespace(acceptance_rate=0.625, proposal_stats=None, sr=None), + model=model, + progress_postfix="acceptance", + ) + assert bar.postfix == {"accept": "0.625"} + _set_evaluation_progress_postfix( bar, model=model, @@ -871,6 +1063,51 @@ def contract_boundary(self, **kwargs): assert network.calls[1]["canonize"] is True +def test_netket_ctmrg_zero_separation_has_flat_symmray_retry(): + import pepsy.vmc.netket as netket_vmc + + flat_array = type( + "Z2FermionicArrayFlat", + (), + {"__module__": "symmray.fake"}, + )() + + class Tensor: + data = flat_array + + class Network: + def __init__(self): + self.calls = [] + + def __iter__(self): + return iter((Tensor(),)) + + def contract_ctmrg(self, **kwargs): + self.calls.append(kwargs) + if len(self.calls) == 1: + raise TypeError("dot_general shape mismatch") + return (1.0, 0.0) + + netket_vmc._FLAT_SYMMRAY_CTMRG_FALLBACK_WARNED = False + network = Network() + with pytest.warns(RuntimeWarning, match="CTMRG.*max_separation=0"): + result = netket_vmc._contract_ctmrg_for_vmc( + network, + max_bond=8, + cutoff=0.0, + method_opts={ + "sequence": ("ymin", "ymax", "xmin", "xmax"), + "max_separation": 0, + "canonize": True, + }, + ) + assert result == (1.0, 0.0) + assert network.calls[0]["max_separation"] == 0 + assert network.calls[1]["max_separation"] == 1 + assert network.calls[1]["sequence"] == ("ymin", "ymax", "xmin", "xmax") + assert network.calls[1]["canonize"] is True + + def test_netket_amplitude_timing_reports_batch_average(): from pepsy.vmc.netket import NetKetAmplitudeTiming @@ -888,22 +1125,32 @@ class Ansatz: n_params = 1 class Model: + def __init__(self): + self.apply_calls = 0 + def apply(self, variables, configs): - assert "params" in variables + self.apply_calls += 1 return configs.sum(axis=-1) class State: samples = np.zeros((1, 2, 2), dtype=np.int8) - parameters = {"x": 1.0} - State.model = Model() + def __init__(self): + self.log_value_batches = [] + + def log_value(self, configs): + self.log_value_batches.append(tuple(configs.shape)) + return configs.sum(axis=-1) + + model = Model() + state = State() setup = NetKetPEPSVMC( hilbert=None, graph=None, hamiltonian=None, sampler=None, - vstate=State(), - model=State.model, + vstate=state, + model=model, ansatz=Ansatz(), config_map=None, preconditioner=None, @@ -912,6 +1159,48 @@ class State: timing = setup.benchmark_amplitude(n_samples=1) assert timing.n_samples == 1 assert timing.amplitude_seconds >= 0.0 + assert state.log_value_batches == [(1, 2), (1, 2)] + assert model.apply_calls == 0 + + +def test_netket_setup_to_peps_unpacks_current_flax_parameters(): + import jax + import jax.numpy as jnp + import quimb.tensor as qtn + from pepsy.vmc.netket import NetKetPEPSVMC + + original = qtn.TensorNetwork( + [qtn.Tensor(np.array([1.0, 2.0]), inds=("physical",), tags=("I0",))] + ) + params, skeleton = qtn.pack(original) + leaves, treedef = jax.tree_util.tree_flatten(params) + ansatz = SimpleNamespace( + leaves=tuple(leaves), + treedef=treedef, + skeleton=skeleton, + n_sites=1, + n_params=2, + ) + updated = {"params": {"t0": jnp.array([3.0, 4.0])}} + state = SimpleNamespace(variables=updated) + setup = NetKetPEPSVMC( + hilbert=None, + graph=None, + hamiltonian=None, + sampler=None, + vstate=state, + model=None, + ansatz=ansatz, + config_map=None, + preconditioner=None, + ) + + restored = setup.to_peps() + np.testing.assert_allclose(restored.tensor_map[0].data, [3.0, 4.0]) + np.testing.assert_allclose( + setup.to_peps(updated["params"]).tensor_map[0].data, + [3.0, 4.0], + ) def test_netket_vmc_config_validates_shared_settings(): @@ -977,6 +1266,79 @@ def test_netket_compiler_lowers_common_fermion_terms(): assert compiled.hilbert is hilbert +def test_netket_declarative_observables_request_conserving_operators(monkeypatch): + """User-supplied symbolic observables get the fixed-sector fast path.""" + import pepsy.vmc.netket as netket_vmc + + terms_calls = [] + common_calls = [] + + def fake_fermion_operator(hilbert, terms, *, constant=0.0, conserving=False): + terms_calls.append((hilbert, terms, constant, conserving)) + return "fermion-observable" + + def fake_common_operator(hilbert, terms, *, site_order=None, conserving=False): + common_calls.append((hilbert, terms, site_order, conserving)) + return "common-observable" + + monkeypatch.setattr(netket_vmc, "netket_fermion_operator", fake_fermion_operator) + monkeypatch.setattr(netket_vmc, "compile_operator_sum_netket", fake_common_operator) + common = OperatorSum() + resolved = netket_vmc._normalize_fermion_observables( + "hilbert", + {"hopping": [(1.0, ((0, 1, True), (1, 1, False)))], "common": common}, + site_order=((0, 0),), + ) + + assert resolved == { + "hopping": "fermion-observable", + "common": "common-observable", + } + assert terms_calls[0][-1] == "auto" + assert common_calls[0][-1] == "auto" + + +def test_configure_jax_for_vmc_configures_an_optional_private_cache(monkeypatch): + """Compilation caching is opt-in and set before importing JAX.""" + from pepsy.vmc.netket import configure_jax_for_vmc + + for name in ( + "XLA_PYTHON_CLIENT_PREALLOCATE", + "XLA_PYTHON_CLIENT_MEM_FRACTION", + "JAX_PLATFORMS", + "JAX_COMPILATION_CACHE_DIR", + "NETKET_NO_TIPS", + ): + monkeypatch.delenv(name, raising=False) + + configure_jax_for_vmc( + preallocate=True, + mem_fraction=0.8, + platform="cpu", + compilation_cache_dir="/private/jax-cache", + ) + + assert os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] == "true" + assert os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] == "0.8" + assert os.environ["JAX_PLATFORMS"] == "cpu" + assert os.environ["JAX_COMPILATION_CACHE_DIR"] == "/private/jax-cache" + with pytest.raises(ValueError, match="must not be empty"): + configure_jax_for_vmc(compilation_cache_dir="") + + +def test_netket_autochunk_callback_has_a_clear_version_guard(monkeypatch): + """Old supported NetKet versions fail with an actionable feature error.""" + import pepsy.vmc.netket as netket_vmc + + monkeypatch.setattr( + netket_vmc, + "_require_netket", + lambda: SimpleNamespace(__version__="3.10", callbacks=SimpleNamespace()), + ) + with pytest.raises(RuntimeError, match="requires NetKet >= 3.22"): + netket_vmc.make_netket_autochunk_callback() + + def test_common_spinful_fermion_operator_has_matching_exact_local_energies(): """The Torch and NetKet compilers agree in a tiny exact Fock sector. diff --git a/tests/test_vmc_distributed.py b/tests/test_vmc_distributed.py new file mode 100644 index 0000000..5598f21 --- /dev/null +++ b/tests/test_vmc_distributed.py @@ -0,0 +1,198 @@ +"""Focused checks for optional rank-sharded Torch-VMC reductions.""" + +from types import SimpleNamespace + +import pytest + + +def test_rank_sharded_statistics_all_reduce_only_compact_moments(): + torch = pytest.importorskip("torch") + from pepsy.vmc.torch.distributed import distributed_unweighted_statistics + + class FakeDistributed: + class ReduceOp: + SUM = "sum" + MAX = "max" + + def all_reduce(self, tensor, *, op, group): + del group + assert op == self.ReduceOp.SUM + # The second rank owns values 2 and 3, with local ESS two. + tensor.add_(torch.tensor([5.0, 0.0, 13.0, 2.0, 2.0])) + + runtime = SimpleNamespace( + module=FakeDistributed(), + group=None, + world_size=2, + ) + ( + mean, + variance, + stderr, + naive_stderr, + effective_sample_size, + n_samples, + ) = distributed_unweighted_statistics( + torch.tensor([0.0, 1.0]), + local_effective_sample_size=2.0, + runtime=runtime, + ) + + assert mean.item() == pytest.approx(1.5) + assert variance.item() == pytest.approx(1.25) + assert stderr.item() == pytest.approx((1.25 / 4.0) ** 0.5) + assert naive_stderr.item() == pytest.approx((1.25 / 4.0) ** 0.5) + assert effective_sample_size.item() == pytest.approx(4.0) + assert n_samples == 4 + + +def test_driver_measurement_returns_global_rank_sharded_estimate(monkeypatch): + torch = pytest.importorskip("torch") + import pepsy.vmc.torch.driver as driver_module + from pepsy.vmc.torch import TorchConnections, TorchVMCDriver + + class FakeDistributed: + class ReduceOp: + SUM = "sum" + MAX = "max" + + def __init__(self): + self.scalar_sum_calls = 0 + + def all_reduce(self, tensor, *, op, group): + del group + if op == self.ReduceOp.MAX: + return + if tensor.numel() == 5: + tensor.add_(torch.tensor([5.0, 0.0, 13.0, 2.0, 2.0])) + else: + self.scalar_sum_calls += 1 + if self.scalar_sum_calls == 1: + # The remote rank has the same two local chains. + tensor.add_(2) + + class Amplitude: + def __call__(self, configs): + return configs[:, 0].to(dtype=torch.float64) + 1.0 + + def connections(configs, graph): + del graph + return TorchConnections( + configs=configs.clone(), + coeffs=configs[:, 0].to(dtype=torch.float64), + batch_ids=torch.arange(configs.shape[0]), + ) + + runtime = SimpleNamespace( + module=FakeDistributed(), + group=None, + rank=0, + world_size=2, + backend="gloo", + ) + monkeypatch.setattr(driver_module, "resolve_torch_distributed", lambda _: runtime) + driver = TorchVMCDriver( + Amplitude(), + object(), + torch.tensor([[0], [1]]), + connection_fn=connections, + proposal="spin", + ) + + result = driver.measure_samples( + torch.tensor([[[0], [1]]]), + distributed=True, + ) + + assert result.energy_mean.item() == pytest.approx(1.5) + assert result.energy_variance.item() == pytest.approx(1.25) + assert result.n_samples == 4 + assert result.samples_per_second > 0 + assert result.chain_diagnostics is None + assert result.distributed.global_n_chains == 4 + assert result.distributed.local_n_chains == 2 + assert result.distributed.global_n_samples == 4 + + +def test_rank_shard_counts_cover_global_chains_without_empty_ranks(): + from pepsy.vmc.torch.distributed import shard_chain_count + + class Runtime: + world_size = 3 + + def __init__(self, rank): + self.rank = rank + + assert [shard_chain_count(8, Runtime(rank)) for rank in range(3)] == [3, 3, 2] + with pytest.raises(ValueError, match="at least"): + shard_chain_count(2, Runtime(0)) + + +def test_fermion_lazy_setup_uses_rank_local_seed(monkeypatch): + import pepsy.vmc.torch.fermion as fermion_module + from pepsy.vmc import SamplingConfig + from pepsy.vmc.torch import TorchFermionVMC + + runtime = SimpleNamespace(rank=2, world_size=3) + monkeypatch.setattr( + fermion_module, + "resolve_torch_distributed", + lambda _: runtime, + ) + sampling = SamplingConfig( + n_samples_per_chain=2, + n_chains=8, + seed=11, + ) + + returned_runtime, local_sampling = TorchFermionVMC._rank_sharded_sampling_config( + sampling, + True, + ) + + assert returned_runtime is runtime + assert local_sampling.n_chains == 2 + assert local_sampling.seed == 11 + 2 * 104_729 + + +def test_distributed_sample_metadata_survives_portable_conversion(): + torch = pytest.importorskip("torch") + from pepsy.vmc.torch import TorchDistributedMetadata, TorchMCMCSamples + + distributed = TorchDistributedMetadata( + rank=1, + world_size=2, + backend="gloo", + global_n_chains=4, + local_n_chains=2, + global_n_samples=12, + local_n_samples=6, + ) + native = TorchMCMCSamples( + configs=torch.zeros((3, 2, 1), dtype=torch.long), + amplitudes=torch.ones((3, 2), dtype=torch.float64), + n_samples=6, + n_samples_per_chain=3, + n_chains=2, + n_discard_per_chain=0, + sweep_size=1, + acceptance_rate=0.5, + n_proposed=12, + n_accepted=6, + elapsed_seconds=1.0, + samples_per_second=6.0, + distributed=distributed, + ) + + portable = native.to_common() + + assert portable.n_chains == 2 + assert portable.diagnostics["distributed"] == { + "rank": 1, + "world_size": 2, + "backend": "gloo", + "global_n_chains": 4, + "local_n_chains": 2, + "global_n_samples": 12, + "local_n_samples": 6, + } diff --git a/tests/test_vmc_importance.py b/tests/test_vmc_importance.py index b491c95..d1c239d 100644 --- a/tests/test_vmc_importance.py +++ b/tests/test_vmc_importance.py @@ -151,13 +151,7 @@ def test_real_u1u1_mps_sampler_feeds_fermionic_peps_vmc(): from pepsy.vmc import TorchFermionVMC from pepsy.vmc.netket import fermionic_peps_rand - fermion = Fermion( - spinful=True, - symmetry="U1U1", - t=0.2, - U=1.0, - mu=0.1, - ) + fermion = Fermion(spinful=True, symmetry="U1U1") target = fermionic_peps_rand( "U1U1", 2, @@ -198,9 +192,28 @@ def test_real_u1u1_mps_sampler_feeds_fermionic_peps_vmc(): assert occupations[:, :, 0].sum(axis=1).tolist() == [3] * 4 assert occupations[:, :, 1].sum(axis=1).tolist() == [3] * 4 + sites = tuple((x, y) for x in range(target.Lx) for y in range(target.Ly)) + edges = tuple( + ((x, y), (x + 1, y)) + for x in range(target.Lx - 1) + for y in range(target.Ly) + ) + tuple( + ((x, y), (x, y + 1)) + for x in range(target.Lx) + for y in range(target.Ly - 1) + ) + terms = { + edge: -0.2 * fermion.hopping_operator() + for edge in edges + } + terms |= { + site: fermion.onsite_term(site, U=1.0, mu=0.1) + for site in sites + } vmc = TorchFermionVMC( target, fermion=fermion, + terms=terms, contraction="exact", n_walkers=1, seed=5, diff --git a/tests/test_vmc_local_energy.py b/tests/test_vmc_local_energy.py new file mode 100644 index 0000000..eff10c0 --- /dev/null +++ b/tests/test_vmc_local_energy.py @@ -0,0 +1,209 @@ +"""Regression tests for reusable Torch-VMC local-estimator amplitudes.""" + +import pytest + + +def test_local_energy_reuses_matching_parent_amplitudes_across_walkers(): + """A target already retained by another walker needs no PEPS call.""" + torch = pytest.importorskip("torch") + from pepsy.vmc.torch import TorchConnections, local_energy_from_connections + + class Amplitude: + def __init__(self): + self.connections = None + + def connected_amplitudes(self, configs, amplitudes, connections, **kwargs): + del configs, amplitudes, kwargs + self.connections = connections + return connections.configs.sum(dim=1, dtype=torch.float64) + 2.0 + + configs = torch.tensor([[0, 0], [1, 0]], dtype=torch.long) + amplitudes = torch.tensor([2.0, 3.0], dtype=torch.float64) + connections = TorchConnections( + configs=torch.tensor([[1, 0], [0, 1]], dtype=torch.long), + coeffs=torch.ones(2, dtype=torch.float64), + batch_ids=torch.zeros(2, dtype=torch.long), + ) + amplitude = Amplitude() + + values = local_energy_from_connections( + configs, + amplitudes, + connections, + amplitude, + deduplicate_targets=True, + ) + + assert amplitude.connections.configs.tolist() == [[0, 1]] + assert amplitude.connections.batch_ids.tolist() == [0] + assert torch.allclose( + values, + torch.tensor([3.0, 0.0], dtype=torch.float64), + ) + + +def test_boundary_connected_fallbacks_are_dispatched_as_one_cached_batch(): + """Unresolved boundary targets use the cache-aware forward route.""" + torch = pytest.importorskip("torch") + from pepsy.vmc.torch import TorchConnections, TorchPEPSBoundaryAmplitude + + class BoundaryProbe(TorchPEPSBoundaryAmplitude): + def __init__(self): + self.contraction = "boundary" + self._boundary_geometry = object() + self.amplitude_batching = "serial" + self._connection_vmap_enabled = False + self.forward_batches = [] + + def _ensure_boundary_cache_current(self): + pass + + def _unpack_tn(self): + return object() + + def _reference_tensor(self): + return torch.tensor(0.0) + + def _changed_axis_windows(self, parent_config, target_config): + del parent_config, target_config + return () + + def forward(self, configs, params=None, *, chunk_size=None): + del params, chunk_size + self.forward_batches.append(configs.clone()) + return configs.sum(dim=1, dtype=torch.float64) + + model = BoundaryProbe() + configs = torch.tensor([[0], [1]], dtype=torch.long) + amplitudes = torch.tensor([1.0, 2.0], dtype=torch.float64) + connections = TorchConnections( + configs=torch.tensor([[2], [3]], dtype=torch.long), + coeffs=torch.ones(2, dtype=torch.float64), + batch_ids=torch.tensor([0, 1], dtype=torch.long), + ) + + values = model.connected_amplitudes(configs, amplitudes, connections) + + assert len(model.forward_batches) == 1 + assert model.forward_batches[0].tolist() == [[2], [3]] + assert torch.allclose( + values, + torch.tensor([2.0, 3.0], dtype=torch.float64), + ) + assert model.last_connected_reuse_stats["num_fallback"] == 2 + + +def test_boundary_workers_parallelize_no_grad_primary_windows(): + """Independent cached boundary windows can run in parallel on CPU.""" + torch = pytest.importorskip("torch") + from pepsy.vmc.torch import TorchConnections, TorchPEPSBoundaryAmplitude + + class BoundaryProbe(TorchPEPSBoundaryAmplitude): + def __init__(self): + self.contraction = "boundary" + self._boundary_geometry = object() + self.amplitude_batching = "serial" + self.graded_torch = False + self._connection_vmap_enabled = False + self.boundary_workers = 2 + self._boundary_environment_cache = {} + self._boundary_strip_cache = {} + self.last_amplitude_cache_stats = {"stale": 1} + + def _ensure_boundary_cache_current(self): + pass + + def _unpack_tn(self): + return object() + + def _select_config(self, tn, config): + del tn, config + return object() + + def _reference_tensor(self): + return torch.tensor(0.0) + + @staticmethod + def _configuration_key(config): + return tuple(int(value) for value in config.tolist()) + + def _changed_axis_windows(self, parent_config, target_config): + del parent_config, target_config + return (("x", (0,)), ("y", (0,))) + + def _cached_boundary_environments(self, *args, **kwargs): + del args, kwargs + return object(), False + + def _cached_boundary_strip(self, *args, **kwargs): + del args, kwargs + return object(), False + + def _contract_cached_axis_window( + self, + tn, + parent_config, + target_config, + axis, + indices, + envs, + strip_tn, + reference, + ): + del tn, parent_config, axis, indices, envs, strip_tn, reference + return target_config.sum(dtype=torch.float64) + + model = BoundaryProbe() + configs = torch.tensor([[0, 0], [1, 0]], dtype=torch.long) + amplitudes = torch.ones(2, dtype=torch.float64) + connections = TorchConnections( + configs=torch.tensor([[1, 1], [0, 1]], dtype=torch.long), + coeffs=torch.ones(2, dtype=torch.float64), + batch_ids=torch.tensor([0, 1], dtype=torch.long), + ) + + with torch.no_grad(): + values = model.connected_amplitudes(configs, amplitudes, connections) + + assert torch.allclose(values, torch.tensor([2.0, 1.0], dtype=torch.float64)) + assert model.last_amplitude_cache_stats is None + assert model.last_connected_reuse_stats["num_requests"] == 2 + assert model.last_connected_reuse_stats["num_parallel"] == 2 + assert model.last_connected_reuse_stats["num_reused"] == 2 + + +def test_amplitude_benchmark_records_chunks_and_restores_fast_path_state(): + """Benchmarking vmap candidates does not alter the production policy.""" + torch = pytest.importorskip("torch") + from pepsy.vmc.torch import benchmark_torch_amplitudes + + class Amplitude: + amplitude_batching = "auto" + _vmap_forward_enabled = True + last_amplitude_batching = None + + def __init__(self): + self.boundary_cache_size = 7 + self._boundary_amplitude_cache = {"retained": torch.tensor(1.0)} + + def __call__(self, configs): + self.last_amplitude_batching = self.amplitude_batching + return configs.sum(dim=1, dtype=torch.float64) + + amplitude = Amplitude() + run = benchmark_torch_amplitudes( + amplitude, + torch.tensor([[0, 1], [2, 3], [4, 5]], dtype=torch.long), + chunk_sizes=(None, 2), + amplitude_batchings=("serial", "vmap"), + warmup=0, + repeats=1, + ) + + assert len(run.entries) == 4 + assert run.best in run.entries + assert {entry.executed_batching for entry in run.entries} == {"serial", "vmap"} + assert amplitude.amplitude_batching == "auto" + assert amplitude._vmap_forward_enabled + assert amplitude.boundary_cache_size == 7 + assert list(amplitude._boundary_amplitude_cache) == ["retained"] From 123ab2cbbeb7257e9e9592a3104291c1cc14f3ba Mon Sep 17 00:00:00 2001 From: rezaquant Date: Wed, 29 Jul 2026 08:05:31 -0700 Subject: [PATCH 17/70] add cross-simulator planner --- CHANGELOG.md | 6 + docs/api/index.md | 1 + docs/api/optimizers/planning.md | 108 ++++ docs/development/modules/optimizers.md | 3 + src/pepsy/__init__.py | 4 + src/pepsy/optimizers/__init__.py | 5 + src/pepsy/optimizers/planning.py | 666 +++++++++++++++++++++++++ tests/test_package_layout.py | 8 + tests/test_public_api.py | 2 + tests/test_simulator_planner.py | 94 ++++ 10 files changed, 897 insertions(+) create mode 100644 docs/api/optimizers/planning.md create mode 100644 src/pepsy/optimizers/planning.py create mode 100644 tests/test_simulator_planner.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 30cb22c..698ebc0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,12 @@ Pepsy follows [Semantic Versioning](https://semver.org/): Changes for the next release should be added here before the version is bumped. +### Added + +- `SimulatorPlanner` and `recommend_simulator` provide non-executing, + chi-aware rankings across MPS, tree, MPS-stabilizer, and tree-stabilizer + circuit strategies using physical and dressed-frame support geometry. + ## [0.4.0] - 2026-07-27 This release removes obsolete package-layout compatibility layers and keeps diff --git a/docs/api/index.md b/docs/api/index.md index f24820f..4564b29 100644 --- a/docs/api/index.md +++ b/docs/api/index.md @@ -32,6 +32,7 @@ the detailed functions and classes for each area. - [Energy optimization](optimizers/energy.md) - [MERA](optimizers/mera.md) - [Noise and trajectories](optimizers/noise.md) +- [Simulator planning](optimizers/planning.md) - [Stabilizer tensor networks](optimizers/stabilizer_tn.md) - [Symmetry-aware DMRG](optimizers/sym_dmrg.md) - [Tree tensor networks](optimizers/tree.md) diff --git a/docs/api/optimizers/planning.md b/docs/api/optimizers/planning.md new file mode 100644 index 0000000..81a4ef7 --- /dev/null +++ b/docs/api/optimizers/planning.md @@ -0,0 +1,108 @@ +# Simulator planning + +`SimulatorPlanner` compares four Pepsy circuit strategies before any circuit +execution: + +- `MpsOptimizer` +- `TreeOptimizer` +- `MpsStabOptimizer` +- `TreeStabOptimizer` + +The result is typed, ranked advice. It includes the selected MPS order or tree +plan, the physical and stabilizer-frame supports used for pricing, and every +cost component needed to audit the recommendation. + +```python +import pepsy + +stream = [ + ("h", 0), + ("cnot", 0, 1), + ("cnot", 1, 2), + ("rz", 0.31, 2), +] + +advice = pepsy.recommend_simulator( + stream, + n_qubits=3, + chi=64, +) + +print(advice.recommended) +for candidate in advice.candidates: + print( + candidate.optimizer, + candidate.applicable, + candidate.relative_score, + candidate.max_geometry, + ) +``` + +`SimulatorPlanner(...).plan()` and `.recommend()` are equivalent. The +convenience function above returns a `SimulatorPlan`; its `best` property is +the first applicable `SimulatorCandidate`, and +`candidate("TreeStabOptimizer")` selects one record by class name. Both record +types also support mapping-style access and `as_dict()`. + +## What is priced + +The planner first uses `MpsStabOptimizer.analyze_stream` to count Clifford, +injectable, other non-Clifford, structural, and opaque entries. It then builds +two circuit descriptions: + +1. Physical supports price all known circuit events for the ordinary MPS and + tree candidates. +2. A tableau-only dry run removes Clifford events from coefficient-network + work and records the actual support of each current dressed operator + `C† O C` for the stabilizer candidates. + +Each description receives its own layout. MPS candidates use optimized +one-dimensional window widths. Tree candidates use a `TreeLayoutFinder` plan +and the number of nodes in the minimal connected subtree spanning each +support. Consequently, the comparison can distinguish a physically local gate +that becomes a long dressed Pauli from one that remains local in the +stabilizer frame. + +At target bond dimension `chi`, the default work proxy is + +```text +one-site event: weight * chi^2 +routed event: weight * geometry * chi^3 +dressed Pauli event: 16 * the corresponding tensor work +frame bookkeeping: n_qubits * (Clifford entries + dressed events) +``` + +The dressed-MPO factor of 16 is the operation-count constant used for the +bond-dimension-two HSMPO contraction in +[MPStab](https://arxiv.org/abs/2607.24258). Pepsy differs from that global +estimate by pricing the measured dressed support and by updating one live +Clifford frame incrementally. Override `frame_mpo_factor` or +`tableau_factor` when calibration against local benchmarks supports a +different machine-specific ratio. + +`weight_mode="count"` is the default. `"angle"` and `"auto"` reduce the weight +of named rotations with small absolute angles. The default MPS search is +deterministic and does not invoke optional Nevergrad or KaHyPar backends. +`tree_layout_kwargs` can override ordinary `TreeLayoutFinder` structure and +objective options; the circuit, qubit count, `chi`, and event weights remain +planner-owned. + +## Interpretation and limits + +Scores are relative static work proxies, not runtime, memory, fidelity, or +error bounds. A recommendation does not replace a `chi` convergence sweep. +Exact cooling and immediate or deferred magic-state injection can materially +change stabilizer peak bonds, so benchmark candidates with similar scores. +The returned stabilizer settings enable exact cooling and infidelity tracking, +but the conservative score itself models direct dressed-operator replay. + +Unknown physical supports are conservatively priced across all qubits. If the +tableau dry run cannot be defined—for example, a `cap` changes the register +length or feed-forward makes the static frame branch-dependent—the ordinary +candidates remain ranked and both stabilizer candidates are returned with +`applicable=False` and an explanatory warning. + +For a purely Clifford circuit, the planner ranks the stabilizer tensor-network +variants cheaply because their coefficient networks are unchanged. Use a +tableau simulator such as Stim instead when no tensor-network state access is +needed. diff --git a/docs/development/modules/optimizers.md b/docs/development/modules/optimizers.md index ec53e5a..7a472e9 100644 --- a/docs/development/modules/optimizers.md +++ b/docs/development/modules/optimizers.md @@ -32,6 +32,9 @@ important downstream time-compression consumer that depends on Pepsy behavior. simulator. See `../plans/stabilizer_tn.md` for its implementation record and `docs/howto/stabilizer_tn_magic.md` for exact cooling, greedy checkpoints, and immediate versus deferred MAST injection. +- `planning.py`: non-executing physical-versus-stabilizer and + MPS-versus-tree circuit advice using measured frame supports and explicit + chi-scaled work proxies. - `mera/`: dense MERA and schedule-first qMERA local-energy objectives, parameter dictionaries, compiled lightcone contractions, schematics, and Symmray-native fermion helpers. diff --git a/src/pepsy/__init__.py b/src/pepsy/__init__.py index 7d8a48f..e4d8a21 100644 --- a/src/pepsy/__init__.py +++ b/src/pepsy/__init__.py @@ -101,6 +101,10 @@ "MpsEnergyOptimizer": ".optimizers", "MpsOptimizer": ".optimizers", "MpsStabOptimizer": ".optimizers", + "SimulatorCandidate": ".optimizers", + "SimulatorPlan": ".optimizers", + "SimulatorPlanner": ".optimizers", + "recommend_simulator": ".optimizers", "DeferredInjectionRecord": ".optimizers", "DeferredInjectionReport": ".optimizers", "DeferredProjectionRecord": ".optimizers", diff --git a/src/pepsy/optimizers/__init__.py b/src/pepsy/optimizers/__init__.py index fcb001e..07ff161 100644 --- a/src/pepsy/optimizers/__init__.py +++ b/src/pepsy/optimizers/__init__.py @@ -24,6 +24,10 @@ "build_qmera_contraction_optimizer": ".mera", "MpoOptimizer": ".mpo", "MpsOptimizer": ".mps", + "SimulatorCandidate": ".planning", + "SimulatorPlan": ".planning", + "SimulatorPlanner": ".planning", + "recommend_simulator": ".planning", "PepsOptimizer": ".peps", "SimpleUpdateGen": ".peps", "SymDMRG2": ".sym_dmrg", @@ -99,6 +103,7 @@ "mpo", "mps", "noise", + "planning", "peps", "stabilizer_tn", "sym_dmrg", diff --git a/src/pepsy/optimizers/planning.py b/src/pepsy/optimizers/planning.py new file mode 100644 index 0000000..deb95aa --- /dev/null +++ b/src/pepsy/optimizers/planning.py @@ -0,0 +1,666 @@ +"""Static cross-simulator advice for circuit gate streams. + +The planner compares Pepsy's ordinary MPS/tree simulators with their +stabilizer-frame counterparts. It never executes the circuit or constructs a +live ordinary tensor-network state. Its scores are transparent operation-count +proxies at a caller-supplied bond dimension, not wall-clock predictions. +""" + +from __future__ import annotations + +from collections.abc import MutableMapping +from dataclasses import dataclass, fields +from numbers import Integral +from typing import Optional + +import numpy as np + +from .mps import MpsGateStreamLayoutFinder +from .stabilizer_tn import MpsStabOptimizer, StreamAnalysisRecord +from .tree import TreeLayoutFinder, TreePlan + +__all__ = [ + "SimulatorCandidate", + "SimulatorPlan", + "SimulatorPlanner", + "recommend_simulator", +] + + +_ROTATION_NAMES = frozenset( + {"rot", "rx", "ry", "rz", "rxx", "ryy", "rzz", "t", "tdg"} +) +_CANDIDATE_ORDER = { + "MpsOptimizer": 0, + "TreeOptimizer": 1, + "MpsStabOptimizer": 2, + "TreeStabOptimizer": 3, +} + + +class _PlannerRecord(MutableMapping): + """Mapping-compatible facade shared by the typed planner records.""" + + @classmethod + def _field_names(cls): + return tuple(field.name for field in fields(cls)) + + def __getitem__(self, key): + if key in self._field_names(): + return getattr(self, key) + raise KeyError(key) + + def __setitem__(self, key, value): + if key not in self._field_names(): + raise KeyError(key) + setattr(self, key, value) + + def __delitem__(self, key): # pragma: no cover - fixed-shape diagnostics + raise TypeError(f"{type(self).__name__} fields cannot be deleted.") + + def __iter__(self): + return iter(self._field_names()) + + def __len__(self): + return len(self._field_names()) + + def as_dict(self) -> dict: + """Return a shallow plain-``dict`` snapshot.""" + return {name: getattr(self, name) for name in self._field_names()} + + +@dataclass +class SimulatorCandidate(_PlannerRecord): + """One ranked simulator choice and its auditable cost components.""" + + optimizer: str + geometry: str + stabilizer_frame: bool + applicable: bool + score: Optional[float] + relative_score: Optional[float] + tensor_work: Optional[float] + tableau_work: Optional[float] + event_count: int + one_site_events: int + multi_site_events: int + weighted_geometry: float + max_geometry: int + layout: object + layout_report: dict + settings: dict + rationale: str + warnings: tuple[str, ...] = () + + +@dataclass +class SimulatorPlan(_PlannerRecord): + """Ranked, non-executing advice for one circuit stream.""" + + recommended: str + candidates: tuple[SimulatorCandidate, ...] + analysis: StreamAnalysisRecord + n_qubits: int + chi: int + physical_events: tuple[dict, ...] + frame_events: tuple[dict, ...] + weight_mode: str + cost_model: str + warnings: tuple[str, ...] + + @property + def best(self) -> SimulatorCandidate: + """Return the first (lowest-score) applicable candidate.""" + return self.candidates[0] + + def candidate(self, optimizer: str) -> SimulatorCandidate: + """Return advice for ``optimizer`` by public class name.""" + for candidate in self.candidates: + if candidate.optimizer == optimizer: + return candidate + choices = ", ".join(item.optimizer for item in self.candidates) + raise KeyError(f"Unknown simulator {optimizer!r}; choices are: {choices}.") + + +def _unique_warnings(warnings): + return tuple(dict.fromkeys(str(warning) for warning in warnings if warning)) + + +def _normalize_weight_mode(weight_mode): + name = str(weight_mode).replace("-", "_").strip().lower() + aliases = {"unit": "count", "uniform": "count", "none": "count"} + name = aliases.get(name, name) + if name not in {"count", "angle", "auto"}: + raise ValueError( + "weight_mode must be 'count', 'angle', or 'auto', " + f"got {weight_mode!r}." + ) + return name + + +def _event_weight(entry, *, weight_mode): + """Return the physical-stream weight used by the static layout search.""" + if weight_mode == "count": + return 1.0 + if ( + isinstance(entry, (tuple, list)) + and len(entry) >= 2 + and isinstance(entry[0], str) + and str(entry[0]).replace("-", "_").strip().lower() in _ROTATION_NAMES + ): + try: + angle = abs(float(entry[1])) + except (TypeError, ValueError): + return 1.0 + if np.isfinite(angle): + return min(1.0, max(0.0, angle)) + return 1.0 + + +def _layout_stream(records): + """Encode weighted supports through the public layout-stream protocol.""" + return [ + ( + "submpo", + {"angle": float(record["weight"])}, + tuple(record["support"]), + ) + for record in records + ] + + +def _layout_weight(_payload, _support, _event_type): + return float(_payload.get("angle", 1.0)) + + +def _mps_geometry(records, plan): + position = {int(site): int(pos) for site, pos in plan["site_map"].items()} + geometry = [] + for record in records: + support = tuple(dict.fromkeys(int(site) for site in record["support"])) + if len(support) <= 1: + geometry.append(1) + continue + positions = [position[site] for site in support] + geometry.append(max(positions) - min(positions) + 1) + return tuple(geometry) + + +def _tree_steiner_size(plan: TreePlan, support): + support = tuple(dict.fromkeys(int(site) for site in support)) + if len(support) <= 1: + return 1 + nodes = [plan.node_of_qubit[site] for site in support] + steiner = {nodes[0]} + for node in nodes[1:]: + steiner.update(plan.node_path(nodes[0], node)) + return len(steiner) + + +def _tree_geometry(records, plan): + return tuple( + _tree_steiner_size(plan, record["support"]) for record in records + ) + + +def _work_components( + records, + geometry, + *, + chi, + operator_factor, + tableau_work=0.0, +): + one_site_events = 0 + multi_site_events = 0 + weighted_geometry = 0.0 + local_weight = 0.0 + routed_weight = 0.0 + max_geometry = 0 + + for record, extent in zip(records, geometry): + support = tuple(dict.fromkeys(record["support"])) + weight = float(record["weight"]) + extent = int(extent) + max_geometry = max(max_geometry, extent) + if len(support) <= 1: + one_site_events += 1 + local_weight += weight + else: + multi_site_events += 1 + weighted_geometry += weight * extent + routed_weight += weight * extent + + chi_float = float(chi) + tensor_work = float(operator_factor) * ( + local_weight * chi_float**2 + routed_weight * chi_float**3 + ) + score = tensor_work + float(tableau_work) + return { + "score": float(score), + "tensor_work": float(tensor_work), + "tableau_work": float(tableau_work), + "event_count": len(records), + "one_site_events": int(one_site_events), + "multi_site_events": int(multi_site_events), + "weighted_geometry": float(weighted_geometry), + "max_geometry": int(max_geometry), + } + + +class SimulatorPlanner: + """Rank MPS, tree, MPS-STN, and tree-STN strategies without executing. + + Parameters + ---------- + gates + A Pepsy-native gate stream accepted by + :meth:`MpsStabOptimizer.analyze_stream`. + n_qubits + Circuit width. It is inferred from known supports when omitted. + chi + Target tensor-network bond dimension used by the work proxy. + weight_mode + ``"count"`` weights every priced event equally. ``"angle"`` and + ``"auto"`` down-weight named rotations with small absolute angles. + mps_order + Deterministic MPS layout candidate used for physical and dressed + supports. The default avoids optional offline search backends. + tree_layout_kwargs + Optional keyword overrides for :class:`TreeLayoutFinder`. ``n``, + ``gates``, ``supports``, ``chi``, and ``weight_mode`` are planner-owned. + frame_mpo_factor + Relative contraction constant for a dressed Pauli MPO. The default 16 + follows the bond-dimension-two HSMPO operation-count estimate. + tableau_factor + Relative cost assigned to each qubit touched by an incremental + Clifford-frame update or dressed-Pauli query. + """ + + def __init__( + self, + gates, + *, + n_qubits=None, + chi=64, + weight_mode="count", + mps_order="recursive_refined", + tree_layout_kwargs=None, + frame_mpo_factor=16.0, + tableau_factor=1.0, + ): + if isinstance(chi, bool) or not isinstance(chi, Integral): + raise TypeError("chi must be a positive integer.") + chi = int(chi) + if chi <= 0: + raise ValueError("chi must be a positive integer.") + if not np.isfinite(float(chi) ** 3): + raise ValueError("chi is too large for the planner work proxy.") + + frame_mpo_factor = float(frame_mpo_factor) + tableau_factor = float(tableau_factor) + if not np.isfinite(frame_mpo_factor) or frame_mpo_factor <= 0.0: + raise ValueError("frame_mpo_factor must be finite and positive.") + if not np.isfinite(tableau_factor) or tableau_factor < 0.0: + raise ValueError("tableau_factor must be finite and nonnegative.") + + if gates is None or ( + isinstance(gates, (tuple, list)) and len(gates) == 0 + ): + self.entries = () + else: + self.entries = tuple(MpsStabOptimizer._as_entries(gates)) + self.analysis = MpsStabOptimizer.analyze_stream( + None if not self.entries else self.entries, + n_qubits=n_qubits, + ) + inferred = self.analysis.estimated_qubits + if inferred is None or inferred <= 0: + raise ValueError( + "Could not infer a positive circuit width; pass n_qubits explicitly." + ) + + tree_layout_kwargs = ( + {} if tree_layout_kwargs is None else dict(tree_layout_kwargs) + ) + forbidden = { + "chi", + "gates", + "n", + "supports", + "weight_mode", + }.intersection(tree_layout_kwargs) + if forbidden: + names = ", ".join(sorted(forbidden)) + raise ValueError( + f"tree_layout_kwargs contains planner-owned option(s): {names}." + ) + + self.n_qubits = int(inferred) + self.chi = chi + self.weight_mode = _normalize_weight_mode(weight_mode) + self.mps_order = str(mps_order) + self.tree_layout_kwargs = tree_layout_kwargs + self.frame_mpo_factor = frame_mpo_factor + self.tableau_factor = tableau_factor + + def _physical_records(self): + records = [] + warnings = [] + for index, entry in enumerate(self.entries): + kind = MpsStabOptimizer._analysis_entry_kind(entry) + if kind == "control": + continue + try: + sites = MpsStabOptimizer._analysis_entry_sites( + entry, + self.n_qubits, + ) + except (IndexError, TypeError, ValueError): + sites = None + estimated = sites is None + if sites is None: + sites = set(range(self.n_qubits)) + warnings.append( + f"Entry {index} has unknown support and was conservatively " + "priced across all qubits." + ) + support = tuple(sorted(int(site) for site in sites)) + if not support: + continue + records.append( + { + "index": int(index), + "kind": str(kind), + "support": support, + "weight": _event_weight( + entry, + weight_mode=self.weight_mode, + ), + "estimated_support": bool(estimated), + } + ) + return tuple(records), tuple(warnings) + + def _frame_records(self): + if not self.entries: + return () + simulator = MpsStabOptimizer( + self.n_qubits, + gates=self.entries, + chi=self.chi, + exact_cooling=False, + ) + return simulator._frame_layout_records( + self.entries, + weight_mode=self.weight_mode, + ) + + def _mps_layout(self, records): + finder = MpsGateStreamLayoutFinder( + _layout_stream(records), + L=self.n_qubits, + ) + order = self.mps_order + if not any(len(record["support"]) >= 2 for record in records): + order = "input" + return finder.run( + order=order, + weight_fn=_layout_weight, + weight_mode="count", + ) + + def _tree_layout(self, records): + kwargs = { + "structure": "quality", + "max_arity": (2, 3, 4), + "objective": "path", + **self.tree_layout_kwargs, + } + finder = TreeLayoutFinder( + _layout_stream(records), + n=self.n_qubits, + chi=self.chi, + weight_mode="angle", + **kwargs, + ) + plan = finder.run() + return plan, finder.report(plan) + + def _candidate( + self, + *, + optimizer, + geometry_name, + stabilizer_frame, + records, + geometry, + layout, + layout_report, + operator_factor, + tableau_work, + warnings=(), + ): + parts = _work_components( + records, + geometry, + chi=self.chi, + operator_factor=operator_factor, + tableau_work=tableau_work, + ) + basis = "dressed frame" if stabilizer_frame else "physical" + geometry_label = ( + "MPS window width" if geometry_name == "mps" else "tree Steiner size" + ) + rationale = ( + f"Prices {parts['event_count']} {basis} event(s) using " + f"{geometry_label}; {parts['multi_site_events']} require routed " + f"multi-site work." + ) + settings = {"chi": self.chi, "layout": layout} + if stabilizer_frame: + settings.update({"exact_cooling": True, "track_infidelity": True}) + return SimulatorCandidate( + optimizer=optimizer, + geometry=geometry_name, + stabilizer_frame=stabilizer_frame, + applicable=True, + score=parts["score"], + relative_score=None, + tensor_work=parts["tensor_work"], + tableau_work=parts["tableau_work"], + event_count=parts["event_count"], + one_site_events=parts["one_site_events"], + multi_site_events=parts["multi_site_events"], + weighted_geometry=parts["weighted_geometry"], + max_geometry=parts["max_geometry"], + layout=layout, + layout_report=layout_report, + settings=settings, + rationale=rationale, + warnings=_unique_warnings(warnings), + ) + + def _unavailable_stabilizer_candidate(self, optimizer, geometry, reason): + return SimulatorCandidate( + optimizer=optimizer, + geometry=geometry, + stabilizer_frame=True, + applicable=False, + score=None, + relative_score=None, + tensor_work=None, + tableau_work=None, + event_count=0, + one_site_events=0, + multi_site_events=0, + weighted_geometry=0.0, + max_geometry=0, + layout=None, + layout_report={}, + settings={"chi": self.chi}, + rationale="The stabilizer-frame dry run could not price this stream.", + warnings=(str(reason),), + ) + + def plan(self) -> SimulatorPlan: + """Return ranked advice while leaving the circuit and simulators untouched.""" + physical_records, physical_warnings = self._physical_records() + physical_mps_layout = self._mps_layout(physical_records) + physical_tree_layout, physical_tree_report = self._tree_layout( + physical_records + ) + + candidates = [ + self._candidate( + optimizer="MpsOptimizer", + geometry_name="mps", + stabilizer_frame=False, + records=physical_records, + geometry=_mps_geometry(physical_records, physical_mps_layout), + layout=physical_mps_layout, + layout_report={"stats": physical_mps_layout["stats"]}, + operator_factor=1.0, + tableau_work=0.0, + warnings=physical_warnings, + ), + self._candidate( + optimizer="TreeOptimizer", + geometry_name="tree", + stabilizer_frame=False, + records=physical_records, + geometry=_tree_geometry(physical_records, physical_tree_layout), + layout=physical_tree_layout, + layout_report=physical_tree_report, + operator_factor=1.0, + tableau_work=0.0, + warnings=physical_warnings, + ), + ] + + frame_records = () + frame_failure = None + try: + frame_records = self._frame_records() + except (ImportError, IndexError, TypeError, ValueError, RuntimeError) as exc: + frame_failure = ( + "Stabilizer-frame prepass failed: " + f"{type(exc).__name__}: {exc}" + ) + + if frame_failure is None: + frame_mps_layout = self._mps_layout(frame_records) + frame_tree_layout, frame_tree_report = self._tree_layout(frame_records) + tableau_work = self.tableau_factor * self.n_qubits * ( + self.analysis.clifford_entries + len(frame_records) + ) + candidates.extend( + [ + self._candidate( + optimizer="MpsStabOptimizer", + geometry_name="mps", + stabilizer_frame=True, + records=frame_records, + geometry=_mps_geometry(frame_records, frame_mps_layout), + layout=frame_mps_layout, + layout_report={"stats": frame_mps_layout["stats"]}, + operator_factor=self.frame_mpo_factor, + tableau_work=tableau_work, + ), + self._candidate( + optimizer="TreeStabOptimizer", + geometry_name="tree", + stabilizer_frame=True, + records=frame_records, + geometry=_tree_geometry(frame_records, frame_tree_layout), + layout=frame_tree_layout, + layout_report=frame_tree_report, + operator_factor=self.frame_mpo_factor, + tableau_work=tableau_work, + ), + ] + ) + else: + candidates.extend( + [ + self._unavailable_stabilizer_candidate( + "MpsStabOptimizer", + "mps", + frame_failure, + ), + self._unavailable_stabilizer_candidate( + "TreeStabOptimizer", + "tree", + frame_failure, + ), + ] + ) + + candidates.sort( + key=lambda candidate: ( + not candidate.applicable, + float("inf") if candidate.score is None else candidate.score, + _CANDIDATE_ORDER[candidate.optimizer], + ) + ) + applicable_scores = [ + candidate.score + for candidate in candidates + if candidate.applicable and candidate.score is not None + ] + minimum = min(applicable_scores) + for candidate in candidates: + if not candidate.applicable or candidate.score is None: + continue + if minimum == 0.0: + candidate.relative_score = ( + 1.0 if candidate.score == 0.0 else float("inf") + ) + else: + candidate.relative_score = float(candidate.score / minimum) + + warnings = [ + *self.analysis.warnings, + *physical_warnings, + ( + "Scores are static chi-scaled work proxies, not measured runtime " + "or accuracy guarantees; benchmark close candidates." + ), + ( + "Stabilizer scores model direct dressed-operator replay. Exact " + "cooling and immediate/deferred magic injection can change peak " + "bond dimension and runtime." + ), + ] + if frame_failure is not None: + warnings.append(frame_failure) + if self.analysis.is_clifford_only: + warnings.append( + "For a Clifford-only circuit, a tableau simulator such as Stim " + "is normally preferable when full tensor-network state access " + "is unnecessary." + ) + + return SimulatorPlan( + recommended=candidates[0].optimizer, + candidates=tuple(candidates), + analysis=self.analysis, + n_qubits=self.n_qubits, + chi=self.chi, + physical_events=physical_records, + frame_events=tuple(frame_records), + weight_mode=self.weight_mode, + cost_model=( + "local: weight*chi^2; routed: weight*geometry*chi^3; " + f"dressed-MPO factor: {self.frame_mpo_factor:g}; " + f"tableau factor: {self.tableau_factor:g}" + ), + warnings=_unique_warnings(warnings), + ) + + recommend = plan + + +def recommend_simulator(gates, **kwargs) -> SimulatorPlan: + """Convenience wrapper for :class:`SimulatorPlanner`.""" + return SimulatorPlanner(gates, **kwargs).plan() diff --git a/tests/test_package_layout.py b/tests/test_package_layout.py index 5e9cfbc..fa8da18 100644 --- a/tests/test_package_layout.py +++ b/tests/test_package_layout.py @@ -27,11 +27,15 @@ QMeraBuilder, QMeraGeometry, QMeraParametricEnergyOptimizer, + SimulatorCandidate, + SimulatorPlan, + SimulatorPlanner, SimpleUpdateGen, SymDMRG2, SweepOptimizer, build_qmera_contraction_optimizer, mera as mera_module, + recommend_simulator, ) from pepsy.sampling import MpsSampler, PepsBpSampler from pepsy.solvers import FDSolver @@ -118,7 +122,11 @@ def test_new_namespace_imports_resolve(): assert QMeraBuilder is not None assert QMeraGeometry is not None assert QMeraParametricEnergyOptimizer is not None + assert SimulatorCandidate is not None + assert SimulatorPlan is not None + assert SimulatorPlanner is not None assert callable(build_qmera_contraction_optimizer) + assert callable(recommend_simulator) assert mera_module is not None assert SimpleUpdateGen is not None assert SymDMRG2 is not None diff --git a/tests/test_public_api.py b/tests/test_public_api.py index 7c64c83..8e4adb9 100644 --- a/tests/test_public_api.py +++ b/tests/test_public_api.py @@ -57,6 +57,7 @@ def test_tree_optimizers_are_available_from_high_level_api(): "PepsBpSampler", "MpsSampler", "FermionConfigurationEncoding", "MpsDiagonalEstimate", "MpsBatchSampleResult", "MpsSampleResult", "VecSampler", "gate", "gauge_all", "gauge_all_simple", "one_norm_bp", "tn_fidelity", "tn_norm", "TreeSampler", "TreeBatchSampleResult", "TreeSampleResult", "MpsStabOptimizer", "STNState", "StabilizerMpsSimulator", + "SimulatorCandidate", "SimulatorPlan", "SimulatorPlanner", "recommend_simulator", "TreeEnergyOptimizer", "TreeLayoutFinder", "TreeOptimizer", @@ -127,6 +128,7 @@ def test_internal_symbols_not_exported(): "rxx", "ryy", "rzz", "u3", "su4", "fsim", "fsimg", "haar_random_state", "hrs_to_mps", "hrs_to_peps", "hrs_to_ttn", "ps_to_peps", "ps_to_3dpeps", "expec_mpo", "id_to_mpo", "id_to_pepo", "ps_to_pepo", "ps_to_mpo", "ps_to_ttn", "SweepOptimizer", "FDSolver", "MpsEnergyOptimizer", "MpsOptimizer", "MpoOptimizer", "MpsStabOptimizer", "StabilizerMpsSimulator", + "SimulatorCandidate", "SimulatorPlan", "SimulatorPlanner", "recommend_simulator", "DeferredInjectionRecord", "DeferredInjectionReport", "DeferredProjectionRecord", "ImmediateInjectionReport", "ImmediateProjectionRecord", "MeasurementRecord", "NormEventRecord", "StabilizerMpsSettingsAdvice", "StabilizerMpsRunResult", "StreamAnalysisRecord", diff --git a/tests/test_simulator_planner.py b/tests/test_simulator_planner.py new file mode 100644 index 0000000..e676b06 --- /dev/null +++ b/tests/test_simulator_planner.py @@ -0,0 +1,94 @@ +"""Focused tests for cross-simulator static planning.""" + +import pytest + +import pepsy +from pepsy.optimizers import ( + SimulatorCandidate, + SimulatorPlan, + SimulatorPlanner, + recommend_simulator, +) + + +def test_planner_prices_actual_dressed_supports_and_ranks_four_candidates(): + """Clifford-heavy work should expose and benefit from its dressed support.""" + stream = [("h", 0)] * 300 + [ + ("cnot", 0, 1), + ("cnot", 1, 2), + ("rz", 0.3, 2), + ] + + advice = recommend_simulator(stream, n_qubits=3, chi=4) + + assert isinstance(advice, SimulatorPlan) + assert advice.recommended == "MpsStabOptimizer" + assert advice.best is advice.candidates[0] + assert {candidate.optimizer for candidate in advice.candidates} == { + "MpsOptimizer", + "TreeOptimizer", + "MpsStabOptimizer", + "TreeStabOptimizer", + } + assert all( + isinstance(candidate, SimulatorCandidate) + for candidate in advice.candidates + ) + assert advice.frame_events[-1]["support"] == (0, 1, 2) + assert advice.candidate("MpsStabOptimizer").max_geometry == 3 + assert advice.candidates[0].relative_score == pytest.approx(1.0) + assert advice["analysis"].clifford_entries == 302 + + +def test_planner_exposes_chain_windows_and_tree_steiner_sizes(): + """Geometry diagnostics should use each optimizer's routed structure.""" + stream = [ + ("cnot", 0, 4), + ("cnot", 0, 1), + ("cnot", 0, 2), + ("cnot", 0, 3), + ] + + advice = SimulatorPlanner(stream, n_qubits=5, chi=8).plan() + mps = advice.candidate("MpsOptimizer") + tree = advice.candidate("TreeOptimizer") + + assert mps.geometry == "mps" + assert tree.geometry == "tree" + assert mps.multi_site_events == tree.multi_site_events == 4 + assert mps.max_geometry >= 2 + assert tree.max_geometry >= 2 + assert mps.layout["kind"] == "mps_gate_stream_layout" + assert tree.layout.n == 5 + assert tree.layout_report["n_qubits"] == 5 + + +def test_planner_marks_unprepassable_stabilizer_streams_unavailable(): + """A dynamic-width stream should retain only safely priced candidates.""" + stream = [("h", 0), ("cap", (0,), "left")] + + advice = SimulatorPlanner(stream, n_qubits=2, chi=4).recommend() + + assert advice.candidate("MpsOptimizer").applicable + assert advice.candidate("TreeOptimizer").applicable + assert not advice.candidate("MpsStabOptimizer").applicable + assert not advice.candidate("TreeStabOptimizer").applicable + assert any("cap changes" in warning for warning in advice.warnings) + + +def test_planner_validates_pricing_inputs_and_candidate_lookup(): + """Planner-owned settings and typed candidate lookup should be explicit.""" + with pytest.raises(TypeError, match="positive integer"): + SimulatorPlanner([("t", 0)], chi=None) + with pytest.raises(ValueError, match="planner-owned"): + SimulatorPlanner( + [("t", 0)], + chi=4, + tree_layout_kwargs={"chi": 8}, + ) + with pytest.raises(ValueError, match="circuit width"): + SimulatorPlanner([], chi=4) + + advice = pepsy.recommend_simulator([], n_qubits=2, chi=4) + with pytest.raises(KeyError, match="Unknown simulator"): + advice.candidate("DenseSimulator") From d4c3b70bf4c2ac62411d4c1c5966cc79e2582820 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Wed, 29 Jul 2026 10:28:36 -0700 Subject: [PATCH 18/70] tree: add configurable cutoff mode --- CHANGELOG.md | 3 ++ docs/api/optimizers/tree.md | 13 ++++++++- src/pepsy/optimizers/tree/optimizer.py | 32 +++++++++++++++------ src/pepsy/optimizers/tree/ttn.py | 19 ++++++++++-- tests/test_optimize_tree.py | 40 ++++++++++++++++++++++++++ 5 files changed, 94 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 698ebc0..bbf0d04 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,9 @@ Changes for the next release should be added here before the version is bumped. - `SimulatorPlanner` and `recommend_simulator` provide non-executing, chi-aware rankings across MPS, tree, MPS-stabilizer, and tree-stabilizer circuit strategies using physical and dressed-frame support geometry. +- `TreeOptimizer` and `TreeTensorNetwork.compress_edge_` now accept + `cutoff_mode`, allowing Tree truncations to use the same Quimb + singular-value cutoff conventions as MPS truncations. ## [0.4.0] - 2026-07-27 diff --git a/docs/api/optimizers/tree.md b/docs/api/optimizers/tree.md index 2a2a3e3..293194d 100644 --- a/docs/api/optimizers/tree.md +++ b/docs/api/optimizers/tree.md @@ -28,7 +28,13 @@ plan = finder.run( search_budget=128, progbar=True, ) -opt = TreeOptimizer(gates, tree=plan, chi=64) +opt = TreeOptimizer( + gates, + tree=plan, + chi=64, + cutoff=1e-12, + cutoff_mode="rsum2", +) assert opt.plan.node_of_qubit[4] == opt.plan.root assert set(opt.tn.node_tensor(opt.plan.root).inds) >= {"k4"} ``` @@ -204,6 +210,11 @@ the optimizer's selected two-site mode for that run, later runs, and copies. The old `run(mode="tree")`/`"ttn"` selector is a deprecated no-op retained only for shared frontends. +`TreeOptimizer` accepts Quimb's `cutoff_mode` conventions for every truncating +Tree-edge SVD. Its default `"rel"` preserves historical Tree behavior; +`"rsum2"` applies a relative discarded-squared-weight threshold and matches +the default used by `MpsOptimizer`. + `TreeOptimizer.apply_submpo(...)` is the public form for an explicit MPO of arbitrary support. It losslessly QR-routes its virtual bonds, then uses its supplied (or configured) `max_bond` / `cutoff` in one final canonical sweep over diff --git a/src/pepsy/optimizers/tree/optimizer.py b/src/pepsy/optimizers/tree/optimizer.py index be57606..0afc324 100644 --- a/src/pepsy/optimizers/tree/optimizer.py +++ b/src/pepsy/optimizers/tree/optimizer.py @@ -294,7 +294,12 @@ class TreeOptimizer: ``None`` leaves the bond uncapped; the singular-value ``cutoff`` still applies. cutoff : float - Relative singular-value cutoff for truncations. + Singular-value cutoff for truncations, interpreted according to + ``cutoff_mode``. + cutoff_mode : str + Quimb singular-value cutoff mode. The default ``"rel"`` preserves the + historical TreeOptimizer behavior; use ``"rsum2"`` for a relative + discarded-squared-weight threshold matching Pepsy's MPS default. mode : {"auto", "direct", "mpo", "submpo"} Implementation used for two-site gates and explicit operator streams. ``"direct"`` uses the specialised gate-SVD/QR path. ``"mpo"`` first @@ -421,7 +426,7 @@ def _normalize_max_bond(max_bond): return max_bond def __init__(self, gates=None, n=None, *, chi=64, cutoff=1e-12, - mode="auto", two_site_mode=None, + cutoff_mode="rel", mode="auto", two_site_mode=None, structure="quality", max_arity=(2, 3, 4), community_frac=0.35, star_frac=0.75, layout_objective="path", layout_weight_mode="count", layout=None, tree=None, @@ -549,6 +554,7 @@ def __init__(self, gates=None, n=None, *, chi=64, cutoff=1e-12, self.cutoff = float(cutoff) if self.cutoff < 0.0: raise ValueError("cutoff must be non-negative.") + self.cutoff_mode = cutoff_mode self.mode = self._normalize_mode(mode) if two_site_mode is not None: legacy_mode = self._normalize_mode(two_site_mode) @@ -2427,12 +2433,14 @@ def _spectrum_payload(values, kept_values=None): return payload @staticmethod - def _probe_bond_spectrum(ta, tb, *, max_bond=None, cutoff=0.0): + def _probe_bond_spectrum( + ta, tb, *, max_bond=None, cutoff=0.0, cutoff_mode="rel", + ): """Return full/kept singular spectra for a two-tensor bond.""" info = {} qtn.tensor_compress_bond( ta.copy(), tb.copy(), max_bond=None, cutoff=0.0, - cutoff_mode="rel", absorb=None, info=info, + cutoff_mode=cutoff_mode, absorb=None, info=info, ) values = info.get("singular_values") if values is None: @@ -2444,7 +2452,7 @@ def _probe_bond_spectrum(ta, tb, *, max_bond=None, cutoff=0.0): kept_info = {} qtn.tensor_compress_bond( ta.copy(), tb.copy(), max_bond=max_bond, - cutoff=cutoff, cutoff_mode="rel", absorb=None, + cutoff=cutoff, cutoff_mode=cutoff_mode, absorb=None, info=kept_info, ) kept_values = kept_info.get("singular_values") @@ -2453,6 +2461,7 @@ def _probe_bond_spectrum(ta, tb, *, max_bond=None, cutoff=0.0): @staticmethod def _probe_split_spectrum( tensor, left_inds, *, max_bond=None, cutoff=0.0, + cutoff_mode="rel", ): """Return full/kept singular spectra for a tensor split.""" _, values, _ = tensor.split( @@ -2460,7 +2469,7 @@ def _probe_split_spectrum( method="svd", cutoff=0.0, max_bond=None, - cutoff_mode="rel", + cutoff_mode=cutoff_mode, absorb=None, get="arrays", ) @@ -2473,7 +2482,7 @@ def _probe_split_spectrum( method="svd", cutoff=cutoff, max_bond=max_bond, - cutoff_mode="rel", + cutoff_mode=cutoff_mode, absorb=None, get="arrays", ) @@ -2528,6 +2537,7 @@ def _record_truncation( # lossless split before the final ``chi``-limited path sweep. "max_bond": None if max_bond is None else int(max_bond), "cutoff": float(self.cutoff if cutoff is None else cutoff), + "cutoff_mode": self.cutoff_mode, }) def _split_with_diagnostics( @@ -2545,12 +2555,14 @@ def _split_with_diagnostics( left_inds, max_bond=max_bond, cutoff=cutoff, + cutoff_mode=self.cutoff_mode, ) if self.track_truncation and not lossless else None ) left, right = tensor.split( left_inds=left_inds, method="svd", max_bond=max_bond, - cutoff=cutoff, absorb="right", get="tensors", bond_ind=bond_ind, + cutoff=cutoff, cutoff_mode=self.cutoff_mode, + absorb="right", get="tensors", bond_ind=bond_ind, ) after_bond = self._tensor_ind_size(left, bond_ind) self._record_truncation( @@ -2596,13 +2608,14 @@ def _compress_edge_with_diagnostics( tb, max_bond=max_bond, cutoff=cutoff, + cutoff_mode=self.cutoff_mode, ) # Keep the live canonical-region metadata in one place: the TTN edge # wrapper performs the compression and advances its tracked centre. self.tn.compress_edge_( u, v, max_bond=max_bond, cutoff=cutoff, absorb="right", - reduced=reduced, + cutoff_mode=self.cutoff_mode, reduced=reduced, ) bond_after = self.tn.bond(u, v) after_bond = int(self.tn.ind_size(bond_after)) @@ -4165,6 +4178,7 @@ def copy(self): n=self.n, chi=self.chi, cutoff=self.cutoff, + cutoff_mode=self.cutoff_mode, mode=self.mode, structure=self.structure, max_arity=self.max_arity, diff --git a/src/pepsy/optimizers/tree/ttn.py b/src/pepsy/optimizers/tree/ttn.py index ee964a4..2cd8e92 100644 --- a/src/pepsy/optimizers/tree/ttn.py +++ b/src/pepsy/optimizers/tree/ttn.py @@ -1093,7 +1093,7 @@ def _fermionic_canonize_edge_(self, a, b, absorb): return self def _fermionic_compress_edge_( - self, a, b, *, max_bond, cutoff, absorb, + self, a, b, *, max_bond, cutoff, cutoff_mode, absorb, ): """Compress one native graded tree cut by an explicit two-node SVD.""" if absorb == "right": @@ -1113,6 +1113,7 @@ def _fermionic_compress_edge_( method="svd", max_bond=max_bond, cutoff=cutoff, + cutoff_mode=cutoff_mode, absorb="right", get="tensors", bond_ind=bond, @@ -1169,8 +1170,17 @@ def canonize_edge_(self, a, b, absorb="right"): self._track_edge_center(a, b, absorb, previous=previous) return self - def compress_edge_(self, a, b, *, max_bond=None, cutoff=1e-12, - absorb="right", reduced=True): + def compress_edge_( + self, + a, + b, + *, + max_bond=None, + cutoff=1e-12, + cutoff_mode="rel", + absorb="right", + reduced=True, + ): """Compress the tree edge ``a -> b`` in place. Dense/nonfermionic trees delegate to Quimb's ``compress_between``. @@ -1178,6 +1188,7 @@ def compress_edge_(self, a, b, *, max_bond=None, cutoff=1e-12, The tracked :attr:`orthogonality_center` advances as for :meth:`canonize_edge_`. + ``cutoff_mode`` selects Quimb's singular-value cutoff convention. ``reduced`` is forwarded only on the dense path. Quimb's one-sided ``"left"`` mode is exact when node ``b`` is already isometric on its non-shared legs. Native fermionic compression ignores this option and @@ -1190,6 +1201,7 @@ def compress_edge_(self, a, b, *, max_bond=None, cutoff=1e-12, b, max_bond=max_bond, cutoff=cutoff, + cutoff_mode=cutoff_mode, absorb=absorb, ) else: @@ -1198,6 +1210,7 @@ def compress_edge_(self, a, b, *, max_bond=None, cutoff=1e-12, self.node_tag(b), max_bond=max_bond, cutoff=cutoff, + cutoff_mode=cutoff_mode, absorb=absorb, reduced=reduced, ) diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index 43c80d4..4d78daf 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -144,6 +144,46 @@ def test_tree_two_site_direct_and_mpo_modes_agree(): assert _fidelity(mpo.to_dense(), exact) > 1 - 1e-9 +def test_tree_cutoff_mode_controls_edge_truncation_and_copy(): + """Tree truncations honor the configured Quimb cutoff convention.""" + small = 0.1 + large = np.sqrt(1.0 - small**2) + gate = np.array( + [ + [large, 0.0, 0.0, -small], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.0], + [small, 0.0, 0.0, large], + ], + dtype=complex, + ) + + relative = TreeOptimizer( + [(gate, (0, 1))], + n=2, + chi=2, + cutoff=0.05, + cutoff_mode="rel", + track_truncation=True, + ) + relative_sum2 = TreeOptimizer( + [(gate, (0, 1))], + n=2, + chi=2, + cutoff=0.05, + cutoff_mode="rsum2", + track_truncation=True, + ) + + assert relative.max_bond() == 2 + assert relative_sum2.max_bond() == 1 + assert all( + event["cutoff_mode"] == "rsum2" + for event in relative_sum2.truncation_history + ) + assert relative_sum2.copy().cutoff_mode == "rsum2" + + def test_dense_path_thread_preserves_qr_isometry_metadata(monkeypatch): """Every dense path-thread Q keeps its toward-destination isometry.""" rng = np.random.default_rng(919) From 21c479daa8c01bfedf44fe72468d813bca9abefb Mon Sep 17 00:00:00 2001 From: rezaquant Date: Wed, 29 Jul 2026 14:02:46 -0700 Subject: [PATCH 19/70] tree: preserve nonunitary normalization scale --- .github/skills/tree-optimizer/SKILL.md | 8 +- CHANGELOG.md | 6 + docs/api/optimizers/tree.md | 11 +- src/pepsy/optimizers/tree/optimizer.py | 158 ++++++++++++++++++------- tests/test_optimize_tree.py | 43 ++++++- 5 files changed, 179 insertions(+), 47 deletions(-) diff --git a/.github/skills/tree-optimizer/SKILL.md b/.github/skills/tree-optimizer/SKILL.md index f97e24d..f5df2a4 100644 --- a/.github/skills/tree-optimizer/SKILL.md +++ b/.github/skills/tree-optimizer/SKILL.md @@ -182,7 +182,13 @@ telescopes to identity between bra and ket. doubled tree. Native fermionic states use a one-tensor `TensorNetwork.H` contraction when a centre is known, so Symmray applies the graded outer-leg phase flips; unknown-centre fermionic states use the exact - complete doubled-network contraction. Keep the backend dispatch separate. + complete doubled-network contraction. Known-centre fast paths multiply the + raw centre norm by Quimb's extracted `10 ** tn.exponent`; full contractions + already apply it. During non-unitary replay, `normalize_every` / + `normalize_final` normalize only the raw working centre and accumulate its + removed scale in that exponent, preserving the represented state. Public + `normalize()` is physical renormalization and clears the exponent. Keep the + backend dispatch separate. - Any operation that moves/rebuilds the centre must update the tracked centre (via `self.center = ...`, i.e. `ttn.orthogonality_center`). diff --git a/CHANGELOG.md b/CHANGELOG.md index bbf0d04..e71c65d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,6 +21,12 @@ Changes for the next release should be added here before the version is bumped. `cutoff_mode`, allowing Tree truncations to use the same Quimb singular-value cutoff conventions as MPS truncations. +### Fixed + +- `TreeOptimizer` non-unitary scale control now preserves removed normalization + in the TTN exponent, and fast centre-based norm reads include that exponent, + so `normalize_every=True` no longer changes the represented state. + ## [0.4.0] - 2026-07-27 This release removes obsolete package-layout compatibility layers and keeps diff --git a/docs/api/optimizers/tree.md b/docs/api/optimizers/tree.md index 293194d..3bdd54d 100644 --- a/docs/api/optimizers/tree.md +++ b/docs/api/optimizers/tree.md @@ -696,8 +696,15 @@ For stream control events, `TreeOptimizer.measure_event`, `MpsOptimizer`, including Pauli-basis measurement and reset. Their recorded results are `(pauli, where, outcome, probability)` in `measurements`. `cap(q, vec)` contracts and removes one physical site, shifting the remaining labels -above `q` down by one unless stable labels are requested. `normalize()` rescales the represented state to unit -norm and `max_bond()` reports the largest virtual bond. Truncation details are +above `q` down by one unless stable labels are requested. +For a non-unitary run, `normalize_every=True` (or `normalize_final=True`) keeps +the canonical working tensor numerically normalized and accumulates each +removed base-10 scale in `tn.exponent`; `norm()`, `to_dense()`, copies, and +full contractions continue to represent the original physical scale. +The normalization records expose both the per-event raw scale and the +accumulated exponent. The public `normalize()` method remains a physical +renormalization: it clears that exponent and rescales the represented state to +unit norm. `max_bond()` reports the largest virtual bond. Truncation details are available through `truncation_report()`, `get_infidelities()`, and `get_infidelity_samples()` when spectrum tracking is enabled. diff --git a/src/pepsy/optimizers/tree/optimizer.py b/src/pepsy/optimizers/tree/optimizer.py index 0afc324..f0e8f3e 100644 --- a/src/pepsy/optimizers/tree/optimizer.py +++ b/src/pepsy/optimizers/tree/optimizer.py @@ -1165,6 +1165,10 @@ def _remount_product_state(self, state): root_tensor = target.node_tensor(target.plan.root) root_tensor.modify(data=root_tensor.data * factor) + # Quimb stores extracted global base-10 scale separately from tensor + # data. Preserve it when remounting a geometry-neutral product state. + if hasattr(state, "exponent"): + target.exponent = state.exponent target._with_center(self.plan.root) target._set_isometry_metadata_from_region({self.plan.root}).validate() self._state_backend_info(target) @@ -1663,11 +1667,15 @@ def run(self, gates=None, *, progbar=False, mode=None, non_unitary=False, stream; ``"tree"``/``"ttn"`` are deprecated no-op compatibility selectors for shared coefficient frontends. non_unitary : bool, default=False - Mark the stream as non-unitary when using automatic normalization. + Mark the stream as non-unitary when using automatic working-scale + control. normalize_every : bool, default=False - Normalize after every replay event. Requires ``non_unitary=True``. + Normalize the canonical working tensor after every replay event + and accumulate the removed global scale in ``tn.exponent``. + Requires ``non_unitary=True``. normalize_final : bool, default=False - Normalize once after replay. Requires ``non_unitary=True``. + Apply working-scale control once after replay. Requires + ``non_unitary=True``. normalize_eps : float, default=1e-15 Zero-state threshold used by automatic normalization. seed : int | None, default=None @@ -1779,17 +1787,12 @@ def run(self, gates=None, *, progbar=False, mode=None, non_unitary=False, self._finish_update() if normalize_every: - old_norm = self.norm() - self.normalize(eps=normalize_eps) - self.normalizations.append({ - "step": step, - "old_norm": float(old_norm * old_norm), - "span": tuple(support), - "insert": self.center, - "sites": tuple(support), - "scales": (float(old_norm),), - "reason": "step", - }) + self._normalize_and_record_working_scale( + step=step, + support=support, + reason="step", + eps=normalize_eps, + ) if pbar is not None: postfix = { @@ -1827,17 +1830,12 @@ def run(self, gates=None, *, progbar=False, mode=None, non_unitary=False, if pbar is not None: pbar.close() if normalize_final and self.G: - old_norm = self.norm() - self.normalize(eps=normalize_eps) - self.normalizations.append({ - "step": len(self.G), - "old_norm": float(old_norm * old_norm), - "span": (), - "insert": self.center, - "sites": (), - "scales": (float(old_norm),), - "reason": "final", - }) + self._normalize_and_record_working_scale( + step=len(self.G), + support=(), + reason="final", + eps=normalize_eps, + ) return self def set_gates(self, gates): @@ -3916,20 +3914,94 @@ def norm(self): a canonical center is known, with the complete doubled-network contraction as the unknown-gauge fallback. Dense/nonfermionic trees use the ordinary single-centre contraction when known and otherwise - the full doubled-tree path. + the full doubled-tree path. The fast one-tensor paths explicitly + restore Quimb's extracted base-10 ``tn.exponent``; full contractions + already include it. """ + center = self.center if self.tn.fermionic: with self._thread_ctx(): val = self.tn._fermionic_center_norm_squared() - return float(np.sqrt(abs(to_float(val, real=True)))) - if self.center is not None: - t = self.tn.tensor_map[self._tid(self.center)] + nrm = float(np.sqrt(abs(to_float(val, real=True)))) + if center is not None: + nrm *= self._represented_scale() + return nrm + if center is not None: + t = self.tn.tensor_map[self._tid(center)] val = qtn.tensor_contract(t.H, t, output_inds=[]) - return float(np.sqrt(abs(to_float(val, real=True)))) + nrm = float(np.sqrt(abs(to_float(val, real=True)))) + return nrm * self._represented_scale() with self._thread_ctx(): val = (self.tn.H & self.tn).contract(output_inds=[]) return float(np.sqrt(abs(to_float(val, real=True)))) + def _represented_scale(self): + """Return Quimb's extracted global base-10 state scale.""" + exponent = float(getattr(self.tn, "exponent", 0.0)) + try: + return float(10.0 ** exponent) + except OverflowError: + return np.inf + + def _working_norm(self): + """Return the raw canonical-centre norm without ``tn.exponent``. + + Non-unitary replay uses this to keep tensor entries numerically scaled + while preserving the removed global factor in Quimb's exponent. + Establishing a centre here is lossless and only needed for an + unmanaged/unknown canonical gauge. + """ + center = self.center + if center is None: + region = self.tn.canonical_region + center = min(region) if region else self.plan.root + self._move_center(center) + if self.tn.fermionic: + with self._thread_ctx(): + val = self.tn._fermionic_center_norm_squared(center) + else: + tensor = self.tn.tensor_map[self._tid(center)] + val = qtn.tensor_contract(tensor.H, tensor, output_inds=[]) + return float(np.sqrt(abs(to_float(val, real=True)))) + + def _normalize_working_scale(self, eps=1e-15): + """Normalize canonical working data and preserve represented scale.""" + working_norm = self._working_norm() + if working_norm > float(eps) and np.isfinite(working_norm): + self._invalidate_state_norm_cache() + tensor = self.tn.tensor_map[self._tid(self.center)] + tensor.modify(data=tensor.data / working_norm) + self.tn.exponent = ( + float(getattr(self.tn, "exponent", 0.0)) + + float(np.log10(working_norm)) + ) + return working_norm + + def _normalize_and_record_working_scale( + self, *, step, support, reason, eps, + ): + """Apply one non-unitary scale-control event and record its factor.""" + old_norm = self._normalize_working_scale(eps=eps) + log10_scale = ( + float(np.log10(old_norm)) if old_norm > 0.0 else -np.inf + ) + support = tuple(support) + event = { + "step": int(step), + "old_norm": float(old_norm * old_norm), + "span": support, + "insert": self.center, + "sites": support, + "scales": (float(old_norm),), + "log10_scale": log10_scale, + "log10_scales": (log10_scale,), + "reason": str(reason), + "method": "canonical_center", + "exponent": float(getattr(self.tn, "exponent", 0.0)), + } + self.normalizations.append(event) + return event + def _leaf_canonical_norm(self): """Return a cheap dense canonical norm or the exact fermionic norm. @@ -3951,23 +4023,22 @@ def normalize(self, eps=1e-15, insert=None): ``eps`` and ``insert`` are accepted for compatibility with :meth:`MpsOptimizer.normalize`; a tree has no chain insertion site, so - ``insert`` is intentionally ignored. + ``insert`` is intentionally ignored. Unlike the non-unitary replay + scale-control path, this is a physical renormalization: any accumulated + Quimb ``tn.exponent`` is cleared so the represented state has unit norm. """ eps = float(eps) if eps < 0.0: raise ValueError("eps must be non-negative.") nrm = self.norm() - if nrm > eps: + working_norm = self._working_norm() + if working_norm > eps and np.isfinite(working_norm): self._invalidate_state_norm_cache() - if self.center is not None: - target = self.center - elif self.tn.canonical_region is not None: - target = next(iter(self.tn.canonical_region)) - else: - target = self.plan.root - tid = self._tid(target) + tid = self._tid(self.center) t = self.tn.tensor_map[tid] - t.modify(data=t.data / nrm) + t.modify(data=t.data / working_norm) + if hasattr(self.tn, "exponent"): + self.tn.exponent = 0.0 return nrm def _measure_pauli(self, pauli, where, outcome=None, *, renormalize=True, @@ -4235,9 +4306,10 @@ def get_infidelity_samples(self): def get_normalizations(self): """Return automatic normalization records. - Tree normalization is explicit rather than an MPS run-time scale - controller, so this list remains empty unless a higher-level wrapper - records its own normalization events here. + Non-unitary replay records each raw canonical-centre scale removed + from tensor data and accumulated into ``tn.exponent``. Thus the + working tensors remain normalized without changing the represented + state norm. """ return self.normalizations diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index 4d78daf..90ff718 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -593,12 +593,51 @@ def test_tree_run_supports_shared_non_unitary_normalization_controls(): half = 0.5 * np.eye(2, dtype=complex) opt = TreeOptimizer([(half, 0), (half, 1)], n=3, run=False) opt.run(non_unitary=True, normalize_every=True) - assert opt.norm() == pytest.approx(1.0) + assert opt.norm() == pytest.approx(0.25) + assert np.linalg.norm(opt.to_dense()) == pytest.approx(0.25) + assert opt.tn.exponent == pytest.approx(np.log10(0.25)) assert len(opt.get_normalizations()) == 2 + assert [event["old_norm"] for event in opt.get_normalizations()] == pytest.approx( + [0.25, 0.25] + ) with pytest.raises(ValueError, match="non_unitary"): opt.run(normalize_every=True) +def test_tree_nonunitary_scale_control_preserves_represented_state(): + """Per-step scale control changes only the TTN working-data gauge.""" + twice = 2.0 * np.eye(2, dtype=complex) + gates = [(twice, 0), (twice, 1)] + raw = TreeOptimizer(gates, n=3) + controlled = TreeOptimizer(gates, n=3, run=False) + + controlled.run(non_unitary=True, normalize_every=True) + + assert np.allclose(controlled.to_dense(), raw.to_dense()) + assert controlled.norm() == pytest.approx(raw.norm()) + assert controlled.tn.exponent == pytest.approx(np.log10(4.0)) + center = controlled.tn.node_tensor(controlled.center) + assert np.linalg.norm(np.asarray(center.data)) == pytest.approx(1.0) + assert [event["exponent"] for event in controlled.get_normalizations()] == ( + pytest.approx([np.log10(2.0), np.log10(4.0)]) + ) + + +def test_tree_physical_normalize_clears_accumulated_scale(): + """Public normalize still makes the represented state unit norm.""" + twice = 2.0 * np.eye(2, dtype=complex) + opt = TreeOptimizer([(twice, 0)], n=2, run=False) + opt.run(non_unitary=True, normalize_every=True) + assert opt.norm() == pytest.approx(2.0) + + old_norm = opt.normalize() + + assert old_norm == pytest.approx(2.0) + assert opt.tn.exponent == pytest.approx(0.0) + assert opt.norm() == pytest.approx(1.0) + assert np.linalg.norm(opt.to_dense()) == pytest.approx(1.0) + + def test_tree_logical_position_helpers_are_identity_mps_compatibility(): """Tree backends expose identity logical/physical mapping helpers.""" opt = TreeOptimizer(None, n=4, run=False) @@ -1296,12 +1335,14 @@ def test_product_ttn_is_remounted_exactly_on_a_requested_new_layout(): target_plan = TreePlan.from_order((0, 2, 1, 3), structure="balanced") h = np.array([[1.0, 1.0], [1.0, -1.0]], dtype=complex) / np.sqrt(2.0) source = TreeOptimizer([(h, 1)], tree=source_plan, chi=8).tn.copy() + source.exponent = np.log10(3.0) with pytest.warns(UserWarning, match="product TreeTensorNetwork"): opt = TreeOptimizer(None, state=source, tree=target_plan, run=False) assert opt.plan is target_plan assert opt.max_bond() == 1 + assert opt.tn.exponent == pytest.approx(source.exponent) assert np.allclose( np.asarray(opt.to_dense()).reshape(-1), np.asarray(source.to_dense()).reshape(-1), From f2aae038154ed05f07962762193dada802aa59d0 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Wed, 29 Jul 2026 15:29:07 -0700 Subject: [PATCH 20/70] tree: share native MPO routing across optimizers --- docs/api/optimizers/tree.md | 8 +-- docs/api/optimizers/tree_stabilizer.md | 8 +++ src/pepsy/optimizers/tree/optimizer.py | 20 ++++-- src/pepsy/optimizers/tree/ttn.py | 2 +- .../optimizers/tree_stabilizer/optimizer.py | 12 +++- tests/test_optimize_tree.py | 41 ++++++++++++ tests/test_optimize_tree_stabilizer.py | 62 +++++++++++++++++++ 7 files changed, 140 insertions(+), 13 deletions(-) diff --git a/docs/api/optimizers/tree.md b/docs/api/optimizers/tree.md index 3bdd54d..532a9ab 100644 --- a/docs/api/optimizers/tree.md +++ b/docs/api/optimizers/tree.md @@ -32,7 +32,7 @@ opt = TreeOptimizer( gates, tree=plan, chi=64, - cutoff=1e-12, + cutoff=1e-10, cutoff_mode="rsum2", ) assert opt.plan.node_of_qubit[4] == opt.plan.root @@ -211,9 +211,9 @@ The old `run(mode="tree")`/`"ttn"` selector is a deprecated no-op retained only for shared frontends. `TreeOptimizer` accepts Quimb's `cutoff_mode` conventions for every truncating -Tree-edge SVD. Its default `"rel"` preserves historical Tree behavior; -`"rsum2"` applies a relative discarded-squared-weight threshold and matches -the default used by `MpsOptimizer`. +Tree-edge SVD. Its defaults, `cutoff=1e-10` and `cutoff_mode="rsum2"`, match +Quimb's open-boundary `MatrixProductState.gate_with_submpo` compression path. +`"rel"` remains available as a relative largest-singular-value threshold. `TreeOptimizer.apply_submpo(...)` is the public form for an explicit MPO of arbitrary support. It losslessly QR-routes its virtual bonds, then uses its diff --git a/docs/api/optimizers/tree_stabilizer.md b/docs/api/optimizers/tree_stabilizer.md index c0cd5fb..5eccca4 100644 --- a/docs/api/optimizers/tree_stabilizer.md +++ b/docs/api/optimizers/tree_stabilizer.md @@ -10,6 +10,14 @@ milestone. It represents the state as where `C` is a Stim tableau Clifford and `|p>` is a dense two-level `TreeTensorNetwork` evolved by `TreeOptimizer`. +TreeStab forwards `cutoff` and `cutoff_mode` to that same coefficient +optimizer. The defaults are `cutoff=1e-10` and `cutoff_mode="rsum2"`, matching +Quimb's open-boundary `gate_with_submpo` compression convention. Its +`mode="mpo"` path and explicit coefficient-frame `submpo` events therefore +reuse TreeOptimizer's Quimb MPO tag lookup, lossless QR routing, and one final +subtree compression sweep; a payload without that MPO interface is the only +case that uses the bounded dense fallback. + Canonical and compression state has the same single owner as ordinary tree simulation: local isometry proofs live on the coefficient tensors' ``left_inds`` and are interpreted by ``TreeTensorNetwork``. TreeStab delegates diff --git a/src/pepsy/optimizers/tree/optimizer.py b/src/pepsy/optimizers/tree/optimizer.py index f0e8f3e..9eb1718 100644 --- a/src/pepsy/optimizers/tree/optimizer.py +++ b/src/pepsy/optimizers/tree/optimizer.py @@ -232,6 +232,8 @@ def _operator_schmidt_rank(op, where, left_where): "Z": np.array([[1.0, 0.0], [0.0, -1.0]], dtype=complex), } _RESET_FLIP_AXES = {"X": "Z", "Y": "X", "Z": "X"} +_DEFAULT_CUTOFF = 1e-10 +_DEFAULT_CUTOFF_MODE = "rsum2" _DEFAULT_MAX_OPERATOR_QUBITS = 8 _DEFAULT_MAX_SUBTREE_NODES = 128 @@ -297,9 +299,9 @@ class TreeOptimizer: Singular-value cutoff for truncations, interpreted according to ``cutoff_mode``. cutoff_mode : str - Quimb singular-value cutoff mode. The default ``"rel"`` preserves the - historical TreeOptimizer behavior; use ``"rsum2"`` for a relative - discarded-squared-weight threshold matching Pepsy's MPS default. + Quimb singular-value cutoff mode. The default ``"rsum2"`` matches + Quimb's open-boundary ``MatrixProductState.gate_with_submpo`` path; + use ``"rel"`` for a relative largest-singular-value threshold. mode : {"auto", "direct", "mpo", "submpo"} Implementation used for two-site gates and explicit operator streams. ``"direct"`` uses the specialised gate-SVD/QR path. ``"mpo"`` first @@ -425,8 +427,10 @@ def _normalize_max_bond(max_bond): raise ValueError("max_bond must be a positive integer or None.") return max_bond - def __init__(self, gates=None, n=None, *, chi=64, cutoff=1e-12, - cutoff_mode="rel", mode="auto", two_site_mode=None, + def __init__(self, gates=None, n=None, *, chi=64, + cutoff=_DEFAULT_CUTOFF, + cutoff_mode=_DEFAULT_CUTOFF_MODE, mode="auto", + two_site_mode=None, structure="quality", max_arity=(2, 3, 4), community_frac=0.35, star_frac=0.75, layout_objective="path", layout_weight_mode="count", layout=None, tree=None, @@ -3142,7 +3146,11 @@ def _try_apply_native_submpo( state_t = self.tn.tensor_map[self._tid(nid)].copy() state_inds[nid] = set(state_t.inds) q = self.plan.qubit_of_node.get(nid) - if q is None: + # A physical root can be an internal Steiner node without being + # acted on by the MPO. Keep its state tensor untouched and route + # it as ordinary state data. Only target physical nodes need an + # MPO site tensor and the associated site-tag lookup. + if q is None or q not in payload_for_compact: local[nid] = state_t operator_inds[nid] = set() continue diff --git a/src/pepsy/optimizers/tree/ttn.py b/src/pepsy/optimizers/tree/ttn.py index 2cd8e92..4ddce86 100644 --- a/src/pepsy/optimizers/tree/ttn.py +++ b/src/pepsy/optimizers/tree/ttn.py @@ -1176,7 +1176,7 @@ def compress_edge_( b, *, max_bond=None, - cutoff=1e-12, + cutoff=1e-10, cutoff_mode="rel", absorb="right", reduced=True, diff --git a/src/pepsy/optimizers/tree_stabilizer/optimizer.py b/src/pepsy/optimizers/tree_stabilizer/optimizer.py index f56c2d9..5e7306a 100644 --- a/src/pepsy/optimizers/tree_stabilizer/optimizer.py +++ b/src/pepsy/optimizers/tree_stabilizer/optimizer.py @@ -39,7 +39,11 @@ from ..stabilizer_tn.stn_state import _CLIFFORD_GATES, _validate_bits from ..mps.optimizer import conditional_event_parts, submpo_event_parts from ..tree.layout import TreeLayoutFinder, TreePlan -from ..tree.optimizer import TreeOptimizer +from ..tree.optimizer import ( + TreeOptimizer, + _DEFAULT_CUTOFF, + _DEFAULT_CUTOFF_MODE, +) from ..tree.ttn import TreeTensorNetwork __all__ = ["TreeStabOptimizer", "run_stabilizer_tree_stream"] @@ -620,7 +624,8 @@ def __init__( *, n=None, chi=None, - cutoff=1e-12, + cutoff=_DEFAULT_CUTOFF, + cutoff_mode=_DEFAULT_CUTOFF_MODE, tree=None, layout=None, structure="quality", @@ -818,6 +823,7 @@ def __init__( n=n, chi=chi, cutoff=cutoff, + cutoff_mode=cutoff_mode, mode=mode, structure=structure, max_arity=max_arity, @@ -1455,6 +1461,7 @@ def apply_frame_layout(self, plan="auto", *, layout_kwargs=None): n=self.n, chi=self._tree.chi, cutoff=self._tree.cutoff, + cutoff_mode=self._tree.cutoff_mode, mode=self._tree.mode, structure=self._tree.structure, max_arity=self._tree.max_arity, @@ -3512,6 +3519,7 @@ def cap(self, where, vec, *, absorb="left") -> "TreeStabOptimizer": n=reduced_n, chi=old_tree.chi, cutoff=old_tree.cutoff, + cutoff_mode=old_tree.cutoff_mode, mode=old_tree.mode, structure=old_tree.structure, max_arity=old_tree.max_arity, diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index 90ff718..22be4c9 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -184,6 +184,16 @@ def test_tree_cutoff_mode_controls_edge_truncation_and_copy(): assert relative_sum2.copy().cutoff_mode == "rsum2" +def test_tree_cutoff_defaults_match_quimb_open_mps_submpo_path(): + """TreeOptimizer defaults match Quimb's open-chain MPO compression.""" + opt = TreeOptimizer(None, n=2, run=False) + + assert opt.cutoff == pytest.approx(1e-10) + assert opt.cutoff_mode == "rsum2" + assert opt.copy().cutoff == pytest.approx(1e-10) + assert opt.copy().cutoff_mode == "rsum2" + + def test_dense_path_thread_preserves_qr_isometry_metadata(monkeypatch): """Every dense path-thread Q keeps its toward-destination isometry.""" rng = np.random.default_rng(919) @@ -2342,6 +2352,37 @@ def fail_to_dense(): assert report["events"][0]["crossing_edges"] +def test_tree_native_submpo_keeps_unacted_physical_root_structured(monkeypatch): + """A physical root on the Steiner subtree need not be an MPO site.""" + where = (1, 2, 3) + plan = TreePlan.from_order( + range(1, 5), structure="balanced", root_qubit=0 + ) + gate = _rand_unitary(len(where), np.random.default_rng(53)) + mpo = qtn.MatrixProductOperator.from_dense( + gate.reshape((2,) * (2 * len(where))), + dims=(2,) * len(where), + sites=where, + L=5, + max_bond=None, + cutoff=0.0, + ) + expected = _sv_apply_kq( + np.eye(2**5, dtype=complex)[:, 0], gate, where, 5 + ) + + def fail_to_dense(): + raise AssertionError("sub-MPO was unexpectedly materialized") + + monkeypatch.setattr(mpo, "to_dense", fail_to_dense) + opt = TreeOptimizer( + None, n=5, tree=plan, chi=64, cutoff=0.0, run=False + ) + opt.apply_submpo(mpo, where) + + assert _fidelity(expected, opt.to_dense()) > 1 - 1e-10 + + def test_tree_submpo_mode_declares_and_validates_mpo_streams(): """The explicit sub-MPO mode accepts MPO events and rejects dense gates.""" mpo = _two_branch_flip_submpo(L=4, sites=(0, 3), targets=(0, 3)) diff --git a/tests/test_optimize_tree_stabilizer.py b/tests/test_optimize_tree_stabilizer.py index a05e26b..050597a 100644 --- a/tests/test_optimize_tree_stabilizer.py +++ b/tests/test_optimize_tree_stabilizer.py @@ -85,6 +85,14 @@ def test_tree_stab_is_public_and_cliffords_are_tableau_only(): _assert_same_state(opt.to_statevector(), expected) +def test_tree_stab_cutoff_defaults_match_tree_quimb_path(): + """TreeStab forwards the TreeOptimizer/Quimb compression defaults.""" + opt = pepsy.TreeStabOptimizer(2) + + assert opt.tree_optimizer.cutoff == pytest.approx(1e-10) + assert opt.tree_optimizer.cutoff_mode == "rsum2" + + def test_tree_stab_isometry_api_and_backend_conversion_preserve_proofs(): """TreeStab delegates one live map and backend conversion keeps it valid.""" def converter(array): @@ -280,6 +288,60 @@ def test_tree_stab_submpo_matches_mps_coefficient_frame_contract(): _assert_same_state(tree.to_statevector(), mps.to_statevector()) +def test_tree_stab_submpo_uses_native_tree_router_for_unacted_root(monkeypatch): + """TreeStab keeps structured MPO application on the coefficient tree.""" + from pepsy.optimizers.tree import TreePlan + + where = (1, 2, 3) + plan = TreePlan.from_order( + range(1, 5), structure="balanced", root_qubit=0 + ) + dense_operator = np.diag( + np.array([1.0, 0.8, 0.6, 0.4, 0.3, 0.2, 0.1, -0.2]) + ).astype(complex) + submpo = qtn.MatrixProductOperator.from_dense( + dense_operator.reshape((2,) * 6), + dims=(2, 2, 2), + sites=where, + L=5, + max_bond=None, + cutoff=0.0, + ) + expected = _apply_local( + np.eye(2**5, dtype=complex)[:, 0], dense_operator, where, 5 + ) + + def fail_to_dense(): + raise AssertionError("TreeStab sub-MPO was unexpectedly materialized") + + monkeypatch.setattr(submpo, "to_dense", fail_to_dense) + opt = pepsy.TreeStabOptimizer(5, tree=plan, chi=64, cutoff=0.0) + opt.apply([("submpo", submpo, where)]) + + _assert_same_state(opt.to_statevector(), expected) + assert opt.tree_optimizer.update_history[0]["kind"] == "submpo" + + +def test_tree_stab_mpo_mode_reuses_tree_two_factor_kernel(monkeypatch): + """TreeStab's MPO mode uses the same TreeOptimizer factor kernel.""" + opt = pepsy.TreeStabOptimizer(2, mode="mpo") + calls = [] + apply_factors = opt.tree_optimizer._apply_2q_factors_impl + + def traced_apply_factors(*args, **kwargs): + calls.append(True) + return apply_factors(*args, **kwargs) + + monkeypatch.setattr( + opt.tree_optimizer, "_apply_2q_factors_impl", traced_apply_factors + ) + opt.apply([("cnot", 0, 1)]) + opt.measure_pauli("Z", 1, outcome=+1, absorb_basis=True) + + assert calls + assert opt.tree_optimizer.mode == "mpo" + + def test_tree_stab_amplitude_probability_match_dense_readout(): opt = pepsy.TreeStabOptimizer( 2, gates=[("h", 0), ("cnot", 0, 1)] From 4a251aac89b30ea4d886cd0ec9c54f10728b28e8 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Wed, 29 Jul 2026 16:52:05 -0700 Subject: [PATCH 21/70] tree: preserve capped MPO replay components --- docs/api/optimizers/tree.md | 15 ++-- history/2026-07-29-tree-submpo-cutoff.md | 33 +++++++++ src/pepsy/optimizers/tree/optimizer.py | 88 +++++++++++++++++++++--- tests/test_optimize_tree.py | 53 ++++++++++++++ 4 files changed, 174 insertions(+), 15 deletions(-) create mode 100644 history/2026-07-29-tree-submpo-cutoff.md diff --git a/docs/api/optimizers/tree.md b/docs/api/optimizers/tree.md index 532a9ab..9debf35 100644 --- a/docs/api/optimizers/tree.md +++ b/docs/api/optimizers/tree.md @@ -165,8 +165,12 @@ for the whole Steiner subtree is formed. Each dense routed Q tensor retains its `left_inds` isometry metadata, so canonical recovery recognizes that it already points toward the hub instead of repeating the same QR; native fermionic trees retain explicit graded QR recovery. Once all MPO factors have arrived, every -touched edge is SVD-compressed once. Thus every truncation sees the complete -operator in an isometric environment. +touched edge is compressed once. A bond that remains within its configured +`max_bond` uses a lossless QR, avoiding repeated cutoff loss of tiny state +components across successive sub-MPO events; when the MPO expands an edge +past its cap, the configured Quimb `cutoff` and `cutoff_mode` are applied to +the truncating SVD. Thus every actual truncation sees the complete operator in +an isometric environment. `op` acts on `len(where)` qubits: an array reshaped to `(2,) * 2k` with output indices first, `op[o_0..o_{k-1}, i_0..i_{k-1}]` (a `(2**k, 2**k)` matrix is @@ -203,7 +207,9 @@ fermionic gates. Select `mode="mpo"` explicitly to inspect or benchmark Quimb's operator-TN factorization. Direct and MPO share the update kernel and defer truncation until the complete gate has reached the affected path, so at an exact `chi` they differ only by the factorization gauge and numerical -roundoff. +roundoff. MPO replay also uses the cap-aware cutoff rule: already-within-cap +path bonds take lossless QR, while expanded bonds use the configured cutoff +mode. `run(mode=...)` has the same persistent semantics as `MpsOptimizer`: it updates the optimizer's selected two-site mode for that run, later runs, and copies. @@ -218,7 +224,8 @@ Quimb's open-boundary `MatrixProductState.gate_with_submpo` compression path. `TreeOptimizer.apply_submpo(...)` is the public form for an explicit MPO of arbitrary support. It losslessly QR-routes its virtual bonds, then uses its supplied (or configured) `max_bond` / `cutoff` in one final canonical sweep over -the affected subtree. +the affected subtree. Existing bonds at or below `max_bond` take a lossless +QR; only bonds expanded past the cap invoke the configured cutoff mode. The tree backend also exposes numerical Pauli primitives used by a future stabilizer frontend: `apply_pauli_rotation(...)`, `apply_pauli_sum(...)`, `expectation_pauli(...)`, `measure_pauli(...)`, and `project_pauli(...)`. These diff --git a/history/2026-07-29-tree-submpo-cutoff.md b/history/2026-07-29-tree-submpo-cutoff.md new file mode 100644 index 0000000..9ee1ba3 --- /dev/null +++ b/history/2026-07-29-tree-submpo-cutoff.md @@ -0,0 +1,33 @@ +# 2026-07-29 — cap-aware Tree MPO cutoffs + +- Milestone: Tree/MPO replay accuracy and shared TreeStab routing +- Branch / commit: `develop` + +## What changed + +- Tree subtree and MPO path replays now use lossless QR on bonds already at or + below the active bond cap, while configured Quimb cutoff modes apply when a + route expands a bond past that cap. +- Added a deterministic regression for tiny pre-existing state components and + updated the Tree optimizer documentation. + +## Why + +- Distance-5, two-cycle SurfaceCode replay with `chi=64`, + `cutoff=1e-12`, `cutoff_mode="rsum2"` changed from roughly 21.9% Tree + logical errors to 0.098% while keeping the native sub-MPO route and the same + layout. + +## How it was validated + +- Focused Tree/MPO/sub-MPO tests: passed. +- `tests/test_optimize_tree_stabilizer.py`: 57 passed. +- Notebook-derived distance-5, two-cycle native replay: 0.0009756 error rate, + max bond 64. +- `py_compile`: passed. +- Ruff was unavailable in `/Users/rezah/envs/genpy`. + +## Open questions / blockers + +- Two broader Tree tests hit sandbox semaphore-permission errors while cotengra + attempted to create loky workers; they were unrelated to this change. diff --git a/src/pepsy/optimizers/tree/optimizer.py b/src/pepsy/optimizers/tree/optimizer.py index 9eb1718..99cbb41 100644 --- a/src/pepsy/optimizers/tree/optimizer.py +++ b/src/pepsy/optimizers/tree/optimizer.py @@ -2096,7 +2096,12 @@ def _apply_2q_mpo_impl( "with one local factor per requested site." ) return self._apply_2q_factors_impl( - *factors, qa, qb, max_bond=max_bond, cutoff=cutoff + *factors, + qa, + qb, + max_bond=max_bond, + cutoff=cutoff, + preserve_subcap=True, ) def _two_site_mpo_factors(self, submpo, qa, qb, *, site_where=None): @@ -2198,7 +2203,7 @@ def _apply_2q_path_thread_impl( def _apply_2q_factors_impl( self, factors, outputs, thread_ind, qa, qb, *, max_bond=None, - cutoff=None, + cutoff=None, preserve_subcap=False, ): """Apply two local operator factors with one shared path-thread kernel. @@ -2220,7 +2225,9 @@ def _apply_2q_factors_impl( ): return self._apply_2q_sibling_factors( factors, outputs, qa, qb, la, lb, parent, - max_bond=max_bond, cutoff=cutoff, + max_bond=max_bond, + cutoff=cutoff, + preserve_subcap=preserve_subcap, ) source, destination = ( @@ -2253,14 +2260,19 @@ def _apply_2q_factors_impl( data=merged_destination.data, inds=merged_destination.inds, ) self.center = destination_node - self._compress_path(path, max_bond=max_bond, cutoff=cutoff) + self._compress_path( + path, + max_bond=max_bond, + cutoff=cutoff, + preserve_subcap=preserve_subcap, + ) finally: self._thread_ind = None return self def _apply_2q_sibling_factors( self, factors, outputs, qa, qb, la, lb, parent, *, max_bond=None, - cutoff=None, + cutoff=None, preserve_subcap=False, ): """Apply two local factors to sibling leaves through their parent. @@ -2292,11 +2304,13 @@ def _apply_2q_sibling_factors( blob, [pa], edge=(la, parent), bond_ind=e_la, max_bond=self.chi if max_bond is None else max_bond, cutoff=self.cutoff if cutoff is None else cutoff, + preserve_subcap=preserve_subcap, ) lb_t, p_t = self._split_with_diagnostics( rem, [pb], edge=(lb, parent), bond_ind=e_lb, max_bond=self.chi if max_bond is None else max_bond, cutoff=self.cutoff if cutoff is None else cutoff, + preserve_subcap=preserve_subcap, ) tla.modify( data=la_t.data, @@ -2544,9 +2558,14 @@ def _record_truncation( def _split_with_diagnostics( self, tensor, left_inds, *, edge, bond_ind, max_bond, cutoff, + preserve_subcap=False, ): """Split ``tensor`` and record the resulting virtual edge.""" before_bond = self._split_rank_bound(tensor, left_inds) + if preserve_subcap: + cutoff = self._subtree_cutoff_for_size( + before_bond, max_bond=max_bond, cutoff=cutoff, + ) # An uncapped zero-cutoff routing split is provably lossless. Avoid an # otherwise redundant full SVD spectrum probe when diagnostics are on; # the subsequent final compression remains fully tracked. @@ -2576,6 +2595,21 @@ def _split_with_diagnostics( ) return left, right + def _subtree_cutoff_for_size(self, before_bond, *, max_bond, cutoff): + """Return the cutoff for a Tree/MPO update at ``before_bond`` size.""" + requested_cutoff = self.cutoff if cutoff is None else float(cutoff) + effective_max_bond = ( + self.chi + if max_bond is None + else self._normalize_max_bond(max_bond) + ) + if ( + effective_max_bond is not None + and int(before_bond) <= effective_max_bond + ): + return 0.0 + return requested_cutoff + def _compress_edge_with_diagnostics( self, u, v, *, max_bond=None, cutoff=None, reduced=True, ): @@ -2640,7 +2674,9 @@ def _metadata_aware_reduction(self, u, v): return "left" return True - def _compress_path(self, path, *, max_bond=None, cutoff=None): + def _compress_path( + self, path, *, max_bond=None, cutoff=None, preserve_subcap=False, + ): """Canonically compress every bond along ``path`` down to ``chi``. The orthogonality centre sits at ``path[-1]`` on entry; sweeping back to @@ -2650,8 +2686,15 @@ def _compress_path(self, path, *, max_bond=None, cutoff=None): sweep of Seitz et al. (Fig. 6) applied along the gate geodesic. """ for v, u in zip(path[::-1], path[-2::-1]): + edge_cutoff = cutoff + if preserve_subcap: + edge_cutoff = self._subtree_cutoff_for_size( + self.tn.ind_size(self.tn.bond(v, u)), + max_bond=max_bond, + cutoff=cutoff, + ) self._compress_edge_with_diagnostics( - v, u, max_bond=max_bond, cutoff=cutoff, + v, u, max_bond=max_bond, cutoff=edge_cutoff, reduced=self._metadata_aware_reduction(v, u), ) self.center = path[0] @@ -2661,12 +2704,30 @@ def _compress_subtree(self, snodes, hub, *, max_bond=None, cutoff=None): Starting at ``hub``, descend each branch. Compressing ``node -> child`` moves the centre onto the child; a lossless QR move returns it before - the next branch. Thus every SVD sees the completed operator update with - an isometric environment, while every edge is truncated exactly once. + the next branch. Thus every actual SVD sees the completed operator + update with an isometric environment, while every affected edge is + compressed exactly once. """ snodes = frozenset(snodes) self._move_center(hub) + def edge_cutoff(node, child): + """Keep existing sub-cap bonds lossless during subtree replay. + + The routed subtree already contains the complete state and MPO + update. Reapplying a positive cutoff to a bond that is still + within the active bond cap can repeatedly remove tiny + *pre-existing* state components on every MPO event. Keep this + Tree/MPO route stable by using the configured Quimb cutoff mode on + over-cap bonds, and using a lossless QR on bonds that remain within + the cap. + """ + return self._subtree_cutoff_for_size( + self.tn.ind_size(self.tn.bond(node, child)), + max_bond=max_bond, + cutoff=cutoff, + ) + def descend(node, parent): children = sorted( neighbor @@ -2675,7 +2736,10 @@ def descend(node, parent): ) for child in children: self._compress_edge_with_diagnostics( - node, child, max_bond=max_bond, cutoff=cutoff, + node, + child, + max_bond=max_bond, + cutoff=edge_cutoff(node, child), reduced=self._metadata_aware_reduction(node, child), ) descend(child, node) @@ -2819,7 +2883,9 @@ def _apply_submpo_resolved(self, submpo, where, *, max_bond=None, if factors is not None: self._apply_2q_factors_impl( *factors, where[0], where[1], - max_bond=max_bond, cutoff=cutoff, + max_bond=max_bond, + cutoff=cutoff, + preserve_subcap=True, ) applied = True if applied is None: diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index 22be4c9..095e934 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -393,6 +393,59 @@ def test_tree_multisite_submpo_qr_routes_before_one_subtree_sweep(): assert all(event["max_bond"] == 1 for event in opt.truncation_history) +def test_tree_submpo_does_not_retruncate_existing_within_cap_bonds(): + """A native MPO replay keeps tiny pre-existing state components.""" + eps = 1e-8 + large = np.sqrt(1.0 - eps**2) + operator = np.zeros((4, 4), dtype=complex) + operator[:, 0] = (large, 0.0, 0.0, eps) + plan = TreePlan.from_order(range(3), structure="balanced") + + seed = TreeOptimizer( + None, n=3, tree=plan, chi=4, cutoff=0.0, mode="direct", run=False, + ) + seed.apply_2q(operator, 0, 1) + expected = seed.to_dense() + + identity = qtn.MatrixProductOperator.from_dense( + np.eye(8, dtype=complex).reshape((2,) * 6), + dims=(2, 2, 2), + sites=(0, 1, 2), + L=3, + max_bond=None, + cutoff=0.0, + ) + replay = TreeOptimizer( + None, + n=3, + tree=plan, + state=seed.tn, + chi=4, + cutoff=1e-12, + cutoff_mode="rsum2", + mode="submpo", + run=False, + ) + replay.apply_submpo(identity, (0, 1, 2)) + + np.testing.assert_allclose(expected, replay.to_dense(), atol=1e-14) + assert replay.to_dense()[6] == pytest.approx(eps) + + mpo_replay = TreeOptimizer( + None, + n=3, + tree=plan, + state=seed.tn, + chi=4, + cutoff=1e-12, + cutoff_mode="rsum2", + mode="mpo", + run=False, + ) + mpo_replay.apply_2q(np.eye(4, dtype=complex), 0, 1) + np.testing.assert_allclose(expected, mpo_replay.to_dense(), atol=1e-14) + + def test_dense_subtree_hub_recovery_reuses_routed_q_metadata(monkeypatch): """Dense routed Q tensors recover the hub without another numerical QR.""" import quimb.tensor.tensor_core as qtc From 906d7f448418be91a44c903698ba288fc0ea156f Mon Sep 17 00:00:00 2001 From: rezaquant Date: Wed, 29 Jul 2026 18:24:15 -0600 Subject: [PATCH 22/70] qMERA graded compiled backend --- .../skills/qmera-energy-optimizer/SKILL.md | 164 ++- .../qmera-energy-optimizer/agents/openai.yaml | 2 +- .../references/design.md | 172 ++-- docs/api/index.md | 2 +- docs/api/optimizers/{mera.md => qmera.md} | 144 ++- docs/api/package.md | 2 +- docs/development/modules/optimizers.md | 4 +- examples/qmera_fermion_hubbard_2d.py | 2 +- examples/qmera_fermion_hubbard_4x4_pbc.py | 2 +- examples/qmera_majorana_2d.py | 2 +- examples/qmera_scale_plan_6x6.py | 2 +- src/pepsy/experimental/__init__.py | 3 +- src/pepsy/optimizers/__init__.py | 25 +- src/pepsy/optimizers/mera/__init__.py | 159 +-- src/pepsy/optimizers/mera/optimizer.py | 468 --------- src/pepsy/optimizers/qmera/__init__.py | 132 +++ .../optimizers/{mera => qmera}/builders.py | 51 +- src/pepsy/optimizers/{mera => qmera}/cache.py | 2 +- .../optimizers/{mera => qmera}/compiled.py | 254 ++++- .../optimizers/{mera => qmera}/fermions.py | 0 src/pepsy/optimizers/{mera => qmera}/gates.py | 0 .../optimizers/{mera => qmera}/geometry.py | 0 src/pepsy/optimizers/qmera/layout.py | 451 ++++++++ .../optimizers/{mera => qmera}/lightcones.py | 10 +- .../optimizers/{mera => qmera}/parametric.py | 23 +- src/pepsy/optimizers/qmera/prototype.py | 153 +++ .../optimizers/{mera => qmera}/schedules.py | 112 +- .../optimizers/{mera => qmera}/schematics.py | 95 +- src/pepsy/optimizers/{mera => qmera}/terms.py | 2 +- tests/conftest.py | 2 +- ...ptimize_mera.py => test_optimize_qmera.py} | 963 ++++++++++++++---- tests/test_package_layout.py | 10 +- 32 files changed, 2333 insertions(+), 1080 deletions(-) rename docs/api/optimizers/{mera.md => qmera.md} (77%) delete mode 100644 src/pepsy/optimizers/mera/optimizer.py create mode 100644 src/pepsy/optimizers/qmera/__init__.py rename src/pepsy/optimizers/{mera => qmera}/builders.py (95%) rename src/pepsy/optimizers/{mera => qmera}/cache.py (97%) rename src/pepsy/optimizers/{mera => qmera}/compiled.py (52%) rename src/pepsy/optimizers/{mera => qmera}/fermions.py (100%) rename src/pepsy/optimizers/{mera => qmera}/gates.py (100%) rename src/pepsy/optimizers/{mera => qmera}/geometry.py (100%) create mode 100644 src/pepsy/optimizers/qmera/layout.py rename src/pepsy/optimizers/{mera => qmera}/lightcones.py (98%) rename src/pepsy/optimizers/{mera => qmera}/parametric.py (89%) create mode 100644 src/pepsy/optimizers/qmera/prototype.py rename src/pepsy/optimizers/{mera => qmera}/schedules.py (91%) rename src/pepsy/optimizers/{mera => qmera}/schematics.py (83%) rename src/pepsy/optimizers/{mera => qmera}/terms.py (98%) rename tests/{test_optimize_mera.py => test_optimize_qmera.py} (67%) diff --git a/.github/skills/qmera-energy-optimizer/SKILL.md b/.github/skills/qmera-energy-optimizer/SKILL.md index 4a2369d..28479a1 100644 --- a/.github/skills/qmera-energy-optimizer/SKILL.md +++ b/.github/skills/qmera-energy-optimizer/SKILL.md @@ -1,12 +1,12 @@ --- name: qmera-energy-optimizer -description: "Design, implement, review, or extend Pepsy MERA/qMERA energy optimization in src/pepsy/optimizers/mera, including QMeraGeometry, QMeraBuilder schedules, parameterized two-qubit gate registries, reverse-lightcone energy chunks, compiled JAX/Torch losses, Symmray-native fermion gates/terms, schematics, and MeraEnergyOptimizer APIs." +description: "Design, implement, review, or extend Pepsy qMERA energy optimization in src/pepsy/optimizers/qmera, including QMeraGeometry, QMeraBuilder schedules, qMERA RG-layout search and scoring, parameterized gate registries, reverse-lightcone energy chunks, compiled JAX/Torch losses, Symmray-native fermion gates/terms, schematics, and QMeraEnergyOptimizer APIs." --- # qMERA Energy Optimizer in Pepsy -Use this skill for qMERA/QMERA-B, dense MERA, and `MeraEnergyOptimizer` work in -Pepsy. The current implementation is Pepsy-owned and schedule-first: Pepsy +Use this skill for qMERA/QMERA-B work in Pepsy. The current implementation is +Pepsy-owned and schedule-first: Pepsy defines geometry, RG blocking, gate placement, parameter dictionaries, local lightcone chunks, and optimizer shells; quimb/cotengra provide tensor-network storage and contraction. @@ -15,8 +15,8 @@ storage and contraction. - Repository rules: `AGENTS.md`. - Design reference: [references/design.md](references/design.md). -- Current source: `src/pepsy/optimizers/mera/`. -- Focused tests: `tests/test_optimize_mera.py`. +- Current source: `src/pepsy/optimizers/qmera/`. +- Focused tests: `tests/test_optimize_qmera.py`. - Public exports: `src/pepsy/optimizers/__init__.py`, `src/pepsy/__init__.py`, `tests/test_public_api.py`, and `tests/test_package_layout.py`. @@ -34,9 +34,6 @@ copied prototype scripts. ## Implementation Map -- `optimizer.py`: dense/isometric `MeraEnergyOptimizer` over existing MERA-like - tensor networks, with local lightcone energy chunks and quimb `TNOptimizer` - integration. - `terms.py`: `LocalTerm` normalization and backend conversion for local Hamiltonian inputs. - `geometry.py`: `QMeraGeometry` with explicit lattice labels, boundary, @@ -58,12 +55,19 @@ copied prototype scripts. `qmera_parametric_lightcone_tn(...)`, and `contract_qmera_lightcone_tn(...)` to rebuild and contract only the scheduled local cone for each Hamiltonian term. -- `compiled.py`: `cotengra.array_contract_expression` wrappers for dense - qMERA local cones. These freeze contraction topology so Torch/JAX see a pure - array loss over parameter dictionaries. -- `parametric.py`: `QMeraParametricEnergyOptimizer`, a - `GradientOptimizer`-based shell for parameter dictionaries, including - compiled-loss runs. +- `compiled.py`: `cotengra.array_contract_expression` wrappers for dense and + native graded qMERA local cones. Native Symmray expressions keep their + product-state/operator constants as Symmray arrays and freeze only the + contraction topology, so Torch/JAX parameter dictionaries remain + differentiable without dropping charge or fermionic-order metadata. +- `parametric.py`: `QMeraEnergyOptimizer`, a `GradientOptimizer`-based shell + for parameter dictionaries, including compiled-loss runs. The old + `QMeraParametricEnergyOptimizer` name remains only as a compatibility alias. +- `layout.py`: `QMeraLayoutFinder` and immutable candidate/score/report objects + for structural pre-ranking of RG architectures. +- `prototype.py`: a loader and stream-level scorer for serialized + `~/mera/U_q3_l*` placement streams; prototype streams remain diagnostics, + not Pepsy schedules. - `fermions.py`: Symmray-native fermion helpers, including `QMeraSymmrayFermionBackend`, `qmera_symmray_fermi_hubbard_terms(...)`, and `symmray_fermion_gate_registry(...)`. @@ -94,9 +98,9 @@ copied prototype scripts. only backend-native arrays in the parameter dictionary. - Use Pepsy backend helpers and autoray-compatible arrays. Do not introduce a qMERA-specific backend abstraction. -- Keep dense MERA projection and explicit qMERA circuit unitarity separate. - Dense MERA uses isometric tensor projection; qMERA uses unitary or - symmetry-preserving gate families. +- Keep qMERA gate unitarity and symmetry in the gate families. Do not add a + separate dense tensor projection or normalization path to the qMERA + optimizer. - Keep optional dependencies optional. Symmray-specific tests must use `pytest.importorskip("symmray")`. @@ -113,26 +117,102 @@ copied prototype scripts. neighbors should be adjacent in the active register. - Use `QMeraSymmrayFermionBackend.product_state(...)` through the builder's `product_state_factory` when contracting native Symmray lightcones. -- The compiled dense path in `compiled.py` is not automatically a Symmray - compiled path. Keep native Symmray contractions explicit until a tested - compiled graded-array route exists. -- Two-dimensional multi-mode schedules are intentionally guarded until the RG - design is explicit; do not silently flatten this case as if it were a normal - spin schedule. - -## Next Useful Work - -- Add user-facing docs/examples for `QMeraBuilder`, schematics, parametric - lightcone losses, compiled JAX/Torch loss usage, and Symmray-native - Fermi-Hubbard terms. -- Add local-cone grouping or reusable path-cache helpers around cotengra once - one-term-per-chunk correctness remains stable. -- Extend tests that compare direct full-qMERA TN energy to schedule-only - lightcone energy for more schedules and boundary conditions. -- Design the explicit 2D multi-mode/Fermi-Hubbard RG blocking before enabling - it in `build_qmera_schedule(...)`. -- Decide whether any qMERA symbols should become top-level `pepsy.*` exports; - if yes, update docs and public API tests in the same patch. +- Native Symmray compilation is supported when the builder receives a graded + `product_state_factory`, normally + `QMeraSymmrayFermionBackend.product_state`. The compiler must use Symmray's + autoray dispatch for pairwise contractions; never densify the frozen + constants or runtime gate blocks. `QMeraCompiledLightconeChunk.is_graded`, + `.symmetry`, and `.contraction_backend` expose this choice. +- Two-dimensional multi-mode schedules are supported when geometry and block + shapes are explicit. Preserve mode labels and same-flavor pairing; true + rectangular isometries remain unsupported and must not be mislabeled as + unitary completions. + +## Remaining Work + +- Add optional actual cotengra FLOP/peak-memory estimates to layout search + after its cheap structural pre-ranking, reusing the path cache. +- Extend compiled graded-array coverage to larger later-scale schedules and + optional device-specific benchmarks while retaining direct native cones as + the correctness oracle. +- Extend the prototype adapter with level-to-scale inference only when the + serialized stream format is formally specified; do not infer RG semantics + from a flat placement list by guesswork. +- Add broader later-scale direct-versus-lightcone comparisons and validate any + future true-isometry implementation separately from unitary completion. + +## qMERA RG Layout Finder + +The layout finder searches immutable RG architecture candidates and +return a reproducible `scales` plan that can be passed directly to +`QMeraBuilder`. It should not replace the schedule-first energy path or mutate +the builder while searching. + +The public objects are `QMeraLayoutCandidate`, `QMeraLayoutScore`, +`QMeraLayoutReport`, and `QMeraLayoutFinder`. A candidate records, for +every RG scale: + +- isometry block shape and orientation; +- disentangler placement (`boundary-faces`, `boundary-square`, or + `within-block`), corner policy, periodic wrapping, and executable rounds; +- internal circuit depths, gate families, and parameter-sharing policy; +- the resulting `QMeraScaleSpec` values and a stable candidate id. + +Candidate generation must validate the existing qMERA invariants before +scoring: + +- isometry blocks form a non-overlapping covering partition; +- concurrent disentangler supports are disjoint, while sequential circuit + rounds may reuse a block support; +- boundary disentanglers connect neighboring isometry regions and cover the + relevant interaction boundaries; +- reverse lightcones are finite, reproducible, and compatible with periodic + geometry, explicit modes, and the selected fermion symmetry; +- native Symmray candidates preserve mode labels, charge maps, and graded + contraction semantics without adding Jordan-Wigner strings. + +The score must expose components rather than hiding all decisions in one +opaque number. The initial components should include: + +- structural cost: gate count, circuit depth, number of placements, and + maximum/mean local-cone width; +- contraction cost: cotengra estimated FLOPs, peak intermediate size, and + path-search cost for representative local-cone topologies; +- interaction coverage: weighted coverage of Hamiltonian supports and + important interaction boundaries by the candidate's blocks and + disentanglers; +- optional entanglement coverage: weighted coverage of a user-supplied + mutual-information, entropy, correlation, or covariance map. + +Without state-derived data, report interaction coverage as a Hamiltonian proxy; +do not call it measured physical entanglement. For expensive searches, use a +cheap structural pre-ranking followed by actual cotengra path estimates for +the top candidates, reusing `QMeraContractionPathCache`. Return both a scalar +weighted score and a Pareto front so users can inspect cost-versus-coverage +tradeoffs. + +The intended workflow is: + +1. generate and structurally validate candidate scale plans; +2. pre-rank candidates using cone width, gate count, depth, and interaction + coverage; +3. evaluate cached contraction estimates for the top candidates; +4. optionally run a short pilot optimization and rescore with measured + entanglement coverage; +5. return the best plan, Pareto candidates, component scores, and schematic +metadata without silently changing mode order or fermion conventions. + +For comparison with the research prototype, call +`load_qmera_prototype_layout(...)` and +`finder.score_prototype_layout(...)`. That path reports flat-stream gate +count/depth and support coverage separately; it never converts `U_q3_l*` +into a `QMeraScaleSpec` without an explicit RG mapping. + +The finder should integrate with `QMeraBuilder`, `QMeraSchedule`, and +`draw_schematic` so a selected architecture can be inspected at every +`rg_step`. Add focused tests for candidate validity, deterministic ranking, +bounded lightcones, contraction-cost reporting, OBC/PBC layouts, and native +Fermi-Hubbard mode/symmetry preservation before exposing top-level exports. ## Validation @@ -140,7 +220,7 @@ Run focused validation after qMERA edits: ```bash env NUMBA_CACHE_DIR=/tmp/numba_cache MPLCONFIGDIR=/tmp/mplconfig PYTHONPYCACHEPREFIX=/tmp \ - /home/reza.haghshenas@quantinuum.com/envs/py312/bin/python -m pytest -q tests/test_optimize_mera.py + /home/reza.haghshenas@quantinuum.com/envs/py312/bin/python -m pytest -q tests/test_optimize_qmera.py ``` For API/export changes, also run: @@ -153,5 +233,11 @@ env NUMBA_CACHE_DIR=/tmp/numba_cache MPLCONFIGDIR=/tmp/mplconfig PYTHONPYCACHEPR For syntax-only checks: ```bash -/home/reza.haghshenas@quantinuum.com/envs/py312/bin/python -m pyflakes src/pepsy/optimizers/mera tests/test_optimize_mera.py +/home/reza.haghshenas@quantinuum.com/envs/py312/bin/python -m pyflakes src/pepsy/optimizers/qmera tests/test_optimize_qmera.py ``` + +The focused suite also covers independent 2D PBC Jordan-Wigner Fock-space +checks for every native Hubbard term, compiled native Symmray Hubbard +lightcones and Torch gradients, prototype-stream loading/scoring, +and canonical `pepsy.optimizers.qmera` imports with the temporary `mera` +compatibility alias. diff --git a/.github/skills/qmera-energy-optimizer/agents/openai.yaml b/.github/skills/qmera-energy-optimizer/agents/openai.yaml index 3c5251b..7668844 100644 --- a/.github/skills/qmera-energy-optimizer/agents/openai.yaml +++ b/.github/skills/qmera-energy-optimizer/agents/openai.yaml @@ -1,4 +1,4 @@ interface: display_name: "qMERA Energy Optimizer" - short_description: "Design and extend Pepsy qMERA/MERA optimization." + short_description: "Design and extend Pepsy qMERA optimization." default_prompt: "Use $qmera-energy-optimizer to design or extend Pepsy qMERA energy optimization from the current schedule-first implementation." diff --git a/.github/skills/qmera-energy-optimizer/references/design.md b/.github/skills/qmera-energy-optimizer/references/design.md index e35255a..51e5e06 100644 --- a/.github/skills/qmera-energy-optimizer/references/design.md +++ b/.github/skills/qmera-energy-optimizer/references/design.md @@ -1,4 +1,4 @@ -# qMERA / MERA Energy Optimizer Design Notes +# qMERA Energy Optimizer Design Notes ## Research Anchor @@ -8,23 +8,21 @@ doi:10.1103/PhysRevX.12.011047. Key design points for Pepsy: -- Dense MERA uses unitary disentanglers and isometric coarse-graining tensors. -- qMERA/QMERA-B keeps the MERA global structure but decomposes dense unitary and - isometric blocks into finite-depth local quantum circuits. +- qMERA/QMERA-B keeps the MERA RG structure but represents every block as a + finite-depth local quantum circuit. - The useful controls are `q` or local block width, internal circuit depth, gate count, and local circuit structure such as brickwall. - The energy objective used in the paper is exact contraction of local Hamiltonian terms, with global gradients supplied by automatic differentiation through quimb tensor networks. -- Unitary/isometric constraints should be enforced by differentiable projection - from parameters to constrained tensors or by explicit circuit gates. +- Unitary and symmetry constraints should be enforced by the parameterized + circuit gate families themselves. ## Current Pepsy State -As of July 2026, Pepsy has a real `src/pepsy/optimizers/mera/` package rather +As of July 2026, Pepsy has a real `src/pepsy/optimizers/qmera/` package rather than only a design sketch. The important implemented pieces are: -- dense/isometric `MeraEnergyOptimizer` for MERA-like tensor networks; - explicit `QMeraGeometry` lattice/register objects, including per-site mode expansion for fermionic registers; - RG-style `QMeraSchedule` objects with non-overlapping isometry blocks, @@ -38,10 +36,11 @@ than only a design sketch. The important implemented pieces are: - explicit local-cone tensor networks through `QMeraLightconeTN`, `qmera_parametric_lightcone_tn(...)`, and `contract_qmera_lightcone_tn(...)`; -- compiled dense local-cone contractions through cotengra - `array_contract_expression`, suitable for JAX/Torch array losses when the - schedule and chunks are static; -- `QMeraParametricEnergyOptimizer`, which routes parameter dictionaries through +- compiled local-cone contractions through cotengra + `array_contract_expression`, including a native graded Symmray route that + keeps Symmray constants and gate blocks intact for JAX/Torch autodiff when + the schedule and chunks are static; +- `QMeraEnergyOptimizer`, which routes parameter dictionaries through Pepsy `GradientOptimizer`; - Symmray-native fermion support for two-state modes, Fermi-Hubbard local terms, and a contextual `symmray-fsim` gate family. @@ -54,14 +53,14 @@ parametric loss route. ## quimb Anchors -quimb has two relevant layers: +quimb provides tensor-network storage and contraction for qMERA. Its circuit +and MERA examples are conceptual references only; Pepsy owns the schedule, +gate registry, parameter dictionary, and lightcone selection. -- `quimb.tensor.MERA`: stable 1D MERA class used in the public MERA example. - It supports causal-cone selection with site tags and works with - `qtn.TNOptimizer`. -- `quimb.experimental.merabuilder.TensorNetworkGenIso`: arbitrary-geometry - isometric builder with `layer_gate_fill_fn(operation="uni"|"iso"|"cap")`. - It is useful for 2D/qMERA prototypes but should remain behind a Pepsy adapter. +- `quimb.tensor.TensorNetwork` stores explicit direct-gate qMERA states and + local-cone tensor networks. +- `quimb.experimental.merabuilder.TensorNetworkGenIso` may inform geometry + prototypes but remains behind a Pepsy qMERA adapter. The quimb example computes local terms by: @@ -71,10 +70,6 @@ The quimb example computes local terms by: 4. join with the conjugate selected cone; 5. contract exactly with a reusable contraction optimizer. -For constrained global optimization, quimb uses a `norm_fn` that projects the -state with `unitize(method="exp")`; in new Pepsy code prefer the current -`isometrize(method="exp")` spelling. - For qMERA circuits, use the quimb quantum-circuit guide as conceptual guidance only: @@ -142,9 +137,8 @@ backend-native arrays in the parameter dictionary. Current layout: ```text -src/pepsy/optimizers/mera/ +src/pepsy/optimizers/qmera/ __init__.py - optimizer.py terms.py lightcones.py geometry.py @@ -155,11 +149,12 @@ src/pepsy/optimizers/mera/ parametric.py schematics.py fermions.py + layout.py + prototype.py ``` The implemented responsibilities are: -- `optimizer.py`: `MeraEnergyOptimizer` for dense/isometric MERA-like states. - `terms.py`: `LocalTerm`, input normalization, and local operator conversion. - `geometry.py`: explicit lattice labels, register sites, mapper support, and optional mode labels such as spin-up/spin-down. @@ -167,19 +162,21 @@ The implemented responsibilities are: - `gates.py`: parameterized gate registry and context-aware gate generation. - `builders.py`: `QMeraBuilder`, `QMeraAnsatz`, parameter casting, direct-gate state construction, and local-cone loss helpers. -- `lightcones.py`: tag-based dense MERA cones plus schedule-first qMERA cones. -- `compiled.py`: static contraction expressions for dense local qMERA cones. -- `parametric.py`: parameter-dict optimizer shell over `GradientOptimizer`. +- `lightcones.py`: schedule-first qMERA cones and direct-gate validation oracles. +- `compiled.py`: static contraction expressions for local qMERA cones. +- `parametric.py`: `QMeraEnergyOptimizer` over `GradientOptimizer`. - `schematics.py`: inspection drawings of disentangler/isometry blocking. - `fermions.py`: Symmray-native fermion mode backend and Fermi-Hubbard terms. +- `layout.py`: deterministic qMERA RG candidate generation and structural + scoring. +- `prototype.py`: diagnostic loading of serialized research-prototype gate + streams without treating them as Pepsy schedules. -Still-open package work: +Remaining package work: -- document examples for the implemented public path; -- local-cone grouping and cotengra path-cache ergonomics; -- explicit 2D multi-mode/Fermi-Hubbard RG design; -- broader comparison tests between direct-gate TNs and schedule-only chunks; -- possible top-level export decisions for selected qMERA symbols. +- optional actual cotengra cost estimates after structural layout ranking; +- broader later-scale direct-versus-lightcone comparisons; +- a formally specified mapping from flat prototype streams to RG scales. ## Design Data Flow @@ -225,9 +222,9 @@ enough for: - fermionic local terms after a caller-selected encoding such as Jordan-Wigner or a fermionic gate-aware ansatz. -`MeraEnergyOptimizer` should not build the Hamiltonian from model names. That -belongs in separate helpers or examples. The optimizer consumes normalized -local terms. +`QMeraBuilder` should not build the Hamiltonian from model names. That belongs +in separate helpers or examples. The qMERA optimizer consumes normalized local +terms. For Fermi-Hubbard, prefer Pepsy's existing symbolic/symmetric layer: @@ -308,9 +305,8 @@ like `iso_(...)` form covering blocks, `uni_(...)` forms shifted/wrapped boundary blocks, and 2D builders alternate horizontal/vertical unitary and isometry sublayers before capping. -For dense MERA, block tensors can be direct quimb isometries. For qMERA, dense -blocks are replaced by local parameterized circuit layers whose contraction -acts as the block tensor. Keep those modes explicit. +qMERA blocks are local parameterized circuit layers whose contraction acts as +the block tensor. Keep fermionic modes explicit. ## Parametrized Gate Registry @@ -381,22 +377,12 @@ solvers the natural optimization layer. ## JAX and Pepsy Gradient Route -There are two useful JAX patterns: - -1. Dense/isometric MERA tensor optimization: - - create a quimb MERA-like tensor network; - - call `qtn.pack(state)` once to get `(params, skeleton)`; - - define `loss_fn(params)` as `state = qtn.unpack(params, skeleton)`, optional - `state.isometrize(method=...)`, then local energy contraction; - - run either external Flax/Optax or Pepsy `GradientOptimizer` with a - `jax-*` solver. +qMERA gate-parameter optimization follows one schedule-first pattern: -2. qMERA gate-parameter optimization: - create a Pepsy `QMeraSchedule` from geometry/disentangler/isometry specs; - initialize a JAX pytree/dict of gate parameters; - - define `loss_fn(params)` by generating all gate tensors from params, - assembling the fixed qMERA tensor skeleton/chunks, and contracting local - lightcone chunks; + - define `loss_fn(params)` by generating parameterized gate tensors and + contracting static local qMERA lightcone chunks; - run `GradientOptimizer(params, loss_fn, solver="jax-adam"|"jax-adamw")` once the loss is pure and JAX-compatible. @@ -408,61 +394,26 @@ Implementation rule: JAX tracing should see a static schedule and static contraction/chunk topology. Dynamic inputs should be arrays in `params`, not new Python tensor-network objects, changing tag sets, or file-loaded schedules. -## `MeraEnergyOptimizer` Shape +## `QMeraEnergyOptimizer` Shape -Mirror `PepsEnergyOptimizer` where possible: +The canonical optimizer owns a parameter dictionary and a static schedule: ```python -class MeraEnergyOptimizer: - def __init__( - self, - state, - hamiltonian, - *, - normalized=True, - energy_per_site=True, - real=True, - isometrize_method="exp", - contraction_opt="auto-hq", - backend="auto", - jit=False, - compute_kwargs=None, - loss_kwargs=None, - ): ... - - def loss(self, state=None, *, hamiltonian=None, terms=None, **kwargs): ... - def energy(self, state=None, *, hamiltonian=None, terms=None, **kwargs): ... - def make_tn_optimizer(...): ... - def optimize(...): ... +class QMeraEnergyOptimizer: + builder: QMeraBuilder + schedule: QMeraSchedule + hamiltonian: object + parameters: Mapping[str, object] + + def loss(self, parameters=None, **kwargs): ... + def compile(self, **kwargs): ... + def compiled_loss(self, parameters=None, **kwargs): ... + def run(self, params_init=None, *, solver="torch-adam", **kwargs): ... ``` -Candidate loss kwargs: - -- `normalized`: whether to divide by local norm if state is not guaranteed - isometric. Default can be `True` only if implemented cheaply and tested. -- `energy_per_site`: divide by inferred number of physical sites. -- `real`: return `autoray.real(...)`. -- `isometrize_method`: method used by `norm_fn`, commonly `"exp"`. -- `contraction_opt`: exact contraction optimizer for causal cones. -- `precompute_tags`: cache causal-cone tags for fixed term supports. -- `chunk_terms`: group compatible local terms by identical or similar - lightcone selectors. -- `jit`: request backend JIT when the loss graph is static, especially JAX. -- `solver`: Pepsy solver name such as `"jax-adam"`, `"jax-adamw"`, - `"torch-adam"`, or `"torch-lbfgs"` when optimizing explicit parameter dicts. -- `simplify`: optional local simplification such as `full_simplify(seq="R")`. - -`make_tn_optimizer()` should pass: - -- `loss_fn`: static adapter around `_loss_state`. -- `norm_fn`: `lambda state: state.isometrize(method=isometrize_method)` when - using dense isometric tensors. -- `loss_constants`: normalized terms and any cached tag selectors. -- `loss_kwargs`: contraction and scalar-format options. - -For qMERA circuit builders with explicit unitary gates, `norm_fn` might be -unneeded because circuit gates already enforce unitarity. Keep this a builder -or ansatz property rather than auto-detecting silently. +Built-in qMERA gates are unitary or symmetry-preserving by construction, so +the optimizer defaults to `normalized=False`. A caller may explicitly enable +normalization for a custom non-unitary gate family. ## Term Normalization @@ -666,8 +617,8 @@ than rebuilding tensor networks inside a traced function. - Jordan-Wigner Fermi-Hubbard hopping is a two-site gate only for adjacent mapped bonds. For non-adjacent 2D mapped bonds, use the Pepsy MPO/term path or an explicitly provided long-string gate representation. -- Dense MERA projection and explicit qMERA circuit unitarity are different - constraint mechanisms. Do not mix them silently in one code path. +- qMERA circuit unitarity and fermionic symmetry preservation are gate-family + constraints. Do not add a separate tensor projection path. - A DMRG-like local gate optimizer from the paper is a later milestone. Start with global autodiff since it matches existing Pepsy energy optimizer style. @@ -677,12 +628,11 @@ Run the focused qMERA suite after implementation changes: ```bash env NUMBA_CACHE_DIR=/tmp/numba_cache MPLCONFIGDIR=/tmp/mplconfig PYTHONPYCACHEPREFIX=/tmp \ - /home/reza.haghshenas@quantinuum.com/envs/py312/bin/python -m pytest -q tests/test_optimize_mera.py + /home/reza.haghshenas@quantinuum.com/envs/py312/bin/python -m pytest -q tests/test_optimize_qmera.py ``` The focused suite should continue to cover: -- dense 1D MERA local expectation and energy smoke tests; - causal-cone metadata and bounded schedule-width diagnostics; - `QMeraGeometry` lattice, mapper, register, and mode ordering behavior; - 1D and 2D RG schedules with non-overlapping isometry blocks and boundary @@ -692,8 +642,8 @@ The focused suite should continue to cover: - direct-gate TN construction as a debugging/comparison path; - schedule-first parametric lightcone chunks that rebuild only local cones; - explicit `QMeraLightconeTN` construction and cotengra contraction; -- compiled dense local-cone contractions matching rebuilt local cones; -- Torch optimizer smoke through `QMeraParametricEnergyOptimizer`; +- compiled qMERA local-cone contractions matching rebuilt local cones; +- Torch optimizer smoke through `QMeraEnergyOptimizer`; - JAX JIT smoke for the compiled parameter-dict loss, skipped cleanly when JAX or Optax is unavailable; - Symmray-native Fermi-Hubbard local terms and `symmray-fsim` lightcone tests, @@ -709,5 +659,5 @@ env NUMBA_CACHE_DIR=/tmp/numba_cache MPLCONFIGDIR=/tmp/mplconfig PYTHONPYCACHEPR For syntax-only checks: ```bash -/home/reza.haghshenas@quantinuum.com/envs/py312/bin/python -m pyflakes src/pepsy/optimizers/mera tests/test_optimize_mera.py +/home/reza.haghshenas@quantinuum.com/envs/py312/bin/python -m pyflakes src/pepsy/optimizers/qmera tests/test_optimize_qmera.py ``` diff --git a/docs/api/index.md b/docs/api/index.md index 4564b29..eda9c68 100644 --- a/docs/api/index.md +++ b/docs/api/index.md @@ -30,7 +30,7 @@ the detailed functions and classes for each area. - [Sweep optimization](optimizers/sweep.md) - [Global optimization](optimizers/global_opt.md) - [Energy optimization](optimizers/energy.md) -- [MERA](optimizers/mera.md) +- [qMERA](optimizers/qmera.md) - [Noise and trajectories](optimizers/noise.md) - [Simulator planning](optimizers/planning.md) - [Stabilizer tensor networks](optimizers/stabilizer_tn.md) diff --git a/docs/api/optimizers/mera.md b/docs/api/optimizers/qmera.md similarity index 77% rename from docs/api/optimizers/mera.md rename to docs/api/optimizers/qmera.md index 24ff6d0..2998b2d 100644 --- a/docs/api/optimizers/mera.md +++ b/docs/api/optimizers/qmera.md @@ -1,29 +1,14 @@ -# `pepsy.optimizers.mera` +# `pepsy.optimizers.qmera` -MERA and qMERA energy helpers evaluate local Hamiltonian terms through reverse -lightcones rather than by contracting a full state for every term. - -The dense MERA path wraps an existing MERA-like tensor network: +Pepsy's optimizer surface is qMERA-only: parameterized gate families are +placed by a static RG schedule, and local Hamiltonian terms are evaluated by +rebuilding only their reverse lightcones. ```python import numpy as np -import quimb.tensor as qtn - -from pepsy.optimizers import MeraEnergyOptimizer zz = np.diag([1.0, -1.0]) h2 = np.kron(zz, zz).reshape(2, 2, 2, 2) -mera = qtn.MERA.rand(L=8, max_bond=2, dtype="complex128", seed=1) - -opt = MeraEnergyOptimizer(mera, {(0, 1): h2, (2, 3): h2}) -estimate = opt.energy() -``` - -The qMERA path is schedule-first. `QMeraBuilder` creates the geometry, RG -schedule, parameter dictionary, and local-cone chunks. The loss then rebuilds -only the selected lightcone for each local term. - -```python from pepsy.optimizers import QMeraBuilder, build_qmera_contraction_optimizer builder = QMeraBuilder( @@ -46,27 +31,7 @@ energy = builder.parametric_loss( ) ``` -For a fixed MERA-like state, `lightcone_energy(...)` exposes the same local -contraction primitive directly. It selects only the reverse cone, applies each -local operator with `.gate`, and can reuse one Pepsy/cotengra path per cone -topology: - -```python -from pepsy.optimizers.mera import lightcone_energy - -path_cache = builder.contraction_path_cache(max_repeats=16) -energy = lightcone_energy( - mera, - {(0, 1): h2, (2, 3): h2}, - energy_per_site=False, - path_cache=path_cache, -) -``` - -This is the fixed-state analogue of `builder.parametric_loss(...)`; the latter -rebuilds the qMERA cone from a parameter dictionary on every evaluation. - -For repeated optimization, compile static local-cone contractions once and +For repeated optimization, compile static qMERA local-cone contractions once and reuse them from NumPy, Torch, or JAX-compatible parameter dictionaries: ```python @@ -82,8 +47,49 @@ loss_fn = builder.compiled_parametric_loss_fn( energy = loss_fn(params) ``` -`QMeraParametricEnergyOptimizer` routes the same compiled loss through Pepsy's -gradient solvers: +Native Symmray qMERA uses the same compiled API, but the builder must receive +the graded product-state factory. The compiled object then reports +`contraction_backend="symmray"` and `is_graded=True`; its frozen constants and +runtime gate blocks stay block-sparse Symmray arrays, so fermionic signs, +charge maps, duals, and Torch autodiff are retained: + +```python +from pepsy.optimizers.qmera import ( + QMeraBuilder, + QMeraGeometry, + QMeraSymmrayFermionBackend, + symmray_fermion_gate_registry, +) + +backend = QMeraSymmrayFermionBackend(symmetry="U1U1") +builder = QMeraBuilder( + geometry=QMeraGeometry( + shape=(4, 4), + boundary="periodic", + site_modes=("up", "down"), + mode_order="mode-major", + ), + gate_registry=symmray_fermion_gate_registry(backend=backend), + gate_family="symmray-fsim", + product_state_factory=backend.product_state, +) +compiled = builder.compile_parametric_lightcones( + terms, + convert_terms=False, + path_cache=builder.contraction_path_cache(max_repeats=16), +) +assert compiled[0].is_graded +``` + +`path_cache` reuses topology-specific cotengra searches while each compiled +expression reuses its frozen contraction tree on every parameter evaluation. +Use `builder.parametric_loss(...)` or `builder.direct_parametric_loss(...)` as +the explicit native correctness oracle when validating a new periodic layout. + +`QMeraEnergyOptimizer` routes the same compiled loss through Pepsy's gradient +solvers. Its built-in spin and fermion gate families are parameterized and +unitary/symmetry-preserving by construction, so normalization is disabled by +default. Pass `normalized=True` for a custom non-unitary gate family: ```python param_opt = builder.parametric_optimizer( @@ -105,17 +111,51 @@ patches and arrows for the RG flow: ```python drawing = schedule.draw_schematic( + rg_step=0, # inspect one bottom-to-top RG step style="clean", # or "register" for the low-level wiring view - figsize=(14, 5), + figsize=(16, 5), label_sites=True, label_blocks=True, scale_figsize=False, ) ``` -The clean view is intended for explaining a schedule or a fermionic block -layout; `schedule.schematic_blocks()` remains the machine-readable placement -audit. +The clean view is ordered as input → disjoint disentangler subrounds → +covering isometry subrounds → coarse output. Disentangler windows overlap the +neighboring isometry blocks by design, but blocks in the same executable +subround are disjoint. Set `rg_step=1` (or another valid scale) to inspect a +later step; use `rg_step=None` to draw all steps. The older `layer=` selector +remains an alias. `schedule.schematic_blocks()` remains the machine-readable +placement audit. + +### qMERA layout search and prototype comparison + +`QMeraLayoutFinder` searches valid immutable RG scale plans and ranks them by +gate count, executable depth, reverse-lightcone width, and Hamiltonian-support +coverage. The result can be passed directly to `QMeraBuilder`: + +```python +from pepsy.optimizers.qmera import QMeraGeometry, QMeraLayoutFinder + +geometry = QMeraGeometry(shape=(6, 6), boundary="periodic") +report = QMeraLayoutFinder(geometry, max_layers=3).search({(0, 1): h2}) +builder = QMeraBuilder(geometry=geometry, scales=report.best.scales) +``` + +The research prototype's serialized `U_q3_l*` files are flat gate-placement +streams, not RG schedules. Load them only for structural comparison: + +```python +from pepsy.optimizers.qmera import load_qmera_prototype_layout + +prototype = load_qmera_prototype_layout("/home/.../mera/U_q3_l1") +prototype_score = QMeraLayoutFinder( + QMeraGeometry(shape=prototype.num_sites) +).score_prototype_layout(prototype, {(0, 1): h2}) +``` + +This adapter deliberately does not infer qMERA scales from the prototype +stream. Native Symmray fermion helpers are available under this module, but the fermion convention is explicit. The `Fermion` helper can now be supplied to @@ -127,7 +167,7 @@ terms: ```python import pepsy -from pepsy.optimizers.mera import QMeraGeometry +from pepsy.optimizers.qmera import QMeraGeometry fermion = pepsy.Fermion( spinful=True, @@ -150,7 +190,7 @@ For the normal spinful Hubbard workflow, let the builder own the mode expansion and conversion: ```python -from pepsy.optimizers.mera import ( +from pepsy.optimizers.qmera import ( QMeraBuilder, QMeraSymmrayFermionBackend, symmray_fermion_gate_registry, @@ -201,7 +241,7 @@ covering blocks, then reduces 3x3 to one site with a 3x3 covering block and vertical 3-site internal disentangler strips: ```python -from pepsy.optimizers.mera import ( +from pepsy.optimizers.qmera import ( QMeraBuilder, QMeraDisentanglerSpec, QMeraIsometrySpec, @@ -246,7 +286,7 @@ spinful Fermi--Hubbard workflow therefore uses `U1U1` and the Symmray FSIM registry: ```python -from pepsy.optimizers.mera import ( +from pepsy.optimizers.qmera import ( QMeraBuilder, QMeraGeometry, QMeraSymmrayFermionBackend, @@ -278,7 +318,7 @@ For the explicit 4x4 periodic construction, use a 2x2 square disentangler around every inter-block face and a 2x2 covering unitary for each RG block: ```python -from pepsy.optimizers.mera import ( +from pepsy.optimizers.qmera import ( QMeraBuilder, QMeraDisentanglerSpec, QMeraGeometry, @@ -383,7 +423,7 @@ physical site with native `Z2` fermion parity: ```python import pepsy as py -from pepsy.optimizers.mera import ( +from pepsy.optimizers.qmera import ( QMeraGeometry, qmera_symmray_majorana_terms, symmray_majorana_gate_registry, diff --git a/docs/api/package.md b/docs/api/package.md index 1620e6a..91e3aca 100644 --- a/docs/api/package.md +++ b/docs/api/package.md @@ -39,7 +39,7 @@ from pepsy.tensors import OneDMap, ps_to_mps, ps_to_peps, tn_norm | --- | --- | --- | | Belief propagation | `pepsy.bp` | BP, relay gauges, loop corrections, and PNE | | VMC | `pepsy.vmc` | Torch and NetKet/JAX variational Monte Carlo | -| MERA | `pepsy.optimizers.mera` | qMERA geometry and energy optimization | +| qMERA | `pepsy.optimizers.qmera` | qMERA geometry, gates, and energy optimization | | Stabilizer TN | `pepsy.optimizers.stabilizer_tn` | Stim tableau plus coefficient-MPS simulation | | Tree TN | `pepsy.optimizers.tree` | Tree layout and circuit replay | | Tree stabilizer | `pepsy.optimizers.tree_stabilizer` | Tableau plus tree-coefficient simulation | diff --git a/docs/development/modules/optimizers.md b/docs/development/modules/optimizers.md index 7a472e9..4b3bf3b 100644 --- a/docs/development/modules/optimizers.md +++ b/docs/development/modules/optimizers.md @@ -35,8 +35,8 @@ important downstream time-compression consumer that depends on Pepsy behavior. - `planning.py`: non-executing physical-versus-stabilizer and MPS-versus-tree circuit advice using measured frame supports and explicit chi-scaled work proxies. -- `mera/`: dense MERA and schedule-first qMERA local-energy objectives, - parameter dictionaries, compiled lightcone contractions, schematics, and +- `qmera/`: schedule-first qMERA local-energy objectives, parameter + dictionaries, compiled lightcone contractions, schematics, and Symmray-native fermion helpers. - `global_opt.py`: whole-network variational optimization helpers. diff --git a/examples/qmera_fermion_hubbard_2d.py b/examples/qmera_fermion_hubbard_2d.py index 2e0d110..36b4644 100644 --- a/examples/qmera_fermion_hubbard_2d.py +++ b/examples/qmera_fermion_hubbard_2d.py @@ -6,7 +6,7 @@ """ import pepsy as py -from pepsy.optimizers.mera import ( +from pepsy.optimizers.qmera import ( QMeraBuilder, QMeraGeometry, QMeraSymmrayFermionBackend, diff --git a/examples/qmera_fermion_hubbard_4x4_pbc.py b/examples/qmera_fermion_hubbard_4x4_pbc.py index 8a5c670..dc423ae 100644 --- a/examples/qmera_fermion_hubbard_4x4_pbc.py +++ b/examples/qmera_fermion_hubbard_4x4_pbc.py @@ -6,7 +6,7 @@ """ import pepsy as py -from pepsy.optimizers.mera import ( +from pepsy.optimizers.qmera import ( QMeraBuilder, QMeraDisentanglerSpec, QMeraGeometry, diff --git a/examples/qmera_majorana_2d.py b/examples/qmera_majorana_2d.py index ce06414..00d6165 100644 --- a/examples/qmera_majorana_2d.py +++ b/examples/qmera_majorana_2d.py @@ -6,7 +6,7 @@ """ import pepsy as py -from pepsy.optimizers.mera import ( +from pepsy.optimizers.qmera import ( QMeraBuilder, QMeraGeometry, QMeraSymmrayFermionBackend, diff --git a/examples/qmera_scale_plan_6x6.py b/examples/qmera_scale_plan_6x6.py index 09858a9..47ba2b0 100644 --- a/examples/qmera_scale_plan_6x6.py +++ b/examples/qmera_scale_plan_6x6.py @@ -1,6 +1,6 @@ """Generic heterogeneous 6x6 periodic qMERA scale-plan example.""" -from pepsy.optimizers.mera import ( +from pepsy.optimizers.qmera import ( QMeraBuilder, QMeraDisentanglerSpec, QMeraIsometrySpec, diff --git a/src/pepsy/experimental/__init__.py b/src/pepsy/experimental/__init__.py index c50dd31..7f3a95d 100644 --- a/src/pepsy/experimental/__init__.py +++ b/src/pepsy/experimental/__init__.py @@ -9,7 +9,8 @@ _MODULES = { "bp": "pepsy.bp", - "mera": "pepsy.optimizers.mera", + "mera": "pepsy.optimizers.qmera", + "qmera": "pepsy.optimizers.qmera", "stabilizer": "pepsy.optimizers.stabilizer_tn", "symmetry": "pepsy.tensors.symmetric", "tree": "pepsy.optimizers.tree", diff --git a/src/pepsy/optimizers/__init__.py b/src/pepsy/optimizers/__init__.py index 07ff161..5f5cced 100644 --- a/src/pepsy/optimizers/__init__.py +++ b/src/pepsy/optimizers/__init__.py @@ -13,15 +13,21 @@ "PepsEnergyOptimizer": ".energy", "TreeEnergyOptimizer": ".energy", "GlobalOptimizer": ".global_opt", - "MeraEnergyOptimizer": ".mera", - "QMeraBuilder": ".mera", - "QMeraDisentanglerSpec": ".mera", - "QMeraGeometry": ".mera", - "QMeraIsometrySpec": ".mera", - "QMeraParametricEnergyOptimizer": ".mera", - "QMeraScaleSpec": ".mera", - "QMeraUnitarySpec": ".mera", - "build_qmera_contraction_optimizer": ".mera", + "QMeraBuilder": ".qmera", + "QMeraDisentanglerSpec": ".qmera", + "QMeraEnergyOptimizer": ".qmera", + "QMeraGeometry": ".qmera", + "QMeraIsometrySpec": ".qmera", + "QMeraLayoutCandidate": ".qmera", + "QMeraLayoutFinder": ".qmera", + "QMeraLayoutReport": ".qmera", + "QMeraLayoutScore": ".qmera", + "QMeraParametricEnergyOptimizer": ".qmera", + "QMeraPrototypeLayout": ".qmera", + "QMeraScaleSpec": ".qmera", + "QMeraUnitarySpec": ".qmera", + "build_qmera_contraction_optimizer": ".qmera", + "load_qmera_prototype_layout": ".qmera", "MpoOptimizer": ".mpo", "MpsOptimizer": ".mps", "SimulatorCandidate": ".planning", @@ -100,6 +106,7 @@ "energy", "global_opt", "mera", + "qmera", "mpo", "mps", "noise", diff --git a/src/pepsy/optimizers/mera/__init__.py b/src/pepsy/optimizers/mera/__init__.py index 080cf45..39063b4 100644 --- a/src/pepsy/optimizers/mera/__init__.py +++ b/src/pepsy/optimizers/mera/__init__.py @@ -1,128 +1,35 @@ -"""MERA and qMERA energy optimization helpers.""" +"""Compatibility namespace for the former qMERA package name. -from .builders import QMeraAnsatz, QMeraBuilder -from .cache import QMeraContractionPathCache, build_qmera_contraction_optimizer -from .compiled import ( - QMeraCompiledLightconeChunk, - compile_qmera_parametric_lightcone, - compile_qmera_parametric_lightcones, - local_qmera_compiled_lightcone_expectation, - qmera_compiled_parametric_energy, -) -from .fermions import ( - QMeraSymmrayFermionBackend, - qmera_symmray_fermi_hubbard_terms, - qmera_symmray_majorana_terms, - symmray_fermion_gate_registry, - symmray_majorana_gate_registry, -) -from .gates import ( - GateRegistry, - GateSpec, - UserGateFamily, - default_gate_registry, - resolve_gate_spec, -) -from .geometry import QMeraGeometry -from .lightcones import ( - LightconeChunk, - QMeraLightconeTN, - QMeraLightconeGroup, - QMeraParametricLightconeChunk, - build_lightcone_chunks, - build_qmera_lightcone_chunks, - build_qmera_parametric_lightcone_chunks, - contract_qmera_lightcone_group, - contract_qmera_lightcone_tn, - group_qmera_parametric_lightcone_chunks, - lightcone_energy, - local_qmera_parametric_lightcone_expectation, - local_lightcone_expectation, - qmera_parametric_energy, - qmera_direct_parametric_energy, - qmera_parametric_lightcone_group_state, - qmera_parametric_state, - qmera_parametric_lightcone_state, - qmera_parametric_lightcone_tn, - select_lightcone, - site_tags_for_where, -) -from .optimizer import MeraEnergyOptimizer -from .parametric import QMeraParametricEnergyOptimizer -from .schedules import ( - QMeraBlockSpec, - QMeraDisentanglerSpec, - QMeraGatePlacement, - QMeraIsometrySpec, - QMeraLayerSpec, - QMeraScaleSpec, - QMeraSchedule, - QMeraUnitarySpec, - build_qmera_schedule, -) -from .schematics import ( - QMeraSchematicBlock, - draw_qmera_schedule, - qmera_schematic_blocks, -) -from .terms import LocalTerm, normalize_local_terms +Use :mod:`pepsy.optimizers.qmera` for all new code. This module only aliases +the canonical qMERA implementation so existing imports keep working during +the namespace migration. +""" -__all__ = [ - "GateRegistry", - "GateSpec", - "LightconeChunk", - "LocalTerm", - "MeraEnergyOptimizer", - "QMeraAnsatz", - "QMeraBlockSpec", - "QMeraBuilder", - "QMeraCompiledLightconeChunk", - "QMeraContractionPathCache", - "QMeraDisentanglerSpec", - "QMeraGatePlacement", - "QMeraGeometry", - "QMeraIsometrySpec", - "QMeraLayerSpec", - "QMeraLightconeTN", - "QMeraLightconeGroup", - "QMeraParametricLightconeChunk", - "QMeraParametricEnergyOptimizer", - "QMeraSchedule", - "QMeraScaleSpec", - "QMeraSchematicBlock", - "QMeraSymmrayFermionBackend", - "QMeraUnitarySpec", - "UserGateFamily", - "build_lightcone_chunks", - "build_qmera_contraction_optimizer", - "build_qmera_lightcone_chunks", - "build_qmera_parametric_lightcone_chunks", - "build_qmera_schedule", - "compile_qmera_parametric_lightcone", - "compile_qmera_parametric_lightcones", - "contract_qmera_lightcone_tn", - "contract_qmera_lightcone_group", - "default_gate_registry", - "draw_qmera_schedule", - "local_qmera_compiled_lightcone_expectation", - "local_qmera_parametric_lightcone_expectation", - "local_lightcone_expectation", - "normalize_local_terms", - "qmera_compiled_parametric_energy", - "qmera_direct_parametric_energy", - "qmera_parametric_energy", - "qmera_parametric_lightcone_group_state", - "qmera_parametric_lightcone_state", - "qmera_parametric_state", - "qmera_parametric_lightcone_tn", - "qmera_schematic_blocks", - "qmera_symmray_fermi_hubbard_terms", - "qmera_symmray_majorana_terms", - "resolve_gate_spec", - "group_qmera_parametric_lightcone_chunks", - "lightcone_energy", - "select_lightcone", - "site_tags_for_where", - "symmray_fermion_gate_registry", - "symmray_majorana_gate_registry", -] +from importlib import import_module +import sys + +from ..qmera import * # noqa: F401,F403 +from ..qmera import __all__ as __all__ + + +for _module_name in ( + "builders", + "cache", + "compiled", + "fermions", + "gates", + "geometry", + "layout", + "lightcones", + "parametric", + "prototype", + "schedules", + "schematics", + "terms", +): + sys.modules.setdefault( + f"{__name__}.{_module_name}", + import_module(f"..qmera.{_module_name}", __name__), + ) + +del import_module, sys, _module_name diff --git a/src/pepsy/optimizers/mera/optimizer.py b/src/pepsy/optimizers/mera/optimizer.py deleted file mode 100644 index 26a3eca..0000000 --- a/src/pepsy/optimizers/mera/optimizer.py +++ /dev/null @@ -1,468 +0,0 @@ -"""MERA energy objective and optimization shell.""" - -from __future__ import annotations - -from collections.abc import Mapping -from typing import Any - -import autoray as ar -import numpy as np -import quimb.tensor as qtn - -from ...backends import ( - get_default_array_backend, - infer_backend_converter_from_sample, - resolve_backend_sample_data_from_tn, -) -from ...tensors import build_optimizer, reg_rel_svd_jax, reg_rel_svd_torch -from ..energy import EnergyEstimate -from ..global_opt import GlobalOptimizer -from .lightcones import ( - build_lightcone_chunks, - build_qmera_lightcone_chunks, - local_lightcone_expectation, -) -from .terms import convert_local_terms, normalize_local_terms - -__all__ = ["MeraEnergyOptimizer"] - - -class MeraEnergyOptimizer: - """Evaluate and optimize local energies of MERA-like tensor networks. - - This first implementation assumes a fixed MERA-like tensor network with - physical site tags. It contracts each Hamiltonian term over its selected - reverse lightcone rather than over the full state. - """ - - _LOSS_KEYS = frozenset({ - "normalized", - "energy_per_site", - "real", - "contraction_opt", - "array_backend", - "convert_terms", - "precompute_tags", - "simplify", - "gate_contract", - "contract_opts", - }) - - def __init__( - self, - state, - hamiltonian, - *, - normalized: bool = True, - energy_per_site: bool = True, - real: bool = True, - isometrize_method: str | None = "exp", - contraction_opt: Any = "auto-hq", - schedule=None, - array_backend=None, - convert_terms: bool = True, - precompute_tags: bool = True, - simplify: bool | str = False, - gate_contract: bool = True, - contract_opts: Mapping[str, Any] | None = None, - loss_kwargs: Mapping[str, Any] | None = None, - ): - self.schedule = self._resolve_schedule(state, schedule) - self.state = self._as_mera_state(state) - self.hamiltonian = hamiltonian - self.terms = normalize_local_terms(hamiltonian) - self.losses: list[float] = [] - self.isometrize_method = isometrize_method - self.loss_kwargs = { - "normalized": normalized, - "energy_per_site": energy_per_site, - "real": real, - "contraction_opt": contraction_opt, - "array_backend": array_backend, - "convert_terms": convert_terms, - "precompute_tags": precompute_tags, - "simplify": simplify, - "gate_contract": gate_contract, - "contract_opts": {} if contract_opts is None else dict(contract_opts), - } - if loss_kwargs is not None: - self.set_loss_kwargs(**loss_kwargs) - self.lightcones = self._build_chunks_if_requested( - self.state, - self.terms, - schedule=self.schedule, - ) - - @classmethod - def loss_kwarg_names(cls): - """Return supported loss keyword names.""" - return tuple(sorted(cls._LOSS_KEYS)) - - @staticmethod - def _merge_opts(base, extra): - opts = dict(base or {}) - if extra: - opts.update(dict(extra)) - return opts - - @classmethod - def _pick_loss_kwargs(cls, options): - incoming = dict(options or {}) - unknown = sorted(set(incoming) - cls._LOSS_KEYS) - if unknown: - allowed = ", ".join(sorted(cls._LOSS_KEYS)) - raise TypeError( - f"Unknown MERA energy option(s): {', '.join(unknown)}. " - f"Allowed options: {allowed}." - ) - return incoming - - @staticmethod - def _as_mera_state(state): - if hasattr(state, "select") and hasattr(state, "gate"): - return state - ansatz_state = getattr(state, "state", None) - if ( - ansatz_state is not None - and hasattr(ansatz_state, "select") - and hasattr(ansatz_state, "gate") - ): - return ansatz_state - tn = getattr(state, "tn", None) - if tn is not None and hasattr(tn, "select") and hasattr(tn, "gate"): - return tn - raise TypeError( - "state must be a MERA-like TensorNetwork with select() and gate()." - ) - - @staticmethod - def _resolve_schedule(state, schedule=None): - if schedule is not None: - return schedule - return getattr(state, "schedule", None) - - @staticmethod - def _num_sites(state): - num_sites = getattr(state, "num_sites", None) - if num_sites is not None: - return int(num_sites() if callable(num_sites) else num_sites) - sites = getattr(state, "sites", None) - if sites is not None: - return len(tuple(sites)) - length = getattr(state, "L", None) - if length is not None: - return int(length) - raise ValueError("Could not infer the number of MERA sites.") - - @staticmethod - def _max_bond(state): - max_bond = getattr(state, "max_bond", None) - if callable(max_bond): - return max_bond() - return max_bond - - @staticmethod - def _maybe_real(value): - try: - return ar.do("real", value) - except Exception: # pragma: no cover - defensive for unusual scalar types - return value.real - - @staticmethod - def _is_finite_number(value): - try: - return bool(np.isfinite(float(value))) - except (TypeError, ValueError): - return False - - @staticmethod - def _prepare_autodiff_backend(backend): - key = str(backend).strip().lower() - try: - if key == "torch": - reg_rel_svd_torch() - elif key == "jax": - reg_rel_svd_jax() - except ImportError: - return - - @classmethod - def _array_backend_for_state(cls, state, explicit_backend=None): - if explicit_backend is not None: - return explicit_backend - default_backend = get_default_array_backend() - if default_backend is not None: - return default_backend - sample = resolve_backend_sample_data_from_tn(state) - return infer_backend_converter_from_sample(sample) - - @classmethod - def _prepare_terms(cls, state, terms, *, array_backend=None, convert_terms=True): - terms = normalize_local_terms(terms) - if not convert_terms: - return terms - backend = cls._array_backend_for_state(state, array_backend) - return convert_local_terms(terms, backend) - - def _build_chunks_if_requested(self, state, terms, *, schedule=None, opts=None): - opts = self.loss_kwargs if opts is None else opts - if not opts.get("precompute_tags", True): - return None - prepared_terms = self._prepare_terms( - state, - terms, - array_backend=opts.get("array_backend"), - convert_terms=opts.get("convert_terms", True), - ) - if schedule is not None: - return build_qmera_lightcone_chunks(state, schedule, prepared_terms) - return build_lightcone_chunks(state, prepared_terms) - - @classmethod - def _chunks_for_loss(cls, state, terms, *, chunks=None, schedule=None, **opts): - if chunks is not None: - return chunks - prepared_terms = cls._prepare_terms( - state, - terms, - array_backend=opts.get("array_backend"), - convert_terms=opts.get("convert_terms", True), - ) - if schedule is not None: - return build_qmera_lightcone_chunks(state, schedule, prepared_terms) - return build_lightcone_chunks(state, prepared_terms) - - @classmethod - def _loss_state( - cls, - state, - *, - terms, - chunks=None, - schedule=None, - normalized=True, - energy_per_site=True, - real=True, - contraction_opt="auto-hq", - array_backend=None, - convert_terms=True, - precompute_tags=True, - simplify=False, - gate_contract=True, - contract_opts=None, - ): - del precompute_tags - state = cls._as_mera_state(state) - if contraction_opt is None: - contraction_opt = build_optimizer(progbar=False) - chunks = cls._chunks_for_loss( - state, - terms, - chunks=chunks, - schedule=schedule, - array_backend=array_backend, - convert_terms=convert_terms, - ) - value = None - for chunk in chunks: - term_value = local_lightcone_expectation( - state, - chunk, - optimize=contraction_opt, - normalized=normalized, - real=False, - simplify=simplify, - gate_contract=gate_contract, - contract_opts=contract_opts, - ) - value = term_value if value is None else value + term_value - if value is None: - raise ValueError("hamiltonian contains no local terms.") - if energy_per_site: - value = value / cls._num_sites(state) - if real: - value = cls._maybe_real(value) - return value - - @staticmethod - def _tnopt_loss(state, *, terms, chunks=None, schedule=None, **loss_kwargs): - """Adapter for :class:`quimb.tensor.TNOptimizer`.""" - return MeraEnergyOptimizer._loss_state( - state, - terms=terms, - chunks=chunks, - schedule=schedule, - **loss_kwargs, - ) - - def set_loss_kwargs(self, **kwargs): - """Update stored defaults for energy loss evaluation.""" - self.loss_kwargs.update(self._pick_loss_kwargs(kwargs)) - self.lightcones = self._build_chunks_if_requested( - self.state, - self.terms, - schedule=self.schedule, - ) - return self - - def loss(self, state=None, *, hamiltonian=None, terms=None, **kwargs): - """Evaluate the configured MERA local energy loss.""" - state = self.state if state is None else self._as_mera_state(state) - terms_use = self.terms - chunks = self.lightcones if state is self.state else None - if hamiltonian is not None: - terms_use = normalize_local_terms(hamiltonian) - chunks = None - if terms is not None: - terms_use = normalize_local_terms(terms) - chunks = None - opts = self._merge_opts(self.loss_kwargs, self._pick_loss_kwargs(kwargs)) - if chunks is not None and ( - opts.get("array_backend") is not self.loss_kwargs.get("array_backend") - or opts.get("convert_terms") != self.loss_kwargs.get("convert_terms") - ): - chunks = None - return self._loss_state( - state, - terms=terms_use, - chunks=chunks, - schedule=self.schedule, - **opts, - ) - - def energy(self, state=None, *, hamiltonian=None, terms=None, **kwargs): - """Return full and per-site MERA energy estimates.""" - state = self.state if state is None else self._as_mera_state(state) - opts = self._merge_opts(self.loss_kwargs, self._pick_loss_kwargs(kwargs)) - opts_full = dict(opts) - opts_full["energy_per_site"] = False - energy = self.loss( - state, - hamiltonian=hamiltonian, - terms=terms, - **opts_full, - ) - num_sites = self._num_sites(state) - energy_per_site = energy / num_sites - chunks = self.lightcones if state is self.state else None - metadata = { - "real": opts["real"], - "contraction_opt": opts["contraction_opt"], - "simplify": opts["simplify"], - "gate_contract": opts["gate_contract"], - "num_terms": len(self.terms), - } - if chunks is not None: - metadata.update(self.lightcone_diagnostics(chunks=chunks)) - return EnergyEstimate( - energy=energy, - energy_per_site=energy_per_site, - num_sites=num_sites, - chi=self._max_bond(state), - boundary_mode="lightcone-exact", - normalized=bool(opts["normalized"]), - metadata=metadata, - ) - - def lightcone_diagnostics(self, *, chunks=None): - """Return compact diagnostics for cached lightcone chunks.""" - chunks = self.lightcones if chunks is None else chunks - if chunks is None: - chunks = self._build_chunks_if_requested(self.state, self.terms) - if not chunks: - return { - "max_lightcone_tensors": 0, - "max_lightcone_indices": 0, - "max_physical_width": 0, - } - return { - "max_lightcone_tensors": max(chunk.num_tensors for chunk in chunks), - "max_lightcone_indices": max(chunk.num_indices for chunk in chunks), - "max_physical_width": max(chunk.physical_width for chunk in chunks), - "max_schedule_width": max(chunk.schedule_width for chunk in chunks), - "lightcone_sources": tuple(chunk.source for chunk in chunks), - "num_tensors_by_term": tuple(chunk.num_tensors for chunk in chunks), - "num_indices_by_term": tuple(chunk.num_indices for chunk in chunks), - "physical_width_by_term": tuple(chunk.physical_width for chunk in chunks), - "schedule_width_by_term": tuple(chunk.schedule_width for chunk in chunks), - } - - def _norm_fn(self, state): - method = self.isometrize_method - if method is None: - return state - isometrize = getattr(state, "isometrize", None) - if not callable(isometrize): - return state - return isometrize(method=method, inplace=False) - - def make_tn_optimizer( - self, - *, - loss_kwargs: Mapping[str, Any] | None = None, - loss_constants: Mapping[str, Any] | None = None, - autodiff_backend: str = "torch", - optimizer: str = "adam", - progbar: bool = True, - device: str = "cpu", - **tnopt_kwargs, - ): - """Construct a configured :class:`quimb.tensor.TNOptimizer`.""" - del device - merged_loss_kwargs = self._merge_opts( - self.loss_kwargs, - self._pick_loss_kwargs(loss_kwargs), - ) - optimizer = GlobalOptimizer._normalize_optimizer_name(optimizer) - self._prepare_autodiff_backend(autodiff_backend) - incoming_constants = dict(loss_constants or {}) - terms = incoming_constants.pop("terms", self.terms) - chunks = incoming_constants.pop( - "chunks", - self.lightcones if merged_loss_kwargs.get("precompute_tags", True) else None, - ) - constants = {"terms": terms, "chunks": chunks, "schedule": self.schedule} - constants.update(incoming_constants) - return qtn.TNOptimizer( - self.state, - self._tnopt_loss, - norm_fn=self._norm_fn if self.isometrize_method is not None else None, - loss_constants=constants, - loss_kwargs=merged_loss_kwargs, - autodiff_backend=autodiff_backend, - optimizer=optimizer, - progbar=progbar, - **tnopt_kwargs, - ) - - def optimize( - self, - *, - n=220, - loss_kwargs: Mapping[str, Any] | None = None, - loss_constants: Mapping[str, Any] | None = None, - autodiff_backend: str = "torch", - optimizer: str = "adam", - progbar: bool = True, - return_losses: bool = False, - **optimize_kwargs, - ): - """Run ``TNOptimizer.optimize`` and store the optimized MERA state.""" - tnopt = self.make_tn_optimizer( - loss_kwargs=loss_kwargs, - loss_constants=loss_constants, - autodiff_backend=autodiff_backend, - optimizer=optimizer, - progbar=progbar, - ) - out = tnopt.optimize(n=n, **optimize_kwargs) - self.losses = list(getattr(tnopt, "losses", ())) - self.state = out - self.lightcones = self._build_chunks_if_requested( - self.state, - self.terms, - schedule=self.schedule, - ) - if return_losses: - return out, tuple(self.losses) - return out diff --git a/src/pepsy/optimizers/qmera/__init__.py b/src/pepsy/optimizers/qmera/__init__.py new file mode 100644 index 0000000..b76538b --- /dev/null +++ b/src/pepsy/optimizers/qmera/__init__.py @@ -0,0 +1,132 @@ +"""Canonical qMERA energy optimization helpers.""" + +from .builders import QMeraAnsatz, QMeraBuilder +from .cache import QMeraContractionPathCache, build_qmera_contraction_optimizer +from .compiled import ( + QMeraCompiledLightconeChunk, + compile_qmera_parametric_lightcone, + compile_qmera_parametric_lightcones, + local_qmera_compiled_lightcone_expectation, + qmera_compiled_parametric_energy, +) +from .fermions import ( + QMeraSymmrayFermionBackend, + qmera_symmray_fermi_hubbard_terms, + qmera_symmray_majorana_terms, + symmray_fermion_gate_registry, + symmray_majorana_gate_registry, +) +from .gates import ( + GateRegistry, + GateSpec, + UserGateFamily, + default_gate_registry, + resolve_gate_spec, +) +from .geometry import QMeraGeometry +from .layout import ( + QMeraLayoutCandidate, + QMeraLayoutFinder, + QMeraLayoutReport, + QMeraLayoutScore, +) +from .lightcones import ( + QMeraLightconeTN, + QMeraLightconeGroup, + QMeraParametricLightconeChunk, + build_qmera_lightcone_chunks, + build_qmera_parametric_lightcone_chunks, + contract_qmera_lightcone_group, + contract_qmera_lightcone_tn, + group_qmera_parametric_lightcone_chunks, + local_qmera_parametric_lightcone_expectation, + qmera_parametric_energy, + qmera_direct_parametric_energy, + qmera_parametric_lightcone_group_state, + qmera_parametric_state, + qmera_parametric_lightcone_state, + qmera_parametric_lightcone_tn, + select_lightcone, + site_tags_for_where, +) +from .parametric import QMeraEnergyOptimizer, QMeraParametricEnergyOptimizer +from .prototype import QMeraPrototypeLayout, load_qmera_prototype_layout +from .schedules import ( + QMeraBlockSpec, + QMeraDisentanglerSpec, + QMeraGatePlacement, + QMeraIsometrySpec, + QMeraLayerSpec, + QMeraScaleSpec, + QMeraSchedule, + QMeraUnitarySpec, + build_qmera_schedule, +) +from .schematics import ( + QMeraSchematicBlock, + draw_qmera_schedule, + qmera_schematic_blocks, +) +from .terms import LocalTerm, normalize_local_terms + +__all__ = [ + "GateRegistry", + "GateSpec", + "LocalTerm", + "QMeraAnsatz", + "QMeraBlockSpec", + "QMeraBuilder", + "QMeraCompiledLightconeChunk", + "QMeraContractionPathCache", + "QMeraDisentanglerSpec", + "QMeraEnergyOptimizer", + "QMeraGatePlacement", + "QMeraGeometry", + "QMeraIsometrySpec", + "QMeraLayerSpec", + "QMeraLayoutCandidate", + "QMeraLayoutFinder", + "QMeraLayoutReport", + "QMeraLayoutScore", + "QMeraLightconeTN", + "QMeraLightconeGroup", + "QMeraParametricLightconeChunk", + "QMeraParametricEnergyOptimizer", + "QMeraPrototypeLayout", + "QMeraSchedule", + "QMeraScaleSpec", + "QMeraSchematicBlock", + "QMeraSymmrayFermionBackend", + "QMeraUnitarySpec", + "UserGateFamily", + "build_qmera_contraction_optimizer", + "build_qmera_lightcone_chunks", + "build_qmera_parametric_lightcone_chunks", + "build_qmera_schedule", + "compile_qmera_parametric_lightcone", + "compile_qmera_parametric_lightcones", + "contract_qmera_lightcone_tn", + "contract_qmera_lightcone_group", + "default_gate_registry", + "draw_qmera_schedule", + "group_qmera_parametric_lightcone_chunks", + "local_qmera_compiled_lightcone_expectation", + "local_qmera_parametric_lightcone_expectation", + "load_qmera_prototype_layout", + "normalize_local_terms", + "qmera_compiled_parametric_energy", + "qmera_direct_parametric_energy", + "qmera_parametric_energy", + "qmera_parametric_lightcone_group_state", + "qmera_parametric_lightcone_state", + "qmera_parametric_state", + "qmera_parametric_lightcone_tn", + "qmera_schematic_blocks", + "qmera_symmray_fermi_hubbard_terms", + "qmera_symmray_majorana_terms", + "resolve_gate_spec", + "select_lightcone", + "site_tags_for_where", + "symmray_fermion_gate_registry", + "symmray_majorana_gate_registry", +] diff --git a/src/pepsy/optimizers/mera/builders.py b/src/pepsy/optimizers/qmera/builders.py similarity index 95% rename from src/pepsy/optimizers/mera/builders.py rename to src/pepsy/optimizers/qmera/builders.py index 22f6dc6..ac63619 100644 --- a/src/pepsy/optimizers/mera/builders.py +++ b/src/pepsy/optimizers/qmera/builders.py @@ -78,13 +78,17 @@ def reverse_lightcone_tags(self, where): """Return schedule tags in the reverse lightcone of ``where``.""" return self.schedule.reverse_lightcone_tags(where) - def schematic_blocks(self, *, layer=None): - """Return display-oriented disentangler/isometry blocks.""" - return self.schedule.schematic_blocks(layer=layer) + def schematic_blocks(self, *, layer=None, rg_step=None): + """Return display blocks for one or more RG steps.""" + return self.schedule.schematic_blocks(layer=layer, rg_step=rg_step) - def draw_schematic(self, *, layer=None, **kwargs): + def draw_schematic(self, *, layer=None, rg_step=None, **kwargs): """Draw qMERA blocking for this ansatz.""" - return self.schedule.draw_schematic(layer=layer, **kwargs) + return self.schedule.draw_schematic( + layer=layer, + rg_step=rg_step, + **kwargs, + ) def _coerce_geometry( @@ -344,13 +348,20 @@ def build_schedule(self): self._validate_unitary_spec(scale.isometry) return schedule - def schematic_blocks(self, *, layer=None): - """Return display-oriented disentangler/isometry blocks.""" - return self.build_schedule().schematic_blocks(layer=layer) + def schematic_blocks(self, *, layer=None, rg_step=None): + """Return display blocks for one or more RG steps.""" + return self.build_schedule().schematic_blocks( + layer=layer, + rg_step=rg_step, + ) - def draw_schematic(self, *, layer=None, **kwargs): + def draw_schematic(self, *, layer=None, rg_step=None, **kwargs): """Draw the qMERA blocking implied by this builder.""" - return self.build_schedule().draw_schematic(layer=layer, **kwargs) + return self.build_schedule().draw_schematic( + layer=layer, + rg_step=rg_step, + **kwargs, + ) def contraction_optimizer(self, **kwargs): """Build a reusable contraction optimizer for repeated local cones.""" @@ -717,6 +728,7 @@ def compile_parametric_lightcones( convert_terms=True, contraction_opt="auto-hq", expression_opts=None, + path_cache=None, ): """Compile static contraction expressions for qMERA local cones.""" schedule = self.build_schedule() if schedule is None else schedule @@ -740,6 +752,8 @@ def compile_parametric_lightcones( physical_dim=self.physical_dim, optimize=contraction_opt, expression_opts=expression_opts, + product_state_factory=self.product_state_factory, + path_cache=path_cache, ) def compiled_parametric_loss( @@ -758,6 +772,7 @@ def compiled_parametric_loss( real=True, contraction_opt="auto-hq", expression_opts=None, + path_cache=None, ): """Evaluate qMERA energy with precompiled local-cone contractions.""" schedule = self.build_schedule() if schedule is None else schedule @@ -778,6 +793,8 @@ def compiled_parametric_loss( energy_per_site=energy_per_site, real=real, expression_opts=expression_opts, + product_state_factory=self.product_state_factory, + path_cache=path_cache, ) def compiled_parametric_loss_fn( @@ -813,8 +830,14 @@ def parametric_optimizer( parameters=None, **loss_kwargs, ): - """Create a parameter-dict qMERA energy optimizer shell.""" - from .parametric import QMeraParametricEnergyOptimizer + """Create the parameterized qMERA energy optimizer. + + Built-in qMERA gates are unitary by construction, so the optimizer + uses the numerator expectation value directly by default. Set + ``normalized=True`` in ``loss_kwargs`` for a custom non-unitary gate + family. + """ + from .parametric import QMeraEnergyOptimizer schedule = self.build_schedule() if schedule is None else schedule if chunks is None: @@ -826,7 +849,9 @@ def parametric_optimizer( ) if parameters is None: parameters = self.initialize_parameters(schedule) - return QMeraParametricEnergyOptimizer( + loss_kwargs = dict(loss_kwargs) + loss_kwargs.setdefault("normalized", False) + return QMeraEnergyOptimizer( builder=self, schedule=schedule, hamiltonian=hamiltonian, diff --git a/src/pepsy/optimizers/mera/cache.py b/src/pepsy/optimizers/qmera/cache.py similarity index 97% rename from src/pepsy/optimizers/mera/cache.py rename to src/pepsy/optimizers/qmera/cache.py index c497c6c..bf0e3d7 100644 --- a/src/pepsy/optimizers/mera/cache.py +++ b/src/pepsy/optimizers/qmera/cache.py @@ -1,4 +1,4 @@ -"""Reusable contraction optimizers for MERA and qMERA local cones.""" +"""Reusable contraction optimizers for qMERA local cones.""" from __future__ import annotations diff --git a/src/pepsy/optimizers/mera/compiled.py b/src/pepsy/optimizers/qmera/compiled.py similarity index 52% rename from src/pepsy/optimizers/mera/compiled.py rename to src/pepsy/optimizers/qmera/compiled.py index 21926de..28d9d34 100644 --- a/src/pepsy/optimizers/mera/compiled.py +++ b/src/pepsy/optimizers/qmera/compiled.py @@ -13,7 +13,7 @@ QMeraParametricLightconeChunk, _gate_for_placement, _maybe_real, - _product_state_on_sites, + _product_state_for_schedule, _site_ind, build_qmera_parametric_lightcone_chunks, ) @@ -31,6 +31,72 @@ _KET_TAG = "_QMERA_KET_COPY" +def _is_native_symmray_spec(spec): + """Return whether ``spec`` builds graded Symmray tensors.""" + return ( + str(getattr(spec, "name", "")).lower().startswith("symmray-") + or str(getattr(spec, "convention", "")).lower().startswith("symmray-") + ) + + +def _is_native_fermionic_operator(operator): + """Return whether an operator payload is a Symmray graded array.""" + return bool(getattr(operator, "fermionic", False)) or ( + "fermionicarray" in type(operator).__name__.lower() + ) + + +def _is_native_symmray_array(value): + """Return whether ``value`` is a native graded Symmray array. + + The check intentionally uses the public Symmray surface rather than an + exact class name so that both block-sparse and flat fermionic arrays are + accepted across Symmray versions. + """ + return bool( + getattr(value, "fermionic", False) + and callable(getattr(value, "tensordot", None)) + and callable(getattr(value, "transpose", None)) + and hasattr(value, "duals") + and hasattr(value, "indices") + ) + + +def _native_symmray_requested(schedule, chunks, gate_registry): + """Return whether ``chunks`` require the graded Symmray route.""" + placements = schedule.placements_by_id() + return any( + _is_native_symmray_spec( + gate_registry.get(placements[gate_id].gate_family) + ) + for chunk in chunks + for gate_id in chunk.schedule_placement_ids + ) or any( + _is_native_fermionic_operator(chunk.term.operator) + for chunk in chunks + ) + + +def _validate_native_symmray_compile( + schedule, + chunks, + gate_registry, + *, + product_state_factory=None, +): + """Validate that native compilation has a graded product-state source.""" + native_requested = _native_symmray_requested(schedule, chunks, gate_registry) + if native_requested and product_state_factory is None: + raise ValueError( + "Native Symmray qMERA compilation requires a graded " + "product_state_factory (for example, " + "QMeraSymmrayFermionBackend.product_state) so that the frozen " + "contraction constants retain charge maps, duals, and fermionic " + "ordering." + ) + return native_requested + + @dataclass(frozen=True) class QMeraCompiledLightconeChunk: """Static contraction expressions for one qMERA local cone.""" @@ -43,6 +109,9 @@ class QMeraCompiledLightconeChunk: num_numerator_tensors: int num_denominator_tensors: int optimize: Any = "auto-hq" + contraction_backend: str = "array" + symmetry: str | None = None + fermionic: bool = False @property def schedule_placement_ids(self): @@ -54,6 +123,11 @@ def num_gates(self): """Number of parametrized gates used by this compiled local cone.""" return self.chunk.num_gates + @property + def is_graded(self): + """Whether this expression uses native graded Symmray arrays.""" + return self.fermionic + def _gate_tag_to_id(tag): text = str(tag) @@ -82,6 +156,23 @@ def _copy_with_tag(tn, tag): return out +def _expression_topology_key(inputs, output, shapes): + """Canonical, hashable topology key for a local contraction expression.""" + labels = {} + + def canonical(label): + if label not in labels: + labels[label] = len(labels) + return labels[label] + + canonical_inputs = tuple( + tuple(canonical(label) for label in tensor_inputs) + for tensor_inputs in inputs + ) + canonical_output = tuple(canonical(label) for label in output) + return canonical_inputs, canonical_output, tuple(tuple(shape) for shape in shapes) + + def _dummy_params(schedule, chunk, gate_registry): placements = schedule.placements_by_id() params = {} @@ -99,12 +190,15 @@ def _static_lightcone_state( gate_registry, array_backend, physical_dim, + product_state_factory=None, ): placements = schedule.placements_by_id() - state = _product_state_on_sites( + state = _product_state_for_schedule( + schedule, chunk.input_sites, physical_dim=physical_dim, array_backend=array_backend, + product_state_factory=product_state_factory, ) params = _dummy_params(schedule, chunk, gate_registry) for gate_id in chunk.schedule_placement_ids: @@ -126,7 +220,7 @@ def _static_lightcone_state( return state -def _expression_from_tn(tn, *, optimize, expression_opts=None): +def _expression_from_tn(tn, *, optimize, expression_opts=None, path_cache=None): inputs = [] shapes = [] constants = {} @@ -149,6 +243,42 @@ def _expression_from_tn(tn, *, optimize, expression_opts=None): f"Could not identify bra/ket copy for qMERA gate {gate_id!r}." ) opts = {} if expression_opts is None else dict(expression_opts) + if path_cache is not None: + optimize = path_cache.resolve( + optimize, + key=_expression_topology_key(inputs, (), shapes), + ) + if any(_is_native_symmray_array(tensor.data) for tensor in tensors): + if opts.get("implementation") is not None: + raise ValueError( + "Native Symmray qMERA compilation must leave the cotengra " + "implementation unset so Symmray's graded autoray dispatch " + "is selected from the runtime arrays." + ) + # Leave the implementation options unset. Cotengra's default expression + # builder traces the static expression and dispatches each runtime + # pairwise contraction through the backend inferred from the native + # Symmray operands. Forcing an implementation at trace time would feed + # NumPy lazy placeholders into the constants folder and lose the graded + # array object before evaluation. + if not slots: + # A local term can have an empty reverse lightcone (for example an + # onsite operator on a schedule with no active gate). Cotengra's + # constants-folding helper expects at least one lazy input, so + # evaluate this immutable scalar once and expose the same zero-arg + # callable interface as a dynamic expression. + static_value = ctg.array_contract( + tuple(tensor.data for tensor in tensors), + inputs, + output=(), + optimize=optimize, + **opts, + ) + + def static_expression(*_arrays): + return static_value + + return static_expression, (), len(tensors) expr = ctg.array_contract_expression( inputs, output=(), @@ -160,13 +290,22 @@ def _expression_from_tn(tn, *, optimize, expression_opts=None): return expr, tuple(slots), len(tensors) -def _compiled_tns(schedule, chunk, *, gate_registry, array_backend, physical_dim): +def _compiled_tns( + schedule, + chunk, + *, + gate_registry, + array_backend, + physical_dim, + product_state_factory=None, +): ket = _static_lightcone_state( schedule, chunk, gate_registry=gate_registry, array_backend=array_backend, physical_dim=physical_dim, + product_state_factory=product_state_factory, ) bra = _copy_with_tag(ket.H, _BRA_TAG) ket_side = _copy_with_tag(ket, _KET_TAG) @@ -180,6 +319,27 @@ def _compiled_tns(schedule, chunk, *, gate_registry, array_backend, physical_dim return numerator, denominator +def _native_tn_metadata(tn): + """Return ``(fermionic, symmetry)`` metadata for a compiled TN.""" + arrays = tuple(tensor.data for tensor in tn) + native = tuple(value for value in arrays if _is_native_symmray_array(value)) + if not native: + return False, None + if len(native) != len(arrays): + raise TypeError( + "Native Symmray qMERA compilation requires every tensor in the " + "frozen local cone to be a graded Symmray array; a dense tensor " + "would drop fermionic signs or charge-sector metadata." + ) + symmetries = {str(getattr(value, "symmetry", "")) for value in native} + if len(symmetries) != 1: + raise ValueError( + "Native Symmray qMERA lightcones must use one compatible symmetry; " + f"found {sorted(symmetries)!r}." + ) + return True, next(iter(symmetries)) + + def compile_qmera_parametric_lightcone( schedule, chunk: QMeraParametricLightconeChunk, @@ -189,25 +349,47 @@ def compile_qmera_parametric_lightcone( physical_dim=2, optimize="auto-hq", expression_opts=None, + product_state_factory=None, + path_cache=None, ): - """Compile static numerator and denominator contractions for ``chunk``.""" + """Compile static numerator and denominator contractions for ``chunk``. + + Native Symmray chunks are compiled as graded-array expressions. Their + static product state and local operator remain Symmray objects, while + cotengra freezes only the contraction topology. + """ gate_registry = default_gate_registry() if gate_registry is None else gate_registry + native_requested = _validate_native_symmray_compile( + schedule, + (chunk,), + gate_registry, + product_state_factory=product_state_factory, + ) numerator, denominator = _compiled_tns( schedule, chunk, gate_registry=gate_registry, array_backend=array_backend, physical_dim=physical_dim, + product_state_factory=product_state_factory, ) + fermionic, symmetry = _native_tn_metadata(numerator) + if native_requested and not fermionic: + raise TypeError( + "Native Symmray qMERA compilation did not produce a fully graded " + "frozen local cone. Check product_state_factory and gate registry." + ) numerator_expr, numerator_slots, num_num_tensors = _expression_from_tn( numerator, optimize=optimize, expression_opts=expression_opts, + path_cache=path_cache, ) denominator_expr, denominator_slots, num_den_tensors = _expression_from_tn( denominator, optimize=optimize, expression_opts=expression_opts, + path_cache=path_cache, ) return QMeraCompiledLightconeChunk( chunk=chunk, @@ -218,6 +400,9 @@ def compile_qmera_parametric_lightcone( num_numerator_tensors=num_num_tensors, num_denominator_tensors=num_den_tensors, optimize=optimize, + contraction_backend="symmray" if fermionic else "array", + symmetry=symmetry, + fermionic=fermionic, ) @@ -230,8 +415,22 @@ def compile_qmera_parametric_lightcones( physical_dim=2, optimize="auto-hq", expression_opts=None, + product_state_factory=None, + path_cache=None, ): - """Compile every qMERA local cone in ``chunks``.""" + """Compile every qMERA local cone in ``chunks``. + + The compiled expressions are static and can be evaluated repeatedly with + new native Symmray gate arrays without rebuilding the qMERA lightcones. + """ + gate_registry = default_gate_registry() if gate_registry is None else gate_registry + chunks = tuple(chunks) + _validate_native_symmray_compile( + schedule, + chunks, + gate_registry, + product_state_factory=product_state_factory, + ) return tuple( compile_qmera_parametric_lightcone( schedule, @@ -241,6 +440,8 @@ def compile_qmera_parametric_lightcones( physical_dim=physical_dim, optimize=optimize, expression_opts=expression_opts, + product_state_factory=product_state_factory, + path_cache=path_cache, ) for chunk in chunks ) @@ -269,6 +470,26 @@ def _arrays_for_slots(slots, gates): return tuple(arrays) +def _validate_compiled_gate_arrays(compiled, gates): + """Ensure a graded expression is evaluated with graded gate arrays.""" + if not compiled.fermionic: + return + if not all(_is_native_symmray_array(gate) for gate in gates.values()): + raise TypeError( + "This compiled qMERA lightcone was built for native Symmray " + "fermionic gates, but evaluation received a dense gate array. " + "Use the same Symmray gate registry used during compilation." + ) + if compiled.symmetry is not None and gates: + found = {str(getattr(gate, "symmetry", "")) for gate in gates.values()} + if found != {compiled.symmetry}: + raise ValueError( + "Compiled Symmray qMERA gate symmetry does not match the " + f"frozen lightcone: expected {compiled.symmetry!r}, found " + f"{sorted(found)!r}." + ) + + def local_qmera_compiled_lightcone_expectation( schedule, compiled: QMeraCompiledLightconeChunk, @@ -288,6 +509,7 @@ def local_qmera_compiled_lightcone_expectation( gate_registry=gate_registry, gate_array_backend=gate_array_backend, ) + _validate_compiled_gate_arrays(compiled, gates) numerator = compiled.numerator_expr( *_arrays_for_slots(compiled.numerator_slots, gates) ) @@ -321,9 +543,20 @@ def qmera_compiled_parametric_energy( energy_per_site=True, real=True, expression_opts=None, + product_state_factory=None, + path_cache=None, ): """Evaluate qMERA energy from precompiled local-cone expressions.""" gate_registry = default_gate_registry() if gate_registry is None else gate_registry + if compiled_chunks is not None: + chunks_for_guard = tuple(chunks or ()) + if chunks_for_guard: + _validate_native_symmray_compile( + schedule, + chunks_for_guard, + gate_registry, + product_state_factory=product_state_factory, + ) if compiled_chunks is None: if chunks is None: if hamiltonian is None: @@ -335,6 +568,13 @@ def qmera_compiled_parametric_energy( if convert_terms: terms = convert_local_terms(terms, array_backend) chunks = build_qmera_parametric_lightcone_chunks(schedule, terms) + chunks = tuple(chunks) + _validate_native_symmray_compile( + schedule, + chunks, + gate_registry, + product_state_factory=product_state_factory, + ) compiled_chunks = compile_qmera_parametric_lightcones( schedule, chunks, @@ -343,6 +583,8 @@ def qmera_compiled_parametric_energy( physical_dim=physical_dim, optimize=optimize, expression_opts=expression_opts, + product_state_factory=product_state_factory, + path_cache=path_cache, ) value = None for compiled in compiled_chunks: diff --git a/src/pepsy/optimizers/mera/fermions.py b/src/pepsy/optimizers/qmera/fermions.py similarity index 100% rename from src/pepsy/optimizers/mera/fermions.py rename to src/pepsy/optimizers/qmera/fermions.py diff --git a/src/pepsy/optimizers/mera/gates.py b/src/pepsy/optimizers/qmera/gates.py similarity index 100% rename from src/pepsy/optimizers/mera/gates.py rename to src/pepsy/optimizers/qmera/gates.py diff --git a/src/pepsy/optimizers/mera/geometry.py b/src/pepsy/optimizers/qmera/geometry.py similarity index 100% rename from src/pepsy/optimizers/mera/geometry.py rename to src/pepsy/optimizers/qmera/geometry.py diff --git a/src/pepsy/optimizers/qmera/layout.py b/src/pepsy/optimizers/qmera/layout.py new file mode 100644 index 0000000..fcba40d --- /dev/null +++ b/src/pepsy/optimizers/qmera/layout.py @@ -0,0 +1,451 @@ +"""Architecture search for schedule-first qMERA RG layouts.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +import hashlib +from itertools import product +from typing import Any + +from .builders import QMeraBuilder +from .schedules import QMeraScaleSpec +from .terms import normalize_local_terms + +__all__ = [ + "QMeraLayoutCandidate", + "QMeraLayoutFinder", + "QMeraLayoutReport", + "QMeraLayoutScore", +] + + +def _freeze_mapping(value): + return dict(value or {}) + + +def _stable_id(config): + payload = repr(tuple(sorted(config.items(), key=lambda item: item[0]))) + digest = hashlib.sha1(payload.encode("utf-8")).hexdigest()[:10] + return f"qmera-layout-{digest}" + + +def _as_options(value, default): + if value is None: + return (default,) + if isinstance(value, (str, bytes)): + return (value,) + if isinstance(value, int): + return (value,) + if ( + isinstance(default, tuple) + and isinstance(value, (tuple, list)) + and len(value) == len(default) + and all(isinstance(item, (int, float)) for item in value) + ): + # In 2D, ``block_shapes=(2, 2)`` is the natural spelling for one + # rectangular shape. Multiple shapes remain explicit as + # ``((2, 2), (3, 3))``. + return (tuple(value),) + try: + values = tuple(value) + except TypeError: + return (value,) + return values or (default,) + + +def _shape_key(value): + if isinstance(value, (tuple, list)): + return tuple(int(x) for x in value) + return int(value) + + +def _candidate_config_repr(config): + return { + "disentangler_block_shape": _shape_key(config["disentangler_block_shape"]), + "isometry_block_shape": _shape_key(config["isometry_block_shape"]), + "disentangler_depth": int(config["disentangler_depth"]), + "isometry_depth": int(config["isometry_depth"]), + "num_scales": int(config["num_scales"]), + } + + +@dataclass(frozen=True) +class QMeraLayoutCandidate: + """One immutable qMERA RG architecture candidate.""" + + candidate_id: str + scales: tuple[QMeraScaleSpec, ...] + metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self): + object.__setattr__(self, "candidate_id", str(self.candidate_id)) + object.__setattr__(self, "scales", tuple(self.scales)) + object.__setattr__(self, "metadata", _freeze_mapping(self.metadata)) + + +@dataclass(frozen=True) +class QMeraLayoutScore: + """Comparable structural score for one qMERA layout candidate.""" + + candidate_id: str + total: float + components: Mapping[str, float] + valid: bool = True + error: str | None = None + + def __post_init__(self): + object.__setattr__(self, "candidate_id", str(self.candidate_id)) + object.__setattr__( + self, + "components", + {str(key): float(value) for key, value in dict(self.components).items()}, + ) + object.__setattr__(self, "total", float(self.total)) + + +@dataclass(frozen=True) +class QMeraLayoutReport: + """Search result containing scores and the non-dominated front.""" + + candidates: tuple[QMeraLayoutCandidate, ...] + scores: tuple[QMeraLayoutScore, ...] + pareto_front: tuple[str, ...] + + @property + def best(self): + """Return the lowest-total valid candidate, or ``None``.""" + valid = [score for score in self.scores if score.valid] + if not valid: + return None + best_id = min(valid, key=lambda score: score.total).candidate_id + return next( + candidate + for candidate in self.candidates + if candidate.candidate_id == best_id + ) + + @property + def best_score(self): + """Return the score associated with :attr:`best`, or ``None``.""" + best = self.best + if best is None: + return None + return next( + score for score in self.scores if score.candidate_id == best.candidate_id + ) + + +class QMeraLayoutFinder: + """Generate and rank valid qMERA RG architectures. + + The finder searches immutable scale plans. It never mutates a builder or + changes geometry/mode order while scoring a candidate. Scores are cheap + structural/contraction proxies intended for pre-ranking; a caller can use + the returned candidate with :class:`QMeraBuilder` for a measured pilot. + """ + + def __init__( + self, + geometry, + *, + gate_family="rxx", + isometry_gate_family=None, + gate_registry=None, + isometry_block_shapes=None, + disentangler_block_shapes=None, + disentangler_depths=None, + isometry_depths=None, + max_layers=None, + top_size=1, + weights=None, + builder_options=None, + ): + self.geometry = geometry + self.gate_family = gate_family + self.isometry_gate_family = isometry_gate_family or gate_family + self.gate_registry = gate_registry + default_shape = (2, 2) if getattr(geometry, "ndim", 1) == 2 else 2 + self.isometry_block_shapes = _as_options( + isometry_block_shapes, + default_shape, + ) + self.disentangler_block_shapes = _as_options( + disentangler_block_shapes, + default_shape, + ) + self.disentangler_depths = tuple( + int(value) for value in _as_options(disentangler_depths, 1) + ) + self.isometry_depths = tuple( + int(value) for value in _as_options(isometry_depths, 1) + ) + self.max_layers = None if max_layers is None else int(max_layers) + self.top_size = int(top_size) + if self.top_size < 1: + raise ValueError("top_size must be >= 1.") + defaults = { + "structural": 1.0, + "contraction": 0.1, + "coverage": 1.0, + } + if weights: + defaults.update({str(key): float(value) for key, value in weights.items()}) + self.weights = defaults + self.builder_options = dict(builder_options or {}) + + def _candidate_configs(self): + for values in product( + self.disentangler_block_shapes, + self.isometry_block_shapes, + self.disentangler_depths, + self.isometry_depths, + ): + yield { + "disentangler_block_shape": values[0], + "isometry_block_shape": values[1], + "disentangler_depth": values[2], + "isometry_depth": values[3], + } + + def _scales_for_config(self, config): + """Find the shortest repeated scale plan that reaches ``top_size``.""" + limit = self.max_layers + if limit is None: + limit = max(1, sum(int(dim).bit_length() for dim in self.geometry.shape)) + for num_scales in range(1, limit + 1): + scales = tuple( + QMeraScaleSpec( + name=f"candidate-scale-{scale}", + disentangler={ + "block_size": config["disentangler_block_shape"], + "circuit_depth": config["disentangler_depth"], + "gate_family": self.gate_family, + }, + isometry={ + "block_size": config["isometry_block_shape"], + "circuit_depth": config["isometry_depth"], + "gate_family": self.isometry_gate_family, + }, + ) + for scale in range(num_scales) + ) + try: + self._build_schedule(scales) + except (ValueError, IndexError, NotImplementedError): + continue + return scales + return None + + def _builder(self, scales): + options = dict(self.builder_options) + options.update( + geometry=self.geometry, + gate_family=self.gate_family, + isometry_gate_family=self.isometry_gate_family, + gate_registry=self.gate_registry, + scales=scales, + top_size=self.top_size, + ) + return QMeraBuilder(**options) + + def _build_schedule(self, scales): + return self._builder(scales).build_schedule() + + def generate_candidates(self): + """Generate valid, deterministically ordered layout candidates.""" + candidates = [] + for config in self._candidate_configs(): + scales = self._scales_for_config(config) + if scales is None: + continue + metadata = _candidate_config_repr( + {**config, "num_scales": len(scales)} + ) + candidates.append( + QMeraLayoutCandidate( + candidate_id=_stable_id(metadata), + scales=scales, + metadata=metadata, + ) + ) + return tuple(candidates) + + @staticmethod + def _lightcone_metrics(schedule, terms): + widths = [] + total_weight = 0.0 + covered_weight = 0.0 + for term in terms: + weight = abs(float(term.weight)) + total_weight += weight + support = set(schedule.geometry.to_register_where(term.where)) + initial_size = len(support) + selected = schedule.reverse_lightcone_placements(term.where) + for placement in selected: + support.update(placement.where) + widths.append(len(support)) + disentangler_support = { + site + for placement in selected + if placement.stage == "disentangler" + for site in placement.where + } + fraction = ( + len(disentangler_support.intersection(set( + schedule.geometry.to_register_where(term.where) + ))) + / max(1, initial_size) + ) + covered_weight += weight * fraction + if not widths: + widths = [0] + coverage = ( + covered_weight / total_weight + if total_weight + else 0.0 + ) + return { + "max_lightcone_width": float(max(widths)), + "mean_lightcone_width": float(sum(widths) / len(widths)), + "interaction_coverage": float(coverage), + } + + def score(self, candidate, hamiltonian=None): + """Score one candidate, returning an invalid score instead of raising.""" + try: + schedule = self._build_schedule(candidate.scales) + terms = () if hamiltonian is None else normalize_local_terms(hamiltonian) + metrics = self._lightcone_metrics(schedule, terms) + gate_count = float(schedule.num_gates) + circuit_depth = float( + sum( + max( + [placement.round for placement in layer.placements], + default=-1, + ) + + 1 + for layer in schedule.layers + ) + ) + structural = gate_count + circuit_depth + contraction = metrics["max_lightcone_width"] ** 3 + total = ( + self.weights["structural"] * structural + + self.weights["contraction"] * contraction + + self.weights["coverage"] + * (1.0 - metrics["interaction_coverage"]) + ) + components = { + "gate_count": gate_count, + "circuit_depth": circuit_depth, + "structural_cost": structural, + "contraction_cost_proxy": contraction, + **metrics, + } + return QMeraLayoutScore(candidate.candidate_id, total, components) + except (TypeError, ValueError, IndexError, KeyError, NotImplementedError) as exc: + return QMeraLayoutScore( + candidate.candidate_id, + float("inf"), + {}, + valid=False, + error=str(exc), + ) + + def score_prototype_layout(self, prototype, hamiltonian=None): + """Score a loaded prototype gate stream for structural comparison. + + Prototype streams are not converted into Pepsy schedules. This method + reports a separate stream-level score so users can compare placement + count, greedy parallel depth, and Hamiltonian-support coverage against + native qMERA candidates without mixing their contraction semantics. + """ + candidate_id = f"prototype:{prototype.name}" + try: + if int(prototype.num_sites) != int(self.geometry.num_modes): + raise ValueError( + "prototype num_sites must equal geometry.num_modes for a " + "direct register-order comparison." + ) + terms = () if hamiltonian is None else normalize_local_terms(hamiltonian) + pair_supports = {frozenset(pair) for pair in prototype.pairs} + total_weight = 0.0 + covered_weight = 0.0 + for term in terms: + weight = abs(float(term.weight)) + total_weight += weight + support = frozenset(self.geometry.to_register_where(term.where)) + if len(support) == 1: + covered = any(next(iter(support)) in pair for pair in pair_supports) + else: + covered = support in pair_supports + if covered: + covered_weight += weight + coverage = covered_weight / total_weight if total_weight else 0.0 + unique_fraction = len(prototype.unique_sites) / max(1, prototype.num_sites) + structural = float(prototype.gate_count + prototype.round_depth) + contraction = float(max(1, prototype.max_support) ** 3) + total = ( + self.weights["structural"] * structural + + self.weights["contraction"] * contraction + + self.weights["coverage"] * (1.0 - coverage) + ) + return QMeraLayoutScore( + candidate_id, + total, + { + "gate_count": float(prototype.gate_count), + "round_depth": float(prototype.round_depth), + "unique_site_fraction": float(unique_fraction), + "interaction_coverage": float(coverage), + "structural_cost": structural, + "contraction_cost_proxy": contraction, + }, + ) + except (TypeError, ValueError, KeyError) as exc: + return QMeraLayoutScore( + candidate_id, + float("inf"), + {}, + valid=False, + error=str(exc), + ) + + @staticmethod + def _pareto_front(scores): + valid = [score for score in scores if score.valid] + front = [] + for score in valid: + dominated = False + for other in valid: + if other is score: + continue + other_components = other.components + components = score.components + no_worse = ( + other.total <= score.total + and other_components.get("interaction_coverage", 0.0) + >= components.get("interaction_coverage", 0.0) + ) + strictly_better = ( + other.total < score.total + or other_components.get("interaction_coverage", 0.0) + > components.get("interaction_coverage", 0.0) + ) + if no_worse and strictly_better: + dominated = True + break + if not dominated: + front.append(score.candidate_id) + return tuple(front) + + def search(self, hamiltonian=None): + """Generate candidates, score them, and return a Pareto report.""" + candidates = self.generate_candidates() + scores = tuple(self.score(candidate, hamiltonian) for candidate in candidates) + return QMeraLayoutReport( + candidates=candidates, + scores=scores, + pareto_front=self._pareto_front(scores), + ) diff --git a/src/pepsy/optimizers/mera/lightcones.py b/src/pepsy/optimizers/qmera/lightcones.py similarity index 98% rename from src/pepsy/optimizers/mera/lightcones.py rename to src/pepsy/optimizers/qmera/lightcones.py index a3792d6..bff7967 100644 --- a/src/pepsy/optimizers/mera/lightcones.py +++ b/src/pepsy/optimizers/qmera/lightcones.py @@ -1,4 +1,4 @@ -"""Reverse-lightcone selection and local expectation kernels for MERA states.""" +"""Reverse-lightcone selection and local qMERA expectation kernels.""" from __future__ import annotations @@ -13,22 +13,18 @@ from .terms import LocalTerm, convert_local_terms, normalize_local_terms __all__ = [ - "LightconeChunk", "QMeraLightconeGroup", "QMeraLightconeTN", "QMeraParametricLightconeChunk", - "build_lightcone_chunks", "build_qmera_lightcone_chunks", "build_qmera_parametric_lightcone_chunks", "group_qmera_parametric_lightcone_chunks", "contract_qmera_lightcone_tn", "contract_qmera_lightcone_group", - "lightcone_energy", "qmera_direct_parametric_energy", "qmera_parametric_state", "qmera_parametric_lightcone_group_state", "local_qmera_parametric_lightcone_expectation", - "local_lightcone_expectation", "qmera_parametric_energy", "qmera_parametric_lightcone_state", "qmera_parametric_lightcone_tn", @@ -293,7 +289,7 @@ def build_qmera_lightcone_chunks(state, schedule, terms, *, validate=True): """Precompute lightcone chunks from an explicit qMERA schedule. Unlike :func:`build_lightcone_chunks`, this follows - :class:`~pepsy.optimizers.mera.QMeraSchedule` placements first and only then + :class:`~pepsy.optimizers.qmera.QMeraSchedule` placements first and only then turns the selected sites/gates into tensor-network tags. This keeps local energy chunks tied to the designed RG blocks rather than to a generic tag query on an already-built network. @@ -975,7 +971,7 @@ def _lightcone_state_and_schedule(state, schedule): ): return candidate, schedule raise TypeError( - "state must be a MERA-like TensorNetwork with select() and gate(), " + "state must be a qMERA TensorNetwork with select() and gate(), " "or an ansatz object exposing .state." ) diff --git a/src/pepsy/optimizers/mera/parametric.py b/src/pepsy/optimizers/qmera/parametric.py similarity index 89% rename from src/pepsy/optimizers/mera/parametric.py rename to src/pepsy/optimizers/qmera/parametric.py index ae251b8..31bba9b 100644 --- a/src/pepsy/optimizers/mera/parametric.py +++ b/src/pepsy/optimizers/qmera/parametric.py @@ -9,7 +9,7 @@ from ...backends import backend_jax, backend_torch from ...solvers import GradientOptimizer, GradSolverResult -__all__ = ["QMeraParametricEnergyOptimizer"] +__all__ = ["QMeraEnergyOptimizer", "QMeraParametricEnergyOptimizer"] def _solver_backend(solver): @@ -40,8 +40,15 @@ def _array_backend_for_train_backend(backend, *, dtype=None, device="cpu"): @dataclass -class QMeraParametricEnergyOptimizer: - """Optimize qMERA parameters using schedule-only local lightcones.""" +class QMeraEnergyOptimizer: + """Optimize parameterized qMERA gates using schedule-only lightcones. + + The built-in qMERA gate families are unitary (and the fermion families + are symmetry/parity preserving) by construction. Consequently the + optimizer does not normalize the state by default. Pass + ``normalized=True`` in ``loss_kwargs`` or to ``loss``/``run`` when using + a custom gate family that is not norm preserving. + """ builder: Any schedule: Any @@ -55,6 +62,7 @@ class QMeraParametricEnergyOptimizer: def __post_init__(self): self.loss_kwargs = {} if self.loss_kwargs is None else dict(self.loss_kwargs) + self.loss_kwargs.setdefault("normalized", False) if self.parameters is None: self.parameters = self.builder.initialize_parameters(self.schedule) else: @@ -132,6 +140,7 @@ def compile( convert_terms=True, contraction_opt="auto-hq", expression_opts=None, + path_cache=None, ): """Compile static contraction expressions for configured local cones.""" chunks = self._chunks_for_backend( @@ -147,6 +156,7 @@ def compile( convert_terms=convert_terms, contraction_opt=contraction_opt, expression_opts=expression_opts, + path_cache=path_cache, ) return self.compiled_chunks @@ -162,6 +172,7 @@ def compiled_loss(self, parameters=None, **kwargs): convert_terms=opts.get("convert_terms", True), contraction_opt=opts.get("contraction_opt", "auto-hq"), expression_opts=opts.get("expression_opts"), + path_cache=opts.get("path_cache"), ) return self.builder.compiled_parametric_loss( params, @@ -230,6 +241,7 @@ def run( convert_terms=opts.get("convert_terms", True), contraction_opt=opts.get("contraction_opt", "auto-hq"), expression_opts=expression_opts, + path_cache=opts.get("path_cache"), ) opts["compiled_chunks"] = self.compiled_chunks @@ -253,3 +265,8 @@ def run( return result optimize = run + + +# Compatibility name retained for callers of the original parameter-dict +# qMERA API. The canonical public name is now QMeraEnergyOptimizer. +QMeraParametricEnergyOptimizer = QMeraEnergyOptimizer diff --git a/src/pepsy/optimizers/qmera/prototype.py b/src/pepsy/optimizers/qmera/prototype.py new file mode 100644 index 0000000..111543f --- /dev/null +++ b/src/pepsy/optimizers/qmera/prototype.py @@ -0,0 +1,153 @@ +"""Adapters for comparing qMERA schedules with the research prototype streams.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +import re +from typing import Any, Callable + +__all__ = [ + "QMeraPrototypeLayout", + "load_qmera_prototype_layout", +] + + +_LEVEL_RE = re.compile(r"(?:^|[_-])(?:l|level)[_-]?(\d+)(?:$|[_-])", re.I) + + +def _normalize_pair(pair, index): + try: + left, right = pair + except (TypeError, ValueError) as exc: + raise ValueError( + f"prototype gate entry {index} must contain exactly two sites." + ) from exc + if isinstance(left, bool) or isinstance(right, bool): + raise TypeError(f"prototype gate entry {index} contains a boolean site.") + try: + left, right = int(left), int(right) + except (TypeError, ValueError) as exc: + raise TypeError( + f"prototype gate entry {index} must contain integer site labels." + ) from exc + if left < 0 or right < 0: + raise ValueError(f"prototype gate entry {index} contains a negative site.") + if left == right: + raise ValueError(f"prototype gate entry {index} is a self-gate.") + return left, right + + +@dataclass(frozen=True) +class QMeraPrototypeLayout: + """A normalized two-site gate stream loaded from the qMERA prototype. + + The prototype ``U_q3_l*`` files are serialized placement streams rather + than Pepsy RG schedules. This object keeps them available for structural + comparison without pretending that their flat stream is a qMERA + ``QMeraScaleSpec``. + """ + + name: str + path: str + pairs: tuple[tuple[int, int], ...] + num_sites: int + level: int | None = None + source: str = "mera-prototype" + + def __post_init__(self): + object.__setattr__(self, "name", str(self.name)) + object.__setattr__(self, "path", str(self.path)) + pairs = tuple((int(left), int(right)) for left, right in self.pairs) + object.__setattr__(self, "pairs", pairs) + num_sites = int(self.num_sites) + if num_sites < 1: + raise ValueError("num_sites must be >= 1.") + if pairs and max(max(pair) for pair in pairs) >= num_sites: + raise ValueError("num_sites must cover every prototype gate site.") + object.__setattr__(self, "num_sites", num_sites) + object.__setattr__(self, "level", None if self.level is None else int(self.level)) + + @property + def gate_count(self): + """Return the number of two-site placements in the stream.""" + return len(self.pairs) + + @property + def unique_sites(self): + """Return the sites touched by at least one prototype placement.""" + return tuple(sorted({site for pair in self.pairs for site in pair})) + + @property + def max_support(self): + """Return the largest placement arity, currently always two.""" + return max((len(pair) for pair in self.pairs), default=0) + + def greedy_rounds(self): + """Partition the stream into deterministic non-overlapping rounds.""" + rounds = [] + for pair in self.pairs: + support = set(pair) + for current in rounds: + if all(not support.intersection(other) for other in current): + current.append(pair) + break + else: + rounds.append([pair]) + return tuple(tuple(current) for current in rounds) + + @property + def round_depth(self): + """Return the greedy non-overlapping stream depth.""" + return len(self.greedy_rounds()) + + +def load_qmera_prototype_layout( + path, + *, + loader: Callable[[str], Any] | None = None, + num_sites: int | None = None, + level: int | None = None, +): + """Load a serialized ``U_q3_l*`` prototype placement stream. + + ``quimb.load_from_disk`` is imported only when no loader is supplied, so + callers can test or adapt the format without making quimb serialization a + new hard dependency. The returned stream is a diagnostic adapter; use + :class:`QMeraLayoutFinder` for Pepsy-native RG schedules. + """ + path = Path(path) + if loader is None: + try: + import quimb as qu # pylint: disable=import-outside-toplevel + except ImportError as exc: # pragma: no cover - optional dependency + raise ImportError( + "Loading qMERA prototype streams requires `quimb` or a custom loader." + ) from exc + loader = qu.load_from_disk + + raw = loader(str(path)) + try: + entries = tuple(raw) + except TypeError as exc: + raise TypeError("prototype layout data must be an iterable of site pairs.") from exc + pairs = tuple(_normalize_pair(pair, index) for index, pair in enumerate(entries)) + inferred_sites = max((max(pair) for pair in pairs), default=-1) + 1 + if num_sites is None: + if inferred_sites < 1: + raise ValueError("num_sites is required when the prototype stream is empty.") + num_sites = inferred_sites + if int(num_sites) < max(1, inferred_sites): + raise ValueError( + f"num_sites={num_sites} does not cover prototype site {inferred_sites - 1}." + ) + if level is None: + match = _LEVEL_RE.search(path.name) + level = None if match is None else int(match.group(1)) + return QMeraPrototypeLayout( + name=path.name, + path=str(path), + pairs=pairs, + num_sites=int(num_sites), + level=level, + ) diff --git a/src/pepsy/optimizers/mera/schedules.py b/src/pepsy/optimizers/qmera/schedules.py similarity index 91% rename from src/pepsy/optimizers/mera/schedules.py rename to src/pepsy/optimizers/qmera/schedules.py index 3f6f87b..143e4ef 100644 --- a/src/pepsy/optimizers/mera/schedules.py +++ b/src/pepsy/optimizers/qmera/schedules.py @@ -1,9 +1,10 @@ """qMERA RG schedules and reverse-lightcone metadata. -The schedule grammar is bottom-to-top MERA-like rather than a generic brickwall -circuit: isometry blocks form a non-overlapping covering partition of active +The schedule grammar is bottom-to-top qMERA RG blocking rather than a generic +brickwall circuit: isometry blocks form a non-overlapping covering partition of active sites, and disentangler blocks are boundary windows between adjacent isometry -blocks. +blocks. In 2D, boundary windows are colored into disjoint executable +subrounds; their supports intentionally overlap neighboring isometry blocks. """ from __future__ import annotations @@ -386,7 +387,13 @@ def arity(self): @dataclass(frozen=True) class QMeraLayerSpec: - """One MERA scale with boundary disentanglers and covering isometries.""" + """One MERA scale with ordered, locally disjoint gate subrounds. + + Isometry blocks form a disjoint covering partition. Disentangler blocks + are boundary windows between that partition; each executable disentangler + round is disjoint, while a disentangler may overlap an isometry block by + design. + """ scale: int input_sites: tuple[int, ...] @@ -467,17 +474,31 @@ def add(tag): add(tag) return tuple(tags) - def schematic_blocks(self, *, layer=None): - """Return display-oriented disentangler/isometry blocks.""" + def schematic_blocks(self, *, layer=None, rg_step=None): + """Return display blocks for one or more RG steps. + + ``rg_step`` is the descriptive alias for ``layer`` used by the + schematic API. Supplying both is an error so a drawing cannot silently + select the wrong scale. + """ from .schematics import qmera_schematic_blocks - return qmera_schematic_blocks(self, layer=layer) + return qmera_schematic_blocks(self, layer=layer, rg_step=rg_step) + + def draw_schematic(self, *, layer=None, rg_step=None, **kwargs): + """Draw qMERA blocking for one or more RG steps. - def draw_schematic(self, *, layer=None, **kwargs): - """Draw a schematic of qMERA blocking for one or more layers.""" + Use ``rg_step=0`` for the first coarse-graining step. ``layer`` is + retained as a backwards-compatible alias. + """ from .schematics import draw_qmera_schedule - return draw_qmera_schedule(self, layer=layer, **kwargs) + return draw_qmera_schedule( + self, + layer=layer, + rg_step=rg_step, + **kwargs, + ) def _nonoverlapping_blocks(active, block_size): @@ -722,14 +743,24 @@ def _stage_placements( mode_order=None, boundary_pairs_by_block=None, block_axes=None, + block_rounds=None, + round_stride=1, periodic=False, ): placements = [] counter = counter_start stage = stage_spec.kind short = "DIS" if stage == "disentangler" else "ISO" - for round_index in range(stage_spec.circuit_depth): + if round_stride < 1: + raise ValueError("round_stride must be >= 1.") + if block_rounds is not None and len(block_rounds) != len(blocks): + raise ValueError("block_rounds must match the number of stage blocks.") + for circuit_round in range(stage_spec.circuit_depth): for block_index, block in _block_ranges(blocks): + block_round = ( + 0 if block_rounds is None else int(block_rounds[block_index]) + ) + round_index = circuit_round * round_stride + block_round if boundary_pairs_by_block is not None: pairs = boundary_pairs_by_block[block_index] axis = None if block_axes is None else block_axes[block_index] @@ -785,6 +816,45 @@ def _stage_placements( return tuple(placements), counter +def _validate_disentangler_subrounds(blocks, placements): + """Ensure concurrent disentangler blocks are mutually disjoint.""" + block_rounds = {} + for placement in placements: + block_rounds.setdefault(placement.block, set()).add(placement.round) + by_round = {} + for block_index, rounds in block_rounds.items(): + for round_index in rounds: + by_round.setdefault(round_index, []).append(block_index) + for round_index, block_indices in by_round.items(): + for position, left_index in enumerate(block_indices): + left = set(blocks[left_index]) + for right_index in block_indices[:position]: + common = left.intersection(blocks[right_index]) + if common: + raise ValueError( + "disentangler blocks must be disjoint within executable " + f"subround {round_index}: {left_index} and {right_index} " + f"share {tuple(sorted(common))!r}." + ) + + +def _disjoint_block_colors(blocks): + """Greedily color overlapping blocks into disjoint execution rounds.""" + colors = [] + supports = [set(block) for block in blocks] + for index, support in enumerate(supports): + used = { + colors[previous] + for previous in range(index) + if support.intersection(supports[previous]) + } + color = 0 + while color in used: + color += 1 + colors.append(color) + return tuple(colors), max(colors, default=-1) + 1 + + def _coarse_grain(isometry_blocks): return tuple(block[0] for block in isometry_blocks if block) @@ -1144,6 +1214,7 @@ def _build_qmera_schedule_1d( geometry.boundary == "periodic" and disentangler.periodic_wrap ), ) + _validate_disentangler_subrounds(disentangler_blocks, dis) else: disentangler_blocks = _boundary_blocks( isometry_blocks, @@ -1152,6 +1223,9 @@ def _build_qmera_schedule_1d( geometry.boundary == "periodic" and disentangler.periodic_wrap ), ) + dis_block_rounds, dis_round_count = _disjoint_block_colors( + disentangler_blocks + ) dis, gate_counter = _stage_placements( disentangler_blocks, scale=scale, @@ -1159,7 +1233,10 @@ def _build_qmera_schedule_1d( counter_start=gate_counter, mode_by_site=mode_by_site, mode_order=mode_order, + block_rounds=dis_block_rounds, + round_stride=max(1, dis_round_count), ) + _validate_disentangler_subrounds(disentangler_blocks, dis) iso, gate_counter = _stage_placements( _placement_blocks(isometry_blocks), scale=scale, @@ -1257,6 +1334,7 @@ def _build_qmera_schedule_2d( geometry.boundary == "periodic" and disentangler.periodic_wrap ), ) + _validate_disentangler_subrounds(disentangler_blocks, dis) else: ( disentangler_blocks, @@ -1276,6 +1354,9 @@ def _build_qmera_schedule_2d( corner_policy=disentangler.corner_policy, ) if disentangler.placement == "boundary-square": + dis_block_rounds, dis_round_count = _disjoint_block_colors( + disentangler_blocks + ) dis, gate_counter = _stage_placements( disentangler_blocks, scale=scale, @@ -1284,8 +1365,14 @@ def _build_qmera_schedule_2d( coords_by_site=coords_by_site, mode_by_site=mode_by_site, mode_order=geometry.site_modes, + block_rounds=dis_block_rounds, + round_stride=max(1, dis_round_count), ) + _validate_disentangler_subrounds(disentangler_blocks, dis) elif disentangler.placement != "within-block": + dis_block_rounds, dis_round_count = _disjoint_block_colors( + disentangler_blocks + ) dis, gate_counter = _stage_placements( disentangler_blocks, scale=scale, @@ -1294,7 +1381,10 @@ def _build_qmera_schedule_2d( boundary_pairs_by_block=dis_pairs_by_block, block_axes=dis_axes, mode_order=geometry.site_modes, + block_rounds=dis_block_rounds, + round_stride=max(1, dis_round_count), ) + _validate_disentangler_subrounds(disentangler_blocks, dis) iso, gate_counter = _stage_placements( _placement_blocks(isometry_blocks), scale=scale, diff --git a/src/pepsy/optimizers/mera/schematics.py b/src/pepsy/optimizers/qmera/schematics.py similarity index 83% rename from src/pepsy/optimizers/mera/schematics.py rename to src/pepsy/optimizers/qmera/schematics.py index 789b3d4..e9b8d27 100644 --- a/src/pepsy/optimizers/mera/schematics.py +++ b/src/pepsy/optimizers/qmera/schematics.py @@ -41,6 +41,13 @@ def _layer_indices(schedule, layer=None): return tuple(int(idx) for idx in layer) +def _resolve_rg_step(layer, rg_step): + """Resolve the public ``layer``/``rg_step`` selectors.""" + if layer is not None and rg_step is not None: + raise TypeError("Specify only one of layer= or rg_step=.") + return rg_step if rg_step is not None else layer + + def _placement_group_key(placement): return (placement.scale, placement.stage, placement.round, placement.block) @@ -56,9 +63,14 @@ def _block_sites_for_group(layer_spec, stage, block_index, fallback): return tuple(fallback) -def qmera_schematic_blocks(schedule, *, layer=None): - """Return qMERA layer blocks grouped for schematic display.""" +def qmera_schematic_blocks(schedule, *, layer=None, rg_step=None): + """Return qMERA blocks grouped for schematic display. + + ``rg_step`` selects a single bottom-to-top coarse-graining step and is a + readable alias for the older ``layer`` argument. + """ layers = schedule.layers + layer = _resolve_rg_step(layer, rg_step) requested = set(_layer_indices(schedule, layer)) blocks = [] for layer_index in requested: @@ -196,11 +208,23 @@ def _unique_physical_sites(geometry, register_sites): return tuple(sites) -def _clean_stage_positions(geometry, sites, *, x0, y0): - """Place 2D physical sites on a stable schematic grid.""" +def _clean_stage_positions(sites, *, x0, y0): + """Place a stage's active sites on a compact local schematic grid. + + Coarse sites retain their physical coordinate labels, e.g. ``(0, 2)``, + but their drawing should occupy adjacent positions in the coarse panel. + Rank-compressing each stage removes misleading gaps between RG scales. + """ + # ``_draw_clean_stage`` has already converted register labels to physical + # site labels before calling this helper. + coords = {site: site for site in sites} + xs = sorted({coord[1] for coord in coords.values()}) + ys = sorted({coord[0] for coord in coords.values()}) + x_rank = {value: pos for pos, value in enumerate(xs)} + y_rank = {value: pos for pos, value in enumerate(ys)} return { - site: (x0 + float(site[1]), y0 - float(site[0])) - for site in sites + site: (x0 + float(x_rank[coord[1]]), y0 - float(y_rank[coord[0]])) + for site, coord in coords.items() } @@ -224,7 +248,7 @@ def _draw_clean_stage( else layer_spec.input_sites ) sites = _unique_physical_sites(geometry, register_sites) - positions = _clean_stage_positions(geometry, sites, x0=x0, y0=y0) + positions = _clean_stage_positions(sites, x0=x0, y0=y0) # Keep the physical graph visible in every stage, like the simple wires in # quimb's manual schematic examples. @@ -275,15 +299,17 @@ def _draw_clean_2d( ): """Draw 2D layers as separated quimb-style input/D/W/output panels.""" height, width = schedule.geometry.shape - panel_width = float(max(width, height)) + 2.0 - panel_gap = 1.5 - row_gap = float(max(height, 1)) + 3.2 - cursor_x = 0.0 + panel_extent = float(max(width, height) - 1) + panel_width = panel_extent + 1.2 + panel_gap = 1.1 + row_gap = panel_extent + 3.0 for row, layer_index in enumerate(layer_indices): layer = schedule.layers[layer_index] + # Each RG scale is its own readable horizontal strip. Keeping the + # origin fixed prevents later scales from drifting to the right. + cursor_x = 0.0 y0 = -row_gap * row - stages = [("input", "input")] disentanglers = tuple( block for block in blocks_by_layer.get(layer_index, ()) @@ -294,15 +320,18 @@ def _draw_clean_2d( for block in blocks_by_layer.get(layer_index, ()) if block.stage == "isometry" ) + stages = [("input", "input", None)] if disentanglers: - stages.append(("disentangler", "D")) + for round_index in sorted({block.round for block in disentanglers}): + stages.append(("disentangler", f"D[r{round_index}]", round_index)) if isometries: - stages.append(("isometry", "W")) - stages.append(("output", "coarse")) + for round_index in sorted({block.round for block in isometries}): + stages.append(("isometry", f"W[r{round_index}]", round_index)) + stages.append(("output", "coarse", None)) layer_start = cursor_x previous_right = None - for stage, label in stages: + for stage, label, round_index in stages: x0 = cursor_x blocks = ( disentanglers @@ -311,6 +340,8 @@ def _draw_clean_2d( if stage == "isometry" else () ) + if round_index is not None: + blocks = tuple(block for block in blocks if block.round == round_index) _draw_clean_stage( drawing, schedule, @@ -322,14 +353,18 @@ def _draw_clean_2d( label_sites=label_sites, label_blocks=label_blocks, ) + if stage == "disentangler": + label = f"{label} ({len(blocks)} blocks)" + elif stage == "isometry": + label = f"{label} ({len(blocks)} blocks)" drawing.text( - (x0 + 0.5 * (width - 1), y0 + 0.85), + (x0 + 0.5 * panel_extent, y0 + 0.85), label, preset="stage_label", ) if previous_right is not None: - arrow_y = y0 - float(height) - 0.7 + arrow_y = y0 - 0.5 * panel_extent start = (previous_right, arrow_y) end = (x0 - 0.35, arrow_y) drawing.line(start, end, preset="flow") @@ -338,14 +373,13 @@ def _draw_clean_2d( cursor_x += panel_width + panel_gap drawing.text( - (layer_start - 0.7, y0), + (layer_start - 0.8, y0 + 0.85), f"L{layer_index}", preset="layer_label", ) - cursor_x += 0.8 drawing.text( - (0.0, -row_gap * len(layer_indices) + 1.0), + (0.5 * panel_extent, -row_gap * len(layer_indices) + 0.9), "D = boundary disentangler W = isometry arrows = coarse-graining", preset="legend", ) @@ -411,6 +445,7 @@ def draw_qmera_schedule( schedule, *, layer=None, + rg_step=None, style="clean", figsize=None, label_sites=True, @@ -424,10 +459,11 @@ def draw_qmera_schedule( Parameters ---------- schedule - A :class:`~pepsy.optimizers.mera.QMeraSchedule`. - layer - Optional layer index or iterable of layer indices. ``None`` draws all - layers. + A :class:`~pepsy.optimizers.qmera.QMeraSchedule`. + layer, rg_step + Optional layer index or iterable of layer indices. ``rg_step`` is the + preferred descriptive name for selecting one RG step; ``layer`` is a + backwards-compatible alias. ``None`` draws all steps. style : {"clean", "register"}, default="clean" ``"clean"`` separates 2D input, disentangler, isometry, and coarse output panels. ``"register"`` keeps the lower-level register wiring @@ -463,14 +499,14 @@ def draw_qmera_schedule( }, "site_label": {"fontsize": 8, "color": neutral_dark}, "disentangler": { - "facecolor": schematic.get_color("orange", alpha=0.34), + "facecolor": schematic.get_color("orange", alpha=0.18), "edgecolor": schematic.get_color("orange"), - "linewidth": 1.5, + "linewidth": 1.4, }, "isometry": { - "facecolor": schematic.get_color("green", alpha=0.30), + "facecolor": schematic.get_color("green", alpha=0.16), "edgecolor": schematic.get_color("green"), - "linewidth": 1.5, + "linewidth": 1.4, }, "block_label": { "fontsize": 9, @@ -500,6 +536,7 @@ def draw_qmera_schedule( drawing = schematic.Drawing(**kwargs) layers = schedule.layers + layer = _resolve_rg_step(layer, rg_step) layer_indices = _layer_indices(schedule, layer) blocks = qmera_schematic_blocks(schedule, layer=layer_indices) blocks_by_layer = {} diff --git a/src/pepsy/optimizers/mera/terms.py b/src/pepsy/optimizers/qmera/terms.py similarity index 98% rename from src/pepsy/optimizers/mera/terms.py rename to src/pepsy/optimizers/qmera/terms.py index f189d8d..0486976 100644 --- a/src/pepsy/optimizers/mera/terms.py +++ b/src/pepsy/optimizers/qmera/terms.py @@ -1,4 +1,4 @@ -"""Local-term normalization for MERA energy objectives.""" +"""Local-term normalization for qMERA energy objectives.""" from __future__ import annotations diff --git a/tests/conftest.py b/tests/conftest.py index f631544..ce8d630 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -26,7 +26,7 @@ "test_bp_relay.py", "test_fh_jw_gates.py", "test_optimize_global.py", - "test_optimize_mera.py", + "test_optimize_qmera.py", "test_optimize_peps.py", "test_optimize_tree.py", "test_optimize_tree_stabilizer.py", diff --git a/tests/test_optimize_mera.py b/tests/test_optimize_qmera.py similarity index 67% rename from tests/test_optimize_mera.py rename to tests/test_optimize_qmera.py index ef9ba40..bb017f1 100644 --- a/tests/test_optimize_mera.py +++ b/tests/test_optimize_qmera.py @@ -1,43 +1,40 @@ -"""Tests for MERA local-energy optimization helpers.""" +"""Tests for qMERA local-energy optimization helpers.""" import numpy as np import pytest -import quimb.tensor as qtn -from pepsy.optimizers.energy import EnergyEstimate -from pepsy.optimizers.mera import ( +from pepsy.optimizers.qmera import ( GateSpec, LocalTerm, - MeraEnergyOptimizer, QMeraBlockSpec, QMeraBuilder, QMeraCompiledLightconeChunk, QMeraContractionPathCache, QMeraDisentanglerSpec, + QMeraEnergyOptimizer, QMeraGeometry, QMeraIsometrySpec, + QMeraLayoutFinder, QMeraLightconeGroup, QMeraLightconeTN, QMeraSchematicBlock, QMeraParametricEnergyOptimizer, QMeraParametricLightconeChunk, + QMeraPrototypeLayout, QMeraSymmrayFermionBackend, QMeraScaleSpec, QMeraUnitarySpec, UserGateFamily, - build_lightcone_chunks, build_qmera_contraction_optimizer, - build_qmera_lightcone_chunks, build_qmera_parametric_lightcone_chunks, compile_qmera_parametric_lightcones, contract_qmera_lightcone_tn, group_qmera_parametric_lightcone_chunks, - lightcone_energy, default_gate_registry, draw_qmera_schedule, local_qmera_compiled_lightcone_expectation, local_qmera_parametric_lightcone_expectation, - local_lightcone_expectation, + load_qmera_prototype_layout, normalize_local_terms, qmera_compiled_parametric_energy, qmera_direct_parametric_energy, @@ -50,9 +47,6 @@ symmray_fermion_gate_registry, symmray_majorana_gate_registry, ) -from pepsy.optimizers.mera.optimizer import ( - MeraEnergyOptimizer as ModuleMeraEnergyOptimizer, -) from pepsy.tensors import Fermion @@ -61,10 +55,6 @@ def _zz_term(): return np.kron(z_op, z_op).reshape(2, 2, 2, 2) -def _small_mera(seed=23): - return qtn.MERA.rand(L=8, max_bond=2, dtype="complex128", seed=seed) - - def test_qmera_contraction_optimizer_helper_delegates(monkeypatch): """qMERA cache helper should expose Pepsy's reusable cotengra builder.""" calls = [] @@ -74,7 +64,7 @@ def fake_build_optimizer(**kwargs): return "optimizer" monkeypatch.setattr( - "pepsy.optimizers.mera.cache.build_optimizer", + "pepsy.optimizers.qmera.cache.build_optimizer", fake_build_optimizer, ) @@ -94,7 +84,7 @@ def fake_build_qmera_contraction_optimizer(**kwargs): return "optimizer" monkeypatch.setattr( - "pepsy.optimizers.mera.builders.build_qmera_contraction_optimizer", + "pepsy.optimizers.qmera.builders.build_qmera_contraction_optimizer", fake_build_qmera_contraction_optimizer, ) @@ -139,6 +129,54 @@ def test_qmera_builder_preserves_explicit_geometry_override(): QMeraBuilder(geometry=bad_geometry, fermion=fermion) +def test_qmera_layout_finder_returns_valid_ranked_candidates(): + """Layout search should rank immutable schedule candidates.""" + geometry = QMeraGeometry(shape=8) + finder = QMeraLayoutFinder( + geometry, + isometry_block_shapes=(2, 4), + disentangler_block_shapes=(2,), + disentangler_depths=(1, 2), + isometry_depths=(1,), + max_layers=4, + ) + + report = finder.search({(0, 1): _zz_term()}) + repeat = finder.search({(0, 1): _zz_term()}) + + assert report.candidates + assert report.scores + assert report.best is not None + assert report.best_score is not None + assert report.pareto_front + assert tuple(candidate.candidate_id for candidate in report.candidates) == tuple( + candidate.candidate_id for candidate in repeat.candidates + ) + assert all(score.valid for score in report.scores) + assert all("max_lightcone_width" in score.components for score in report.scores) + + +def test_qmera_prototype_layout_loader_and_structural_score(): + """Prototype U-streams should load without being mistaken for schedules.""" + layout = load_qmera_prototype_layout( + "/tmp/U_q3_l1", + loader=lambda _path: ((0, 1), (2, 3), (1, 2)), + num_sites=4, + ) + + assert isinstance(layout, QMeraPrototypeLayout) + assert layout.level == 1 + assert layout.gate_count == 3 + assert layout.round_depth == 2 + assert layout.unique_sites == (0, 1, 2, 3) + + finder = QMeraLayoutFinder(QMeraGeometry(shape=4)) + score = finder.score_prototype_layout(layout, {(0, 1): _zz_term()}) + assert score.valid + assert score.components["interaction_coverage"] == pytest.approx(1.0) + assert score.components["gate_count"] == 3 + + def test_normalize_local_terms_accepts_mapping_iterable_and_local_term(): """Hamiltonian input should normalize to explicit LocalTerm objects.""" op = _zz_term() @@ -167,75 +205,6 @@ def test_normalize_local_terms_rejects_bad_supports(): normalize_local_terms([((0, 1), op, 1.0, "extra")]) -def test_mera_lightcone_expectation_matches_quimb_exact_contraction(): - """The cached lightcone kernel should match quimb's full exact oracle.""" - mera = _small_mera() - where = (0, 1) - op = _zz_term() - terms = normalize_local_terms({where: op}) - chunk = build_lightcone_chunks(mera, terms)[0] - - local_value = local_lightcone_expectation( - mera, - chunk, - optimize="auto-hq", - normalized=True, - real=False, - ) - direct = mera.compute_local_expectation_exact( - {where: op}, - optimize="auto-hq", - normalized=True, - ) - - assert complex(local_value) == pytest.approx(complex(direct)) - assert chunk.tags == ("I0", "I1") - assert chunk.physical_width == 2 - - -def test_mera_energy_loss_matches_quimb_exact_sum(): - """MeraEnergyOptimizer.loss() should sum local lightcone contractions.""" - mera = _small_mera(seed=24) - op = _zz_term() - terms = {(0, 1): op, (2, 3): op} - direct = mera.compute_local_expectation_exact( - terms, - optimize="auto-hq", - normalized=True, - ) - opt = MeraEnergyOptimizer( - mera, - terms, - energy_per_site=False, - normalized=True, - contraction_opt="auto-hq", - ) - - assert complex(opt.loss(real=False)) == pytest.approx(complex(direct)) - - -def test_generic_lightcone_energy_groups_select_gate_and_contract(): - """The public fixed-state helper should match the full MERA oracle.""" - mera = _small_mera(seed=241) - terms = {(0, 1): _zz_term(), (2, 3): _zz_term()} - direct = mera.compute_local_expectation_exact( - terms, - optimize="auto-hq", - normalized=True, - ) - - value = lightcone_energy( - mera, - terms, - energy_per_site=False, - normalized=True, - real=False, - group_terms=True, - ) - - assert complex(value) == pytest.approx(complex(direct)) - - def test_qmera_parameter_sharing_per_block_reuses_round_parameters(): """One block can share parameters across its brickwall rounds.""" unitary = QMeraUnitarySpec( @@ -267,72 +236,6 @@ def test_qmera_parameter_sharing_per_block_reuses_round_parameters(): assert len({next(iter(keys)) for keys in by_block.values()}) == len(by_block) -def test_mera_energy_estimate_reports_lightcone_metadata(): - """energy() should return the shared EnergyEstimate dataclass.""" - mera = _small_mera(seed=25) - opt = MeraEnergyOptimizer( - mera, - {(0, 1): _zz_term()}, - energy_per_site=True, - normalized=True, - ) - - estimate = opt.energy() - - assert isinstance(estimate, EnergyEstimate) - assert estimate.num_sites == 8 - assert estimate.boundary_mode == "lightcone-exact" - assert estimate.energy_per_site == pytest.approx(estimate.energy / 8) - assert estimate.metadata["num_terms"] == 1 - assert estimate.metadata["max_physical_width"] == 2 - assert estimate.metadata["max_lightcone_tensors"] >= 1 - - -def test_mera_energy_make_tn_optimizer_and_optimize(monkeypatch): - """TNOptimizer construction should receive loss constants and norm hook.""" - calls = [] - out = _small_mera(seed=27) - - class _FakeTNOptimizer: # pylint: disable=too-few-public-methods - def __init__(self, state, loss_fn, **kwargs): - calls.append((state, loss_fn, kwargs)) - self.losses = [2.0, 1.0] - - def optimize(self, n=220, **kwargs): - calls.append(("optimize", n, kwargs)) - return out - - monkeypatch.setattr( - "pepsy.optimizers.mera.optimizer.qtn.TNOptimizer", - _FakeTNOptimizer, - ) - mera = _small_mera(seed=26) - terms = {(0, 1): _zz_term()} - opt = MeraEnergyOptimizer(mera, terms) - - tnopt = opt.make_tn_optimizer( - optimizer="lbfgs", - autodiff_backend="jax", - progbar=False, - loss_kwargs={"precompute_tags": False}, - ) - assert isinstance(tnopt, _FakeTNOptimizer) - _, loss_fn, kwargs = calls[0] - assert loss_fn is ModuleMeraEnergyOptimizer._tnopt_loss - assert kwargs["loss_constants"]["terms"] == opt.terms - assert kwargs["loss_constants"]["chunks"] is None - assert kwargs["loss_kwargs"]["precompute_tags"] is False - assert kwargs["optimizer"] == "L-BFGS-B" - assert kwargs["autodiff_backend"] == "jax" - assert callable(kwargs["norm_fn"]) - - optimized, losses = opt.optimize(n=3, progbar=False, return_losses=True) - assert optimized is out - assert opt.state is out - assert losses == (2.0, 1.0) - assert calls[-1] == ("optimize", 3, {}) - - def test_qmera_geometry_explicit_lattice_and_mapper(): """Geometry should keep physical labels separate from register sites.""" geom = QMeraGeometry(shape=(2, 3), mapper="snake", boundary="periodic") @@ -458,6 +361,43 @@ def test_qmera_2d_schedule_uses_rg_blocks_and_face_disentanglers(): assert any(placement.axis == "x" for placement in first.isometries) assert any(placement.axis == "y" for placement in first.isometries) + # Boundary disentangler blocks are assigned disjoint executable rounds. + # Their supports intentionally overlap the neighboring isometry blocks, + # which is the MERA boundary-coupling pattern. + d_round_by_block = {} + for placement in first.disentanglers: + d_round_by_block.setdefault(placement.block, placement.round) + assert d_round_by_block[placement.block] == placement.round + for round_index in set(d_round_by_block.values()): + blocks = [ + set(first.disentangler_blocks[block_index]) + for block_index, block_round in d_round_by_block.items() + if block_round == round_index + ] + assert all( + not (left & right) + for index, left in enumerate(blocks) + for right in blocks[:index] + ) + assert all( + not (left & right) + for index, left in enumerate(map(set, first.isometry_blocks)) + for right in map(set, first.isometry_blocks[:index]) + ) + assert any( + set(disentangler) & set(isometry) + for disentangler in first.disentangler_blocks + for isometry in first.isometry_blocks + ) + assert all( + sum( + bool(set(disentangler) & set(isometry)) + for isometry in first.isometry_blocks + ) + >= 2 + for disentangler in first.disentangler_blocks + ) + def test_qmera_explicit_specs_build_4x4_periodic_hubbard_schedule(): """Explicit square layers should include all 4x4 PBC interfaces.""" @@ -869,8 +809,8 @@ def test_qmera_symmray_fermion_lightcone_contracts_native_term(): energy_per_site=False, real=False, ) - fixed_state_value = lightcone_energy( - builder.build(params), + direct_value = builder.direct_parametric_loss( + params, terms[:1], schedule=schedule, convert_terms=False, @@ -886,7 +826,342 @@ def test_qmera_symmray_fermion_lightcone_contracts_native_term(): ) assert complex(value) == pytest.approx(0.0) assert complex(builder_value) == pytest.approx(complex(value)) - assert complex(fixed_state_value) == pytest.approx(complex(value)) + assert complex(direct_value) == pytest.approx(complex(value)) + + +def test_qmera_native_symmray_compilation_preserves_graded_metadata(): + """Compiled native cones keep Symmray grading instead of densifying.""" + pytest.importorskip("symmray") + backend = QMeraSymmrayFermionBackend() + registry = symmray_fermion_gate_registry(backend=backend) + builder = QMeraBuilder( + shape=2, + site_modes=backend.site_modes, + gate_registry=registry, + gate_family="symmray-fsim", + isometry={ + "block_size": 2, + "circuit_depth": 1, + "gate_family": "symmray-fsim", + }, + max_layers=1, + product_state_factory=backend.product_state, + ) + terms = backend.fermi_hubbard_terms( + builder.geometry, + t=0.2, + U=0.0, + mu=0.0, + ) + + schedule = builder.build_schedule() + parameters = builder.initialize_parameters(schedule) + chunks = builder.parametric_lightcone_chunks( + terms[:1], schedule, convert_terms=False, + ) + compiled = builder.compile_parametric_lightcones( + chunks=chunks, + schedule=schedule, + convert_terms=False, + ) + value = builder.compiled_parametric_loss( + parameters, + schedule=schedule, + compiled_chunks=compiled, + energy_per_site=False, + real=False, + ) + direct = builder.parametric_loss( + parameters, + schedule=schedule, + chunks=chunks, + convert_terms=False, + energy_per_site=False, + real=False, + ) + + assert compiled[0].is_graded + assert compiled[0].contraction_backend == "symmray" + assert compiled[0].symmetry == "U1U1" + assert complex(value) == pytest.approx(complex(direct)) + + +def test_qmera_compiled_native_pbc_terms_match_each_explicit_cone(): + """A compiled periodic Hubbard cone agrees term-by-term with Symmray.""" + pytest.importorskip("symmray") + geometry = QMeraGeometry( + shape=(2, 2), + boundary="periodic", + site_modes=("up", "down"), + mode_order="mode-major", + ) + backend = QMeraSymmrayFermionBackend( + symmetry="U1U1", + site_modes=("up", "down"), + mode_order="mode-major", + ) + registry = symmray_fermion_gate_registry(backend=backend) + + def product_state_factory(schedule, sites, **kwargs): + occupations = {0: 1, 1: 0, 2: 0, 3: 1, 4: 0, 5: 0, 6: 0, 7: 0} + return backend.product_state( + schedule, + sites, + occupations=occupations, + **kwargs, + ) + + builder = QMeraBuilder( + geometry=geometry, + gate_registry=registry, + gate_family="symmray-fsim", + disentangler={ + "block_size": 2, + "circuit_depth": 1, + "gate_family": "symmray-fsim", + }, + isometry={ + "block_size": (2, 2), + "circuit_depth": 1, + "gate_family": "symmray-fsim", + }, + max_layers=1, + seed=12, + param_scale=0.07, + product_state_factory=product_state_factory, + ) + schedule = builder.build_schedule() + parameters = builder.initialize_parameters(schedule) + terms = qmera_symmray_fermi_hubbard_terms( + geometry, + backend=backend, + t=0.2, + U=0.5, + mu=0.1, + peierls_angle=np.pi / 3, + ) + chunks = builder.parametric_lightcone_chunks( + terms, + schedule, + convert_terms=False, + ) + + class RecordingPathCache: + def __init__(self): + self.keys = [] + + def resolve(self, optimize, *, key=None): + self.keys.append(key) + return "greedy" if str(optimize).startswith("auto") else optimize + + path_cache = RecordingPathCache() + compiled = builder.compile_parametric_lightcones( + chunks=chunks, + schedule=schedule, + convert_terms=False, + path_cache=path_cache, + ) + compiled_values = [ + local_qmera_compiled_lightcone_expectation( + schedule, + item, + parameters, + gate_registry=registry, + normalized=False, + real=False, + ) + for item in compiled + ] + explicit_values = [ + contract_qmera_lightcone_tn( + builder.parametric_lightcone_tn( + chunk, + parameters, + schedule=schedule, + gate_array_backend=None, + ), + optimize="greedy", + normalized=False, + real=False, + ) + for chunk in chunks + ] + + assert len(compiled) == len(terms) + assert len(set(path_cache.keys)) < len(path_cache.keys) + for compiled_value, explicit_value in zip(compiled_values, explicit_values): + assert complex(compiled_value) == pytest.approx(complex(explicit_value)) + + +def test_qmera_compiled_native_z2_majorana_matches_explicit(): + """The graded compiler also supports the parity-only Majorana symmetry.""" + pytest.importorskip("symmray") + geometry = QMeraGeometry( + shape=(2, 2), + boundary="periodic", + site_modes=("mode",), + mode_order="mode-major", + ) + backend = QMeraSymmrayFermionBackend( + symmetry="Z2", + site_modes=("mode",), + ) + registry = symmray_majorana_gate_registry(backend=backend) + builder = QMeraBuilder( + geometry=geometry, + gate_registry=registry, + gate_family="symmray-majorana", + disentangler={ + "block_size": 2, + "circuit_depth": 1, + "gate_family": "symmray-majorana", + }, + isometry={ + "block_size": (2, 2), + "circuit_depth": 1, + "gate_family": "symmray-majorana", + }, + max_layers=1, + seed=15, + param_scale=0.03, + product_state_factory=backend.product_state, + ) + schedule = builder.build_schedule() + parameters = builder.initialize_parameters(schedule) + terms = qmera_symmray_majorana_terms( + geometry, + fermion=Fermion(spinful=False, symmetry="Z2"), + coupling=0.4, + pairing=0.2, + ) + chunks = builder.parametric_lightcone_chunks( + terms, + schedule, + convert_terms=False, + ) + compiled = builder.compile_parametric_lightcones( + chunks=chunks, + schedule=schedule, + convert_terms=False, + contraction_opt="greedy", + ) + compiled_value = builder.compiled_parametric_loss( + parameters, + schedule=schedule, + compiled_chunks=compiled, + energy_per_site=False, + real=False, + ) + explicit_value = builder.parametric_loss( + parameters, + schedule=schedule, + chunks=chunks, + convert_terms=False, + energy_per_site=False, + real=False, + ) + + assert compiled + assert all(item.is_graded and item.symmetry == "Z2" for item in compiled) + assert complex(compiled_value) == pytest.approx(complex(explicit_value)) + + +def test_qmera_compiled_native_symmray_torch_gradients_match_explicit(): + """Graded compiled contractions keep the qMERA Torch parameter graph.""" + pytest.importorskip("symmray") + torch = pytest.importorskip("torch") + from pepsy.backends import backend_torch + + gate_backend = backend_torch(dtype=torch.complex128) + backend = QMeraSymmrayFermionBackend(to_backend=gate_backend) + registry = symmray_fermion_gate_registry(backend=backend) + + def product_state_factory(schedule, sites, **kwargs): + return backend.product_state( + schedule, + sites, + occupations={0: 1, 1: 0, 2: 0, 3: 1}, + **kwargs, + ) + + builder = QMeraBuilder( + shape=2, + site_modes=backend.site_modes, + mode_order="mode-major", + gate_registry=registry, + gate_family="symmray-fsim", + isometry={ + "block_size": 2, + "circuit_depth": 1, + "gate_family": "symmray-fsim", + }, + max_layers=1, + seed=7, + param_scale=0.04, + product_state_factory=product_state_factory, + ) + schedule = builder.build_schedule() + parameters = builder.cast_params( + builder.initialize_parameters(schedule), + backend="torch", + trainable=True, + dtype=torch.float64, + ) + terms = backend.fermi_hubbard_terms( + builder.geometry, + t=0.2, + U=0.3, + mu=0.1, + ) + chunks = builder.parametric_lightcone_chunks( + terms, + schedule, + convert_terms=False, + ) + compiled = builder.compile_parametric_lightcones( + chunks=chunks, + schedule=schedule, + convert_terms=False, + contraction_opt="greedy", + ) + compiled_value = builder.compiled_parametric_loss( + parameters, + schedule=schedule, + compiled_chunks=compiled, + energy_per_site=False, + real=True, + ) + explicit_value = builder.parametric_loss( + parameters, + schedule=schedule, + chunks=chunks, + convert_terms=False, + energy_per_site=False, + real=True, + ) + compiled_gradients = torch.autograd.grad( + compiled_value, + tuple(parameters.values()), + retain_graph=True, + ) + explicit_gradients = torch.autograd.grad( + explicit_value, + tuple(parameters.values()), + ) + + assert compiled_value.requires_grad + assert explicit_value.requires_grad + assert float(compiled_value) == pytest.approx(float(explicit_value)) + for compiled_gradient, explicit_gradient in zip( + compiled_gradients, + explicit_gradients, + ): + np.testing.assert_allclose( + compiled_gradient.detach().numpy(), + explicit_gradient.detach().numpy(), + rtol=1.0e-10, + atol=1.0e-10, + ) def test_qmera_2d_multi_mode_schedule_retains_modes_and_pairs_like_modes(): @@ -1085,6 +1360,264 @@ def product_state_factory(schedule, sites, **kwargs): assert complex(majorana_lightcone) == pytest.approx(complex(majorana_direct)) +def test_qmera_fermion_every_term_and_grouping_match_direct_oracle(): + """Every native Hubbard term should agree before grouping is enabled.""" + pytest.importorskip("symmray") + geometry = QMeraGeometry( + shape=(2, 2), + site_modes=("up", "down"), + mode_order="mode-major", + ) + backend = QMeraSymmrayFermionBackend( + symmetry="U1U1", + site_modes=("up", "down"), + mode_order="mode-major", + ) + registry = symmray_fermion_gate_registry(backend=backend) + + def product_state_factory(schedule, sites, **kwargs): + occupations = { + site: int( + (sum(schedule.geometry.to_site(site)) % 2 == 0) + == (schedule.geometry.to_mode(site)[1] == "up") + ) + for site in sites + } + return backend.product_state( + schedule, + sites, + occupations=occupations, + **kwargs, + ) + + builder = QMeraBuilder( + geometry=geometry, + gate_registry=registry, + gate_family="symmray-fsim", + disentangler={ + "block_size": 2, + "circuit_depth": 1, + "gate_family": "symmray-fsim", + }, + isometry={ + "block_size": (2, 2), + "circuit_depth": 1, + "gate_family": "symmray-fsim", + }, + max_layers=1, + seed=97, + param_scale=0.01, + product_state_factory=product_state_factory, + ) + schedule = builder.build_schedule() + parameters = builder.initialize_parameters(schedule) + fermion = Fermion(spinful=True, symmetry="U1U1") + terms = builder.fermion_terms(fermion, t=0.2, U=0.7, mu=0.1) + + for term in terms: + chunk = build_qmera_parametric_lightcone_chunks(schedule, (term,)) + local = builder.parametric_loss( + parameters, + (term,), + schedule=schedule, + chunks=chunk, + convert_terms=False, + energy_per_site=False, + group_terms=False, + real=False, + ) + direct = builder.direct_parametric_loss( + parameters, + (term,), + schedule=schedule, + convert_terms=False, + energy_per_site=False, + group_terms=False, + real=False, + ) + assert complex(local) == pytest.approx(complex(direct)) + + grouped = builder.parametric_loss( + parameters, + terms, + schedule=schedule, + convert_terms=False, + energy_per_site=False, + group_terms=True, + real=False, + ) + ungrouped = builder.parametric_loss( + parameters, + terms, + schedule=schedule, + convert_terms=False, + energy_per_site=False, + group_terms=False, + real=False, + ) + assert complex(grouped) == pytest.approx(complex(ungrouped)) + + +def test_qmera_periodic_fermion_terms_match_jordan_wigner_fock_oracle(): + """PBC native qMERA terms must match an independent JW Fock oracle.""" + pytest.importorskip("symmray") + + def jw_annihilate(num_modes, mode): + eye = np.eye(2) + zed = np.diag([1.0, -1.0]) + lower = np.array([[0.0, 1.0], [0.0, 0.0]]) + mats = [zed] * mode + [lower] + [eye] * (num_modes - mode - 1) + out = mats[0] + for matrix in mats[1:]: + out = np.kron(out, matrix) + return out + + def fock_vector(state, geometry, backend): + full = state.contract(all) + labels = tuple(f"k{site}" for site in geometry.register_sites) + permutation = tuple(full.inds.index(label) for label in labels) + # The native contraction already carries the graded swap phases. The + # phase-aware reorder converts its output-index order to the canonical + # qMERA register order without applying a second bosonization gauge. + data = full.data.transpose(permutation, phase=True) + dense = np.asarray(data.to_dense()) + occupation_positions = [] + for axis, register_site in enumerate(geometry.register_sites): + mode = geometry.to_mode(register_site) + occupied_charge = backend.mode_index_map(mode)[1] + charges = [] + for charge, size in data.indices[axis].chargemap.items(): + charges.extend([charge] * int(size)) + occupation_positions.append( + tuple(int(charge == occupied_charge) for charge in charges) + ) + + vector = np.zeros(2 ** geometry.num_modes, dtype=complex) + for index in np.ndindex(dense.shape): + flat = 0 + for axis, local_index in enumerate(index): + flat = (flat << 1) | occupation_positions[axis][local_index] + vector[flat] = dense[index] + return vector + + geometry = QMeraGeometry( + shape=(2, 2), + boundary="periodic", + site_modes=("up", "down"), + mode_order="mode-major", + ) + backend = QMeraSymmrayFermionBackend( + symmetry="U1U1", + site_modes=("up", "down"), + mode_order="mode-major", + ) + registry = symmray_fermion_gate_registry(backend=backend) + + def product_state_factory(schedule, sites, **kwargs): + # The qMERA isometry pairs are (0, 2) and (1, 3) in this register + # ordering, so both native hopping directions are populated. + occupations = { + 0: 1, + 1: 0, + 2: 0, + 3: 1, + 4: 0, + 5: 0, + 6: 0, + 7: 0, + } + return backend.product_state( + schedule, + sites, + occupations=occupations, + **kwargs, + ) + + builder = QMeraBuilder( + geometry=geometry, + gate_registry=registry, + gate_family="symmray-fsim", + disentangler={ + "block_size": 2, + "circuit_depth": 1, + "gate_family": "symmray-fsim", + }, + isometry={ + "block_size": (2, 2), + "circuit_depth": 1, + "gate_family": "symmray-fsim", + }, + max_layers=1, + seed=12, + param_scale=0.2, + product_state_factory=product_state_factory, + ) + schedule = builder.build_schedule() + parameters = builder.initialize_parameters(schedule) + state, _ = builder.build_state(parameters, schedule) + vector = fock_vector(state, geometry, backend) + norm = np.vdot(vector, vector) + assert norm == pytest.approx(1.0) + + t, U, mu, angle = 0.2, 0.5, 0.1, np.pi / 3 + terms = qmera_symmray_fermi_hubbard_terms( + geometry, + backend=backend, + t=t, + U=U, + mu=mu, + peierls_angle=angle, + ) + annihilators = [ + jw_annihilate(geometry.num_modes, mode) + for mode in range(geometry.num_modes) + ] + number_operators = [ + operator.conj().T @ operator for operator in annihilators + ] + + native_total = 0.0j + oracle_total = 0.0j + for term in terms: + kind = term.metadata["kind"] + if kind == "hubbard-onsite": + site = term.metadata["site"] + up = geometry.mode_register(site, "up") + down = geometry.mode_register(site, "down") + oracle_operator = U * number_operators[up] @ number_operators[down] + elif kind == "hubbard-chemical": + site = term.metadata["site"] + mode = geometry.mode_register(site, term.metadata["mode"]) + oracle_operator = -mu * number_operators[mode] + else: + left, right = term.metadata["edge"] + mode = term.metadata["mode"] + left_mode = geometry.mode_register(left, mode) + right_mode = geometry.mode_register(right, mode) + oracle_operator = -t * ( + np.exp(1.0j * angle) + * annihilators[left_mode].conj().T + @ annihilators[right_mode] + + np.exp(-1.0j * angle) + * annihilators[right_mode].conj().T + @ annihilators[left_mode] + ) + + native_term_state = state.gate_inds( + term.operator, + inds=tuple(geometry.site_ind(site) for site in term.where), + contract=False, + inplace=False, + ) + native_value = (state.H & native_term_state).contract(all) / norm + oracle_value = np.vdot(vector, oracle_operator @ vector) / norm + assert complex(native_value) == pytest.approx(complex(oracle_value)) + native_total += native_value + oracle_total += oracle_value + + assert complex(native_total) == pytest.approx(complex(oracle_total)) + + def test_qmera_grouped_and_direct_energy_match_schedule_lightcones(): """Grouping and the full direct-gate oracle must preserve local energy.""" builder = QMeraBuilder(shape=8, seed=12, param_scale=0.02) @@ -1120,7 +1653,7 @@ def fake_builder(**kwargs): return object() monkeypatch.setattr( - "pepsy.optimizers.mera.cache.build_qmera_contraction_optimizer", + "pepsy.optimizers.qmera.cache.build_qmera_contraction_optimizer", fake_builder, ) cache = QMeraContractionPathCache({"directory": False}) @@ -1206,13 +1739,13 @@ def test_qmera_builder_outputs_parameters_gates_state_and_lightcone_tags(): assert "DISENTANGLER" in ansatz.state.tags assert "ISOMETRY" in ansatz.state.tags - opt = MeraEnergyOptimizer( - ansatz.state, + value = builder.parametric_loss( + ansatz.parameters, {(0, 1): _zz_term()}, + schedule=ansatz.schedule, energy_per_site=False, - contraction_opt="auto-hq", ) - assert np.isfinite(float(opt.loss())) + assert np.isfinite(float(value)) def test_qmera_schedule_lightcone_chunks_follow_placements(): @@ -1224,39 +1757,39 @@ def test_qmera_schedule_lightcone_chunks_follow_placements(): seed=19, param_scale=0.1, ) - ansatz = builder.build() + schedule = builder.build_schedule() + params = builder.initialize_parameters(schedule) op = _zz_term() - chunks = build_qmera_lightcone_chunks( - ansatz.state, - ansatz.schedule, - normalize_local_terms({(0, 1): op}), + chunks = builder.parametric_lightcone_chunks( + {(0, 1): op}, + schedule, ) - opt = MeraEnergyOptimizer( - ansatz, + value = builder.parametric_loss( + params, {(0, 1): op}, + schedule=schedule, + chunks=chunks, energy_per_site=False, - contraction_opt="auto-hq", + real=False, ) - direct = ansatz.state.compute_local_expectation_exact( + direct = builder.direct_parametric_loss( + params, {(0, 1): op}, - optimize="auto-hq", - normalized=True, + schedule=schedule, + energy_per_site=False, + real=False, ) expected_ids = tuple( placement.gate_id - for placement in ansatz.schedule.reverse_lightcone_placements((0, 1)) + for placement in schedule.reverse_lightcone_placements((0, 1)) ) - assert opt.schedule is ansatz.schedule - assert opt.lightcones[0].tags == chunks[0].tags - assert opt.lightcones[0].schedule_placement_ids == chunks[0].schedule_placement_ids - assert chunks[0].source == "schedule" + assert chunks[0].source == "parametric-schedule" assert chunks[0].schedule_placement_ids == expected_ids assert any(tag.startswith("GATE_L0_DIS") for tag in chunks[0].tags) assert chunks[0].schedule_width >= chunks[0].support_size - assert complex(opt.loss(real=False)) == pytest.approx(complex(direct)) - assert opt.energy().metadata["lightcone_sources"] == ("schedule",) + assert complex(value) == pytest.approx(complex(direct)) def test_qmera_schedule_lightcone_chunks_map_coordinate_terms(): @@ -1268,29 +1801,39 @@ def test_qmera_schedule_lightcone_chunks_map_coordinate_terms(): seed=20, param_scale=0.03, ) - ansatz = builder.build() + schedule = builder.build_schedule() + params = builder.initialize_parameters(schedule) op = _zz_term() - opt = MeraEnergyOptimizer( - ansatz, - {((0, 0), (0, 1)): op}, + hamiltonian = {((0, 0), (0, 1)): op} + chunks = builder.parametric_lightcone_chunks( + hamiltonian, + schedule, + ) + value = builder.parametric_loss( + params, + hamiltonian, + schedule=schedule, + chunks=chunks, energy_per_site=False, - contraction_opt="auto-hq", + real=False, ) - chunk = opt.lightcones[0] - direct = ansatz.state.compute_local_expectation_exact( - {(0, 1): op}, - optimize="auto-hq", - normalized=True, + direct = builder.direct_parametric_loss( + params, + hamiltonian, + schedule=schedule, + energy_per_site=False, + real=False, ) + chunk = chunks[0] - assert chunk.source == "schedule" + assert chunk.source == "parametric-schedule" assert chunk.term.where == (0, 1) assert chunk.term.metadata["original_where"] == ((0, 0), (0, 1)) assert chunk.term.metadata["register_where"] == (0, 1) assert chunk.schedule_placement_ids assert chunk.schedule_width_by_scale - assert complex(opt.loss(real=False)) == pytest.approx(complex(direct)) + assert complex(value) == pytest.approx(complex(direct)) def test_qmera_parametric_lightcone_loss_rebuilds_local_cone_from_params(): @@ -1533,6 +2076,9 @@ def test_qmera_parametric_optimizer_runs_compiled_torch_solver(): ) assert isinstance(opt, QMeraParametricEnergyOptimizer) + assert isinstance(opt, QMeraEnergyOptimizer) + assert QMeraEnergyOptimizer is QMeraParametricEnergyOptimizer + assert opt.loss_kwargs["normalized"] is False assert np.isfinite(float(initial)) assert result.solver == "torch-adam" assert len(result.history) == 2 @@ -1721,5 +2267,40 @@ def test_qmera_2d_draw_schematic_builds_quimb_drawing(): ) assert isinstance(clean_drawing, schematic.Drawing) assert isinstance(register_drawing, schematic.Drawing) + + # Each RG scale should start from the same left edge. This catches the + # easy-to-miss cursor drift that makes later 2D scales look detached. + all_layers = builder.draw_schematic( + style="clean", + label_sites=False, + label_blocks=False, + scale_figsize=False, + ) + input_labels = [ + text + for text in all_layers.ax.texts + if text.get_text() == "input" + ] + assert len(input_labels) == 2 + assert input_labels[0].get_position()[0] == pytest.approx( + input_labels[1].get_position()[0] + ) + assert input_labels[0].get_position()[1] != pytest.approx( + input_labels[1].get_position()[1] + ) + + first_step = builder.draw_schematic( + rg_step=0, + style="clean", + label_sites=False, + label_blocks=True, + scale_figsize=False, + ) + assert sum(text.get_text() == "input" for text in first_step.ax.texts) == 1 + assert sum(text.get_text() == "coarse" for text in first_step.ax.texts) == 1 + assert any(text.get_text().startswith("D[") for text in first_step.ax.texts) + assert any(text.get_text().startswith("W[") for text in first_step.ax.texts) + with pytest.raises(TypeError, match="only one of layer= or rg_step="): + builder.draw_schematic(layer=0, rg_step=0) with pytest.raises(ValueError, match="style"): builder.draw_schematic(layer=0, style="unknown") diff --git a/tests/test_package_layout.py b/tests/test_package_layout.py index fa8da18..530bea7 100644 --- a/tests/test_package_layout.py +++ b/tests/test_package_layout.py @@ -19,13 +19,14 @@ ) from pepsy.operators import gate, rx from pepsy.optimizers import ( - MeraEnergyOptimizer, MpsEnergyOptimizer, MpsOptimizer, PepsEnergyOptimizer, PepsOptimizer, QMeraBuilder, + QMeraEnergyOptimizer, QMeraGeometry, + QMeraLayoutFinder, QMeraParametricEnergyOptimizer, SimulatorCandidate, SimulatorPlan, @@ -35,6 +36,7 @@ SweepOptimizer, build_qmera_contraction_optimizer, mera as mera_module, + qmera as qmera_module, recommend_simulator, ) from pepsy.sampling import MpsSampler, PepsBpSampler @@ -115,12 +117,13 @@ def test_new_namespace_imports_resolve(): assert callable(gate) assert callable(rx) assert MpsOptimizer is not None - assert MeraEnergyOptimizer is not None assert MpsEnergyOptimizer is not None assert PepsEnergyOptimizer is not None assert PepsOptimizer is not None assert QMeraBuilder is not None + assert QMeraEnergyOptimizer is not None assert QMeraGeometry is not None + assert QMeraLayoutFinder is not None assert QMeraParametricEnergyOptimizer is not None assert SimulatorCandidate is not None assert SimulatorPlan is not None @@ -128,6 +131,9 @@ def test_new_namespace_imports_resolve(): assert callable(build_qmera_contraction_optimizer) assert callable(recommend_simulator) assert mera_module is not None + assert qmera_module is not None + assert qmera_module.QMeraBuilder is QMeraBuilder + assert qmera_module.QMeraBuilder is mera_module.QMeraBuilder assert SimpleUpdateGen is not None assert SymDMRG2 is not None assert SweepOptimizer is not None From 7863755b25783240b009f36f3162e4823ca808cb Mon Sep 17 00:00:00 2001 From: rezaquant Date: Wed, 29 Jul 2026 18:31:07 -0600 Subject: [PATCH 23/70] Add scalable native MPO contraction --- docs/api/operators/hamiltonians.md | 31 + docs/api/optimizers/energy.md | 11 + docs/api/optimizers/mpo.md | 78 +++ docs/api/tensors/symmetric.md | 56 +- src/pepsy/operators/hamiltonians.py | 69 +- src/pepsy/optimizers/energy/peps.py | 139 +++- src/pepsy/optimizers/mpo/optimizer.py | 245 ++++++- src/pepsy/tensors/symm_fermions.py | 9 +- src/pepsy/tensors/symmetric.py | 921 +++++++++++++++++++++++--- tests/test_optimize_mpo.py | 325 +++++++++ tests/test_symmetric_tensors.py | 321 +++++++++ 11 files changed, 2087 insertions(+), 118 deletions(-) diff --git a/docs/api/operators/hamiltonians.md b/docs/api/operators/hamiltonians.md index 3dfadf2..e9a9d42 100644 --- a/docs/api/operators/hamiltonians.md +++ b/docs/api/operators/hamiltonians.md @@ -1,4 +1,35 @@ # `pepsy.operators.hamiltonians` +`ham_tn.build_mpo` retains its original explicit local-operator form. It also +accepts a `Fermion` model for symmetry-aware MPO construction: + +```python +fermion = pepsy.Fermion(spinful=True, symmetry="U1U1") +builder = pepsy.ham_tn(Lx=3, Ly=1) +mpo = builder.build_mpo( + fermion=fermion, + edges=[(0, 1), (1, 2)], + t=1.0, + U=2.0, + mu=0.1, +) +``` + +The native model-facing shorthand is +`fermion.to_mpo(edges, L=3, t=..., U=..., mu=...)`. Couplings remain explicit; +they are not stored on the `Fermion` object. Use `fermion.build_mpo(...)` when +the Jordan-Wigner-compatible MPO convention is wanted. + +`Fermion.to_mpo(...)` and `SymHamiltonian.to_mpo(..., fermionic=True)` return +native graded `FermionicArray` MPO tensors. Explicit mappings can contain +arbitrary neutral multi-site terms; non-contiguous supports are represented by +charged virtual channels. `ham_tn.build_mpo(..., fermionic=True)` selects the +same native path. Pass `to_backend=...` to map the stored Symmray blocks to a +selected array backend. + +Native MPO assembly, replay, and exact energy measurement are supported. The +native energy path applies the MPO sitewise as a factorized graded MPO-MPS +contraction, so it does not materialize the global physical operator. Its cost +is controlled by the MPS and MPO bond dimensions. > API details are maintained as handwritten Markdown in this page. diff --git a/docs/api/optimizers/energy.md b/docs/api/optimizers/energy.md index 432c4a7..d74b4f5 100644 --- a/docs/api/optimizers/energy.md +++ b/docs/api/optimizers/energy.md @@ -30,6 +30,17 @@ the automatic re-encoding can create very large block-sparse intermediates. For a deliberately small or explicitly managed conversion, pass ``allow_encoding_conversion=True`` to ``MpsEnergyOptimizer``. +An MPO returned by ``Fermion.to_mpo(...)`` is native graded and can be measured +directly with a native fermionic MPS. Pepsy applies that MPO sitewise as a +factorized graded MPO-MPS network, preserving Symmray's contraction order +without materializing an exponentially sized operator. + +Repeated native-MPO evaluations reuse a per-optimizer cotengra path cache. The +default is uncompressed and exact. For a controlled approximation, pass for +example ``native_mpo_compression={"max_bond": 64, "cutoff": 1e-12, +"method": "svd"}``; always compare against an uncompressed result when setting +the truncation cap. + ## Tree tensor networks ``TreeEnergyOptimizer`` mirrors the ``MpsEnergyOptimizer`` measurement surface diff --git a/docs/api/optimizers/mpo.md b/docs/api/optimizers/mpo.md index 1a781ca..385c5ca 100644 --- a/docs/api/optimizers/mpo.md +++ b/docs/api/optimizers/mpo.md @@ -1,4 +1,82 @@ # `pepsy.optimizers.mpo.optimizer` +`MpoOptimizer` accepts ordinary Quimb MPOs and Symmray block-sparse MPOs. For +a native graded fermion workflow, use `Fermion.to_mpo(...)` and replay the +matching native gate stream: + +```python +fermion = pepsy.Fermion(spinful=True, symmetry="U1U1") +edges = [(0, 1), (1, 2)] +hamiltonian = fermion.hamiltonian(edges, t=1.0, U=2.0, mu=0.1) +mpo = fermion.to_mpo(hamiltonian=hamiltonian, L=3) + +opt = pepsy.MpoOptimizer( + mpo, + hamiltonian.trotter_gates(0.01), + chi=16, + mode="svd", +) +mpo = opt.run(progbar=False) +``` + +For an explicit neutral term collection, arbitrary one- or multi-site support +is accepted: + +```python +term = fermion.operator_term( + [(1.0, ((0, "create_up"), (2, "annihilate_up")))], + sites=(0, 2), + add_hc=True, +) +mpo = fermion.to_mpo({(0, 2): term}, L=3) +``` + +The Jordan-Wigner compatibility path remains available through +`Fermion.build_mpo(...)` and the matching +`SymHamiltonian.jw_trotter_gates(...)` stream: + +```python +fermion = pepsy.Fermion(spinful=True, symmetry="U1U1") +edges = [(0, 1), (1, 2)] +hamiltonian = fermion.hamiltonian(edges, t=1.0, U=2.0, mu=0.1) +mpo = fermion.build_mpo(edges, L=3, t=1.0, U=2.0, mu=0.1) + +opt = pepsy.MpoOptimizer( + mpo, + hamiltonian.jw_trotter_gates(0.01), + chi=16, + mode="svd", +) +mpo = opt.run(progbar=False) +``` + +Symmray gates keep their charge and dual metadata and are not coerced to dense +arrays. Native graded gates from `Fermion.strang_gate_stream(...)` are accepted +as well; even gates are adapted explicitly to the MPO's Jordan-Wigner +convention. `mode="svd"` is the symmetry-aware compression path. For a +Symmray MPO, `mode="mpo"` and `mode="dmrg"` use that same block-aware path +because generic dense auxiliary MPO compression and bond padding do not support +multi-sector Symmray bonds reliably. + +Native MPO tensors retain their Symmray graded metadata throughout replay and +compression. `mode="svd"`, `mode="mpo"`, and `mode="dmrg"` use the same +block-aware path for native Symmray MPOs; the optimizer does not require a +dense conversion of the input MPO. +`ham_tn.build_mpo(..., fermionic=True)` is also routed to the native +`Fermion.to_mpo(...)` entry point. Use `to_backend=...` on the model-facing +builder when the stored blocks must be moved to Torch or another supported +backend. + +Native MPO assembly/replay is also measurable with a native fermionic MPS. +`MpsEnergyOptimizer` applies the native MPO sitewise as a factorized graded +MPO-MPS network, preserving Symmray ordering while retaining MPO bond scaling. +Repeated evaluations reuse the optimizer's contraction paths. Optional +controlled truncation is available through +``native_mpo_compression={"max_bond": ..., "cutoff": ..., "method": "svd"}``; +the default remains exact and uncompressed. + +To compress an existing MPO without replaying gates, use +`MpoOptimizer(mpo, gates=[], chi=...).compress()`. Symmray may retain a small +sector-multiplicity overshoot above the requested numeric bond cap. > API details are maintained as handwritten Markdown in this page. diff --git a/docs/api/tensors/symmetric.md b/docs/api/tensors/symmetric.md index 743f7c6..d6fc201 100644 --- a/docs/api/tensors/symmetric.md +++ b/docs/api/tensors/symmetric.md @@ -178,6 +178,11 @@ spinful.onsite_gate(dt=0.01, site=0, U=8.0) spinful.gate("interaction", dt=0.01, U=8.0) spinless.gate_stream(edges, dt=0.01, order=2, t=1.0, V=0.5, mu=0.0) spinful.local_terms(edges, t=1.0, U=8.0) # native terms for optimization +spinful.strang_gate_stream( + edges, dt=0.01, t=1.0, U=8.0, field_z=0.2 +) +pairing = py.Fermion(spinful=False, symmetry="Z2") +pairing.strang_gate_stream(edges, dt=0.01, t=1.0, pairing=0.2) ``` The bare ``*_operator`` methods return explicit native fermionic operators, @@ -308,9 +313,17 @@ the native Symmray gate directly. ``imaginary=True`` changes the evolution to ``exp(-dt H)``. The local exponential must be neutral so it remains in one conserved charge sector. -``SpinfulFermion`` and ``SpinfulFermionHubbard`` remain compatibility aliases -for ``Fermion``. ``SymmFermions.spinless(...)`` and -``SymmFermions.spinful(...)`` are factory-style alternatives. +``SpinfulFermion`` and ``SpinfulFermionHubbard`` remain compatibility +constructors that fix ``spinful=True``. ``SymmFermions.spinless(...)`` and +``SymmFermions.spinful(...)`` are factory-style alternatives with the same +local-space guarantees. + +Automatic streams and edge-built Hamiltonians also accept ``field_x``, +``field_y``, and ``field_z``. Transverse fields mix up/down occupation and +therefore require spinful total-``U1`` or ``Z2`` symmetry; a longitudinal +``field_z`` also works with spin-resolved ``U1U1``/``Z2Z2``. The ``pairing`` +and ``pairing_phase`` options describe a parity-preserving spinless pairing +term and require ``Fermion(spinful=False, symmetry="Z2")``. ## Native spinful-fermion helper @@ -507,6 +520,8 @@ Current support is: - spinful Fermi-Hubbard ``model="fermi_hubbard_u1u1"`` with ``symmetry="U1U1"``, hopping, onsite interaction, nearest-neighbor density interaction, and chemical-potential terms. +- native graded MPO conversion for arbitrary neutral one- or multi-site + ``FermionicArray`` terms, including non-contiguous support; Spinful total-particle-number ``model="fermi_hubbard"`` with ``symmetry="U1"`` still raises ``NotImplementedError``; use ``model="fermi_hubbard_u1u1"`` when @@ -553,6 +568,41 @@ mpo = ham.to_mpo(mapper=mapper) mpo = ham.to_mpo(idx2coo=idx2coo, coo2idx=coo2idx) ``` +For the common edge-built Fermi-Hubbard path, the model helper is equivalent +and keeps the selected symmetry attached to the construction: + +```python +fermion = py.Fermion(spinful=True, symmetry="U1U1") +mpo = fermion.build_mpo( + [(0, 1), (1, 2)], L=3, t=1.0, U=8.0, mu=0.0, max_bond=16 +) +``` + +This returns the existing symmetry-preserving Jordan-Wigner-compatible MPO +convention. For the native graded path, use ``to_mpo``: + +```python +mpo = fermion.to_mpo( + [(0, 1), (1, 2)], L=3, t=1.0, U=8.0, mu=0.0, max_bond=16 +) +assert all(type(tensor.data).__name__ == "U1U1FermionicArray" for tensor in mpo) +``` + +Native fermionic gate streams from the same ``Fermion`` model can be passed to +``MpoOptimizer``; the optimizer preserves the graded Symmray tensors and their +charge blocks during replay and compression. ``fermionic=False`` remains the +explicit Jordan-Wigner compatibility choice for ``SymHamiltonian.to_mpo``. + +Pass ``to_backend=`` to ``Fermion.to_mpo`` or ``Fermion.build_mpo`` when the +returned Symmray blocks should use a selected array backend. Native MPO +assembly, replay, and exact energy measurement are supported. Native MPO +energy applies the operator sitewise as a factorized graded MPO-MPS +contraction, so its cost is controlled by the MPS and MPO bond dimensions. + +``Fermion.build_mpo(..., fermionic=True)`` selects the same native construction +as ``Fermion.to_mpo(...)``. Omitting ``fermionic=True`` retains the explicit +Jordan--Wigner compatibility MPO path. + For periodic square lattices encoded as long-range edges in an OBC MPS/MPO, ``mode="folded-snake"`` alternates opposite columns before snaking. On a 6 by 6 torus this lowers the longest nearest-neighbor chain separation from 35 diff --git a/src/pepsy/operators/hamiltonians.py b/src/pepsy/operators/hamiltonians.py index c9b4020..7aed622 100644 --- a/src/pepsy/operators/hamiltonians.py +++ b/src/pepsy/operators/hamiltonians.py @@ -465,7 +465,7 @@ def _zero_mpo(self, *, phys_dim, dtype): def build_mpo( self, - ints, + ints=None, *, phys_dim=2, max_bond=None, @@ -473,12 +473,16 @@ def build_mpo( data_type=None, compress_each=True, mapper=None, + fermion=None, + edges=None, + fermionic=None, + **model_params, ): """Build MPO from user interactions. Parameters ---------- - ints : sequence + ints : sequence | Mapping | pepsy.Fermion | None Sequence of terms. Supported term formats: - ``((op,), (coord,))`` - ``((op1, op2), (coord1, coord2))`` @@ -500,12 +504,73 @@ def build_mpo( mapper : pepsy.tensors.core.OneDMap | None, default=None Optional mapper override used only for this MPO build. When omitted, the builder's configured mapper is used. + fermion : pepsy.Fermion | None, default=None + Optional native fermion model. When supplied, ``ints`` (or the + explicit ``edges`` alias) is passed to ``fermion.build_mpo`` and + the returned Symmray MPO keeps the model's U1/U1U1 symmetry. + edges : sequence | None, default=None + Explicit edge alias for the ``fermion=...`` form. For example, + ``builder.build_mpo(fermion=f, edges=edges, t=..., U=...)``. + fermionic : bool | None, default=None + Native graded encoding flag for the fermion-model form. ``None`` + and ``False`` select the Jordan-Wigner-compatible MPO builder; + ``True`` selects ``Fermion.to_mpo(...)``. + **model_params + Explicit fermion couplings such as ``t``, ``U``/``V``, and ``mu``. Returns ------- qtn.MatrixProductOperator Built Hamiltonian MPO. """ + if ( + fermion is None + and hasattr(ints, "build_mpo") + and hasattr(ints, "hamiltonian") + ): + fermion = ints + ints = None + + if fermion is not None: + if edges is not None: + if ints is not None: + raise TypeError( + "Pass fermion terms through either ints or edges, not both." + ) + ints = edges + if ints is None: + raise ValueError( + "Fermion MPO construction requires terms or an edge sequence." + ) + if not isinstance(phys_dim, Integral) or int(phys_dim) < 1: + raise ValueError("phys_dim must be an integer >= 1.") + if not hasattr(fermion, "build_mpo"): + raise TypeError( + "fermion must provide the Fermion.build_mpo interface." + ) + dtype = self.data_type if data_type is None else np.dtype(data_type) + max_bond_use = self.max_bond if max_bond is None else int(max_bond) + cutoff_use = self.cutoff if cutoff is None else float(cutoff) + mapper_use = self.mapper if mapper is None else mapper + fermionic_use = False if fermionic is None else bool(fermionic) + mpo_builder = fermion.to_mpo if fermionic_use else fermion.build_mpo + return mpo_builder( + ints, + L=self.L, + mapper=mapper_use, + max_bond=max_bond_use, + cutoff=cutoff_use, + compress=bool(compress_each), + dtype=dtype, + fermionic=fermionic_use, + **model_params, + ) + + if edges is not None or model_params or fermionic is not None: + raise TypeError( + "edges, fermion model parameters, and fermionic encoding are " + "only valid with fermion=... ." + ) if ints is None: raise ValueError("ints must be provided.") if not isinstance(phys_dim, Integral) or int(phys_dim) < 1: diff --git a/src/pepsy/optimizers/energy/peps.py b/src/pepsy/optimizers/energy/peps.py index 673d0e3..92803cc 100644 --- a/src/pepsy/optimizers/energy/peps.py +++ b/src/pepsy/optimizers/energy/peps.py @@ -774,13 +774,19 @@ class MpsEnergyOptimizer(PepsEnergyOptimizer): The objective is the normalized local expectation value ``/``. Local-term Hamiltonians are evaluated with - ``MPS.compute_local_expectation_exact(...)``. MPO Hamiltonians are evaluated - directly as ``().contract(all, optimize=...)``, using - ``contraction_opt`` for the full network contraction. Native fermionic - Symmray states use native local terms by default when a mapped - ``SymHamiltonian`` is supplied. A bosonic/Jordan-Wigner Symmray MPO cannot - be silently contracted with a native fermionic state, since the required - re-encoding can create very large block contractions; pass + ``MPS.compute_local_expectation_exact(...)``. Bosonic MPO Hamiltonians are + evaluated directly as ``().contract(all, + optimize=...)``, using ``contraction_opt`` for the full network + contraction. Native fermionic MPOs are applied sitewise as a factorized + MPO-MPS network, which preserves Symmray's graded contraction ordering + without materializing a global operator. Repeated native-MPO evaluations + reuse a per-optimizer contraction path cache. Optional compression can be + requested with ``native_mpo_compression={"max_bond": ..., "cutoff": ...}``; + it is disabled by default so the energy remains exact. + Native fermionic Symmray states use native local terms by default when a + mapped ``SymHamiltonian`` is supplied. A bosonic/Jordan-Wigner Symmray MPO + cannot be silently contracted with a native fermionic state, since the + required re-encoding can create very large block contractions; pass ``allow_encoding_conversion=True`` to explicitly request that conversion. Hamiltonians can be supplied as a ``qtn.MatrixProductOperator``, a ``qtn.LocalHam1D``-like object with ``.terms``, a Pepsy symmetric @@ -798,6 +804,7 @@ class MpsEnergyOptimizer(PepsEnergyOptimizer): "compute_kwargs", "progbar", "allow_encoding_conversion", + "native_mpo_compression", }) def __init__( @@ -814,6 +821,7 @@ def __init__( compute_kwargs: Mapping[str, Any] | None = None, loss_kwargs: Mapping[str, Any] | None = None, allow_encoding_conversion: bool = False, + native_mpo_compression: Mapping[str, Any] | None = None, ): if hamiltonian is not None and terms is not None: raise TypeError("pass either hamiltonian or terms, not both") @@ -830,7 +838,13 @@ def __init__( "progbar": progbar, "compute_kwargs": {} if compute_kwargs is None else dict(compute_kwargs), "allow_encoding_conversion": bool(allow_encoding_conversion), + "native_mpo_compression": ( + None + if native_mpo_compression is None + else dict(native_mpo_compression) + ), } + self._native_mpo_path_optimizer = None if loss_kwargs is not None: self.set_loss_kwargs(**loss_kwargs) @@ -869,6 +883,25 @@ def _as_mps_state(cls, state): "compute_local_expectation_exact()." ) + def _prepare_native_mpo_options(self, state, terms, options): + """Reuse one cotengra optimizer for repeated native MPO losses.""" + options = dict(options) + if not self._is_mpo_hamiltonian(terms): + return options + if self._symmray_encoding(state) != "native_fermionic": + return options + if self._symmray_encoding(terms) != "native_fermionic": + return options + + contraction_opt = options.get("contraction_opt", "auto-hq") + if contraction_opt is None or contraction_opt == "auto-hq": + if self._native_mpo_path_optimizer is None: + self._native_mpo_path_optimizer = build_optimizer( + progbar=bool(options.get("progbar", False)), + ) + options["contraction_opt"] = self._native_mpo_path_optimizer + return options + @staticmethod def _is_mpo_hamiltonian(hamiltonian): return isinstance(hamiltonian, qtn.MatrixProductOperator) @@ -916,6 +949,65 @@ def _symmray_encoding(cls, value): def _mpo_uses_bosonic_symmray(cls, mpo): return cls._symmray_encoding(mpo) == "bosonic_symmray" + @classmethod + def _native_mpo_expectation( + cls, + state, + mpo, + *, + normalized=True, + contraction_opt="auto-hq", + native_mpo_compression=None, + ): + """Evaluate a native graded MPO through a factorized MPO-MPS network. + + Symmray fermionic contractions are order-sensitive. A conventional + MPO sandwich allows Quimb's path optimizer to interleave local MPO + factors with bra and ket tensors, which can change the graded phase. + Applying the MPO sitewise first keeps each local graded contraction + together and leaves the operator bond factorized, so the contraction + scales with the MPS and MPO bond dimensions rather than the global + Hilbert-space dimension. + """ + gated = qtn.tensor_network_apply_op_vec( + mpo, + state, + which_A="lower", + contract=True, + fuse_multibonds=True, + compress=False, + inplace=False, + inplace_A=False, + ) + if native_mpo_compression is not None: + compression_opts = dict(native_mpo_compression) + max_bond = compression_opts.get("max_bond") + if max_bond is None: + raise ValueError( + "native_mpo_compression requires an explicit max_bond." + ) + max_bond = int(max_bond) + if max_bond < 1: + raise ValueError( + "native_mpo_compression max_bond must be positive." + ) + cutoff = compression_opts.get("cutoff", 1e-12) + if cutoff < 0.0: + raise ValueError( + "native_mpo_compression cutoff must be non-negative." + ) + compression_opts["max_bond"] = max_bond + compression_opts["cutoff"] = cutoff + compression_opts.setdefault("method", "svd") + gated.compress(**compression_opts) + value = (state.H | gated).contract(all, optimize=contraction_opt) + if normalized: + norm = (state.H & state).contract(all, optimize=contraction_opt) + if norm == 0.0: + raise ValueError("Cannot compute normalized energy for a zero-norm state.") + value = value / norm + return value + @staticmethod def _symmray_symmetry_name(data): symmetry = getattr(data, "symmetry", None) @@ -1306,6 +1398,7 @@ def _mpo_expectation( normalized=True, contraction_opt="auto-hq", allow_encoding_conversion=False, + native_mpo_compression=None, ): if contraction_opt is None: contraction_opt = build_optimizer(progbar=False) @@ -1325,6 +1418,24 @@ def _mpo_expectation( "pass `allow_encoding_conversion=True` to explicitly request " "the potentially memory-intensive re-encoding." ) + if mpo_encoding == "native_fermionic": + if state_encoding != "native_fermionic": + raise ValueError( + "Native fermionic MPO energy evaluation requires a native " + "fermionic Symmray MPS state." + ) + return cls._native_mpo_expectation( + state, + mpo, + normalized=normalized, + contraction_opt=contraction_opt, + native_mpo_compression=native_mpo_compression, + ) + if native_mpo_compression is not None: + raise ValueError( + "native_mpo_compression is only supported for native " + "fermionic Symmray MPOs." + ) if cls._mpo_uses_bosonic_symmray(mpo): ket = cls._bosonize_fermionic_tn(state) else: @@ -1355,6 +1466,7 @@ def _loss_state( compute_kwargs=None, progbar=False, allow_encoding_conversion=False, + native_mpo_compression=None, ): state = cls._as_mps_state(state) terms = cls._terms_from_hamiltonian(terms) @@ -1370,6 +1482,7 @@ def _loss_state( normalized=normalized, contraction_opt=contraction_opt, allow_encoding_conversion=allow_encoding_conversion, + native_mpo_compression=native_mpo_compression, ) if energy_per_site: value = value / cls._num_sites(state) @@ -1410,6 +1523,7 @@ def loss(self, state=None, *, hamiltonian=None, terms=None, **kwargs): if terms is not None: terms_use = self._terms_from_hamiltonian(terms) opts = self._merge_opts(self.loss_kwargs, self._pick_loss_kwargs(kwargs)) + opts = self._prepare_native_mpo_options(state, terms_use, opts) return self._loss_state(state, terms=terms_use, **opts) def energy(self, state=None, *, hamiltonian=None, terms=None, **kwargs): @@ -1422,6 +1536,7 @@ def energy(self, state=None, *, hamiltonian=None, terms=None, **kwargs): terms_use = self._terms_from_hamiltonian(terms) opts = self._merge_opts(self.loss_kwargs, self._pick_loss_kwargs(kwargs)) + opts = self._prepare_native_mpo_options(state, terms_use, opts) opts_full = dict(opts) opts_full["energy_per_site"] = False energy = self._loss_state(state, terms=terms_use, **opts_full) @@ -1439,6 +1554,11 @@ def energy(self, state=None, *, hamiltonian=None, terms=None, **kwargs): "contraction_opt": opts["contraction_opt"], "progbar": opts["progbar"], "compute_kwargs": dict(opts["compute_kwargs"]), + "native_mpo_compression": ( + None + if opts["native_mpo_compression"] is None + else dict(opts["native_mpo_compression"]) + ), }, ) @@ -1477,6 +1597,11 @@ def make_tn_optimizer( terms = terms.terms else: terms = self._fermionic_hamiltonian_mpo_for_state(terms, self.state) + merged_loss_kwargs = self._prepare_native_mpo_options( + self.state, + terms, + merged_loss_kwargs, + ) if self._is_mpo_hamiltonian(terms): constants = {"terms": terms} constants.update(incoming_constants) diff --git a/src/pepsy/optimizers/mpo/optimizer.py b/src/pepsy/optimizers/mpo/optimizer.py index 4d2561f..af22cc6 100644 --- a/src/pepsy/optimizers/mpo/optimizer.py +++ b/src/pepsy/optimizers/mpo/optimizer.py @@ -20,7 +20,8 @@ * ``mode="svd"`` — apply the gate with ``reduce-split`` then canonicalize + left-compress to ``chi``; * ``mode="mpo"`` — use :func:`pepsy.operators.gates.gate_nonlocal_opt` to - apply each layer independently on the ket and bra families. + apply each layer independently on the ket and bra families. Symmray MPOs + use the block-aware SVD path instead. The class also tracks a running "normalized-norm" proxy ``sqrt( / )`` that equals ``1`` for purely unitary two-sided @@ -47,7 +48,15 @@ def _normalize_gate_queue(gates): if not entries: return [], [] gate_list, where_list = zip(*entries) - return list(gate_list), [tuple(w) if isinstance(w, list) else w for w in where_list] + + def normalize_where(where): + if isinstance(where, Integral): + return (int(where),) + if isinstance(where, list): + return tuple(where) + return where + + return list(gate_list), [normalize_where(where) for where in where_list] class MpoOptimizer: @@ -102,6 +111,26 @@ class MpoOptimizer: _ALLOWED_MODES = frozenset({"dmrg", "svd", "mpo"}) + @staticmethod + def _is_symmray_array(value): + """Return whether ``value`` is a Symmray block-sparse array.""" + return hasattr(value, "blocks") and hasattr(value, "indices") + + @staticmethod + def _is_fermionic_array(value): + """Return whether ``value`` carries Symmray graded-array metadata.""" + return bool(getattr(value, "fermionic", False)) or ( + "fermionicarray" in type(value).__name__.lower() + ) + + @classmethod + def _has_symmray_data(cls, tn): + """Return whether any tensor in ``tn`` stores Symmray data.""" + return any( + cls._is_symmray_array(getattr(tensor, "data", None)) + for tensor in getattr(tn, "tensors", ()) + ) + @classmethod def _normalize_mode(cls, mode): """Lower-case and validate ``mode`` against :attr:`_ALLOWED_MODES`.""" @@ -295,6 +324,13 @@ def _prepare_gate_tensor(gate, n_sites): ``(o1, o2, i1, i2)`` tensors for two-site gates). ``apply_gate`` below expects the opposite ordering, so we transpose accordingly. """ + # Native Symmray gates already carry explicit dual metadata describing + # output and input legs. Transposing them as if they were dense Quimb + # gates changes the charge sectors (and can make a U1U1 split ask for + # impossible virtual charges). Keep their graded/block structure intact. + if MpoOptimizer._is_symmray_array(gate): + return gate + if n_sites == 1: return ar.do("transpose", gate, (1, 0)) elif n_sites == 2: @@ -317,8 +353,93 @@ def _prepare_gate_tensor(gate, n_sites): else: raise ValueError("Each gate location must have one or two sites.") + @classmethod + def _symmray_physical_map(cls, p, site, ind_id): + """Return the dense charge list for one live MPO physical index.""" + ind = ind_id.format(site) + tensor = next( + tensor + for tensor in getattr(p, "tensors", ()) + if ind in getattr(tensor, "inds", ()) + ) + axis = tensor.inds.index(ind) + chargemap = tensor.data.indices[axis].chargemap + return [charge for charge, size in chargemap.items() for _ in range(int(size))] + @staticmethod - def _prepare_gate_pair(gate, n_sites, bra_gate=None): + def _charge_parity(charge): + """Return fermion parity inferred from a scalar or product charge.""" + if isinstance(charge, tuple): + return sum(int(part) for part in charge) % 2 + return int(charge) % 2 + + @classmethod + def _fermionic_gate_to_bosonic(cls, p, gate, where, ind_id): + """Convert an even native gate for a non-graded Symmray MPO. + + ``SymHamiltonian.to_mpo`` deliberately returns a bosonic Symmray MPO + with the Jordan-Wigner parity convention already encoded in its local + channels. Native ``Fermion`` gates use Symmray's graded tensor-product + convention instead. For an even two-site gate, changing conventions + amounts to the endpoint crossing phase on the input ket sectors. + """ + if not cls._is_fermionic_array(gate) or not cls._has_symmray_data(p): + return gate + + sample = next( + tensor.data + for tensor in getattr(p, "tensors", ()) + if cls._is_symmray_array(getattr(tensor, "data", None)) + ) + if cls._is_fermionic_array(sample): + return gate + + try: + dense = gate.to_dense() + except AttributeError: + dense = gate + dense = np.asarray(dense) + where = tuple(where) + physical_maps = [ + cls._symmray_physical_map(p, site, ind_id) for site in where + ] + n_sites = len(where) + if n_sites == 2: + left_odd = np.array( + [cls._charge_parity(charge) for charge in physical_maps[0]], + dtype=bool, + ) + right_odd = np.array( + [cls._charge_parity(charge) for charge in physical_maps[1]], + dtype=bool, + ) + crossing = np.ones( + (len(left_odd), len(right_odd)), + dtype=dense.dtype, + ) + crossing[np.ix_(left_odd, right_odd)] = -1 + dense = dense * crossing[None, None, :, :] + + import symmray.utils as sr_utils # pylint: disable=import-outside-toplevel + + sample_charge = next(iter(physical_maps[0]), 0) + zero = ( + tuple(0 for _ in sample_charge) + if isinstance(sample_charge, tuple) + else 0 + ) + return sr_utils.from_dense( + dense, + symmetry=getattr(sample, "symmetry", None), + index_maps=physical_maps * 2, + duals=(False,) * n_sites + (True,) * n_sites, + fermionic=False, + charge=zero, + ) + + @classmethod + def _prepare_gate_pair(cls, gate, n_sites, bra_gate=None, *, p=None, where=None, + ind_id="k{}"): """Return ``(g_k, g_b)`` ready to be fed to :func:`apply_gate`. ``gate`` becomes ``g_k`` (acts on the ket index family with the @@ -330,13 +451,50 @@ def _prepare_gate_pair(gate, n_sites, bra_gate=None): if gate is None and bra_gate is None: raise ValueError("At least one of ket gate or bra gate must be provided.") - g_k = None if gate is None else MpoOptimizer._prepare_gate_tensor(gate, n_sites) + if p is not None and where is not None: + gate = ( + None + if gate is None + else cls._fermionic_gate_to_bosonic(p, gate, where, ind_id) + ) + bra_gate = ( + None + if bra_gate is None + else cls._fermionic_gate_to_bosonic(p, bra_gate, where, ind_id) + ) + + g_k = None if gate is None else cls._prepare_gate_tensor(gate, n_sites) if bra_gate is None: g_b = None else: - g_b = ar.do("conj", MpoOptimizer._prepare_gate_tensor(bra_gate, n_sites)) + g_b = ar.do("conj", cls._prepare_gate_tensor(bra_gate, n_sites)) return g_k, g_b + @classmethod + def _materialize_split_gate(cls, p, where): + """Contract lazy native split-gate tensors back into site tensors. + + Quimb's ``split-gate`` application is deliberately lazy: for a + non-local gate it temporarily leaves an extra tensor at each endpoint + carrying the site tag. That is useful for general tensor-network + workflows, but ``MatrixProductOperator.right_canonize_site`` expects + exactly one tensor per site. Native Symmray arrays cannot use the + dense swap-gate fallback, so materialize the endpoint tensors before + the MPO canonicalization/compression sweep. + """ + if not cls._has_symmray_data(p): + return p + + for site in set(where): + tag = f"I{site}" + if len(p.tag_map.get(tag, ())) > 1: + p.contract_tags( + tag, + preserve_tensor=True, + inplace=True, + ) + return p + @staticmethod def _parse_gate_entry(G_i, where_i): """Decompose one stream entry into ``(ket_gate, bra_gate, where)``. @@ -387,7 +545,14 @@ def _apply_gate_pair( two index families stay decoupled. """ n_sites = len(where) - g_k, g_b = self._prepare_gate_pair(gate, n_sites, bra_gate=bra_gate) + g_k, g_b = self._prepare_gate_pair( + gate, + n_sites, + bra_gate=bra_gate, + p=p, + where=where, + ind_id=self.ind_id_k, + ) if g_k is not None: apply_gate( @@ -412,6 +577,9 @@ def _apply_gate_pair( inplace=inplace, ) + if contract == "split-gate": + self._materialize_split_gate(p, where) + def _build_dmrg_target(self, p, gate, where, bra_gate, cutoff, cutoff_mode="rsum2"): """Return ``p`` with one two-site gate pair applied via ``split-gate``. @@ -637,6 +805,11 @@ def _run_svd(self, G_seq, where_seq, progbar=False, cutoff=1e-12, cutoff_mode="r elif n_sites == 2: two_qubit_count += 1 xmin, xmax = sorted(where) + contract = ( + "split-gate" + if self._has_symmray_data(p) and (xmax - xmin > 1) + else "reduce-split" + ) self._apply_gate_pair( p, @@ -645,7 +818,7 @@ def _run_svd(self, G_seq, where_seq, progbar=False, cutoff=1e-12, cutoff_mode="r bra_gate=bra_gate, cutoff=cutoff, cutoff_mode=cutoff_mode, - contract="reduce-split", + contract=contract, inplace=True, ) @@ -722,7 +895,14 @@ def _run_mpo(self, G_seq, where_seq, progbar=False, cutoff=1e-12, cutoff_mode="r ) elif n_sites == 2: two_qubit_count += 1 - g_k, g_b = self._prepare_gate_pair(gate, n_sites, bra_gate=bra_gate) + g_k, g_b = self._prepare_gate_pair( + gate, + n_sites, + bra_gate=bra_gate, + p=p, + where=where, + ind_id=self.ind_id_k, + ) if g_k is not None: p = gate_nonlocal_opt( p, g_k, where, @@ -803,6 +983,12 @@ def run( # pylint: disable=too-many-arguments,too-many-positional-arguments ------- qtn.MatrixProductOperator Updated MPO after replaying the queued gate stream. + + Notes + ----- + Symmray MPOs use the block-aware SVD compression implementation for + all three modes. This preserves multi-sector bonds when Quimb's + generic dense auxiliary or bond-padding paths are unavailable. """ if mode is not None: self.set_mode(mode) @@ -813,6 +999,20 @@ def run( # pylint: disable=too-many-arguments,too-many-positional-arguments return self.p if self.mode == "dmrg": + # Quimb's FIT preparation expands bonds with dense-style padding, + # which is not valid for multi-sector Symmray indices. The + # symmetry-aware local SVD path applies the same gate/compression + # contract without destroying charge blocks. + if self._has_symmray_data(self.p): + self._run_svd( + G_seq, + where_seq, + progbar=progbar, + cutoff=cutoff, + cutoff_mode=cutoff_mode, + fidelity_samples=fidelity_samples, + ) + return self.p self._prepare_dmrg_state() self._run_dmrg( G_seq, @@ -838,6 +1038,19 @@ def run( # pylint: disable=too-many-arguments,too-many-positional-arguments return self.p if self.mode == "mpo": + # ``gate_nonlocal_opt`` creates a dense auxiliary sub-MPO and its + # generic compression currently loses multi-sector Symmray bond + # metadata. Reuse the block-aware local SVD route for these MPOs. + if self._has_symmray_data(self.p): + self._run_svd( + G_seq, + where_seq, + progbar=progbar, + cutoff=cutoff, + cutoff_mode=cutoff_mode, + fidelity_samples=fidelity_samples, + ) + return self.p self._run_mpo( G_seq, where_seq, @@ -879,3 +1092,19 @@ def canonize_mpo(self, p, where): def get_fidelities(self): """Return the running loss history.""" return self.losses + + def compress(self, *, cutoff=1e-12, cutoff_mode="rsum2"): + """Compress the current MPO to ``chi`` while preserving its backend. + + Symmray MPOs use Quimb's local block-aware compression path. This is + also useful when the optimizer is constructed with an empty gate + queue and the caller only wants a bond-dimension reduction. + """ + self.p.compress( + form="left", + max_bond=self.chi, + cutoff=cutoff, + cutoff_mode=cutoff_mode, + ) + self._init_canonicalization() + return self.p diff --git a/src/pepsy/tensors/symm_fermions.py b/src/pepsy/tensors/symm_fermions.py index e829ee1..49dadbf 100644 --- a/src/pepsy/tensors/symm_fermions.py +++ b/src/pepsy/tensors/symm_fermions.py @@ -2,14 +2,11 @@ The underlying Symmray adapters live in :mod:`pepsy.tensors.symmetric`. This module is the model-facing home for convenient fermionic building blocks, so -future spinless or multicomponent helpers can live alongside ``SpinfulFermion`` +future multicomponent helpers can live alongside the unified ``Fermion`` without making the symmetry implementation module into a model catalogue. """ -from .symmetric import Fermion - -SpinfulFermion = Fermion -SpinfulFermionHubbard = Fermion +from .symmetric import Fermion, SpinfulFermion, SpinfulFermionHubbard __all__ = [ "Fermion", @@ -34,7 +31,7 @@ def spinful(*args, **kwargs): This namespace makes it possible to add other local fermion spaces later while keeping the direct ``SpinfulFermion(...)`` form concise. """ - return Fermion(*args, **kwargs) + return SpinfulFermion(*args, **kwargs) @staticmethod def spinless(*args, **kwargs): diff --git a/src/pepsy/tensors/symmetric.py b/src/pepsy/tensors/symmetric.py index f3b9e09..bab2e0e 100644 --- a/src/pepsy/tensors/symmetric.py +++ b/src/pepsy/tensors/symmetric.py @@ -4619,7 +4619,7 @@ def _term_mapping_uses_coordinate_sites(terms): def _as_term_where(where, *, coordinate_sites=False): - """Normalize a local-term location to a one- or two-site tuple.""" + """Normalize a local-term location to a one-or-more-site tuple.""" if coordinate_sites and _is_lattice_coordinate(where): # Preserve the coordinate as one site rather than interpreting (x, y) # as an MPS edge. The public mapping retains the flat coordinate key @@ -4629,9 +4629,9 @@ def _as_term_where(where, *, coordinate_sites=False): out = tuple(where) else: out = (where,) - if len(out) not in {1, 2}: + if not out: raise ValueError( - "Hamiltonian term locations must contain one site or two sites." + "Hamiltonian term locations must contain at least one site." ) return out @@ -5111,25 +5111,20 @@ def _is_fermionic_symmray_array(value): return "FermionicArray" in type(value).__name__ -def _first_term_block_sample(terms): +def _dtype_from_hamiltonian_terms(terms, default="complex128"): + dtypes = [] for term in dict(terms).values(): blocks = getattr(term, "blocks", None) - if blocks: - return next(iter(blocks.values())) - if hasattr(term, "dtype"): - return term - return None - - -def _dtype_from_hamiltonian_terms(terms, default="complex128"): - sample = _first_term_block_sample(terms) - dtype = getattr(sample, "dtype", None) - if dtype is None: - return np.dtype(default) - try: - return np.dtype(dtype) - except TypeError: - return np.dtype(default) + values = blocks.values() if blocks else (term,) + for value in values: + dtype = getattr(value, "dtype", None) + if dtype is None: + continue + try: + dtypes.append(np.dtype(dtype)) + except TypeError: + continue + return np.result_type(*dtypes) if dtypes else np.dtype(default) def _as_spin_pair(value, *, name): @@ -5161,6 +5156,13 @@ def _node_parameter(value, site): return value +def _coupling_is_active(value): + """Whether a scalar, site map, or edge map contributes to a model.""" + if callable(value) or isinstance(value, Mapping): + return True + return value != 0 + + def _dense_numpy(value, *, dtype=None): value = _to_dense(value) detach = getattr(value, "detach", None) @@ -5471,6 +5473,7 @@ def _assemble_symmray_mpo( lower_ind_id="b{}", site_tag_id="I{}", to_backend=None, + fermionic=False, ): channel_pos = [ {channel_id: pos for pos, (channel_id, _) in enumerate(cut_channels)} @@ -5536,7 +5539,7 @@ def _assemble_symmray_mpo( symmetry=symmetry, index_maps=index_maps, duals=duals, - fermionic=False, + fermionic=bool(fermionic), charge=zero, ) ) @@ -5550,7 +5553,8 @@ def _assemble_symmray_mpo( ) if to_backend is not None: _apply_to_tensor_network_arrays(mpo, to_backend) - raw_max_bond = int(mpo.max_bond()) + raw_bond = mpo.max_bond() + raw_max_bond = 1 if raw_bond is None else int(raw_bond) did_compress = bool(compress and L > 1) if compress and L > 1: compress_opts = {"cutoff": cutoff} @@ -5559,7 +5563,8 @@ def _assemble_symmray_mpo( mpo.compress(**compress_opts) requested_max_bond = None if max_bond is None else int(max_bond) - final_max_bond = int(mpo.max_bond()) + final_bond = mpo.max_bond() + final_max_bond = 1 if final_bond is None else int(final_bond) report = { "compressed": did_compress, "cutoff": cutoff, @@ -5741,6 +5746,154 @@ def add_channel(edge_pos, i, j, label, left_charge, left_op, right_op): ) +def _add_native_term_to_mpo( + term, + sites, + *, + term_pos, + channels, + transitions, + symmetry, + dtype, + zero, +): + """Add one neutral native fermion term as graded MPO transitions. + + The term is split by operator Schmidt decompositions over the ordered + support sites. This preserves Symmray's fermionic bond phases while + allowing arbitrary term rank and non-contiguous support. + """ + sites = tuple(int(site) for site in sites) + if len(set(sites)) != len(sites): + raise ValueError("Hamiltonian term supports must contain unique sites.") + if any(site < 0 or site >= len(transitions) for site in sites): + raise ValueError(f"Hamiltonian term support {sites!r} is outside MPO bounds.") + start = ("start",) + done = ("done",) + + order = tuple(sorted(range(len(sites)), key=sites.__getitem__)) + sites = tuple(sorted(sites)) + indices = getattr(term, "indices", None) + if indices is None or len(indices) == 0 or len(indices) % 2: + raise TypeError( + "SymHamiltonian.to_mpo requires an even-rank Symmray operator." + ) + n_sites = len(sites) + if len(indices) != 2 * n_sites: + raise ValueError("Term metadata and support size disagree.") + + output_maps = [_expanded_index_charges(index) for index in indices[:n_sites]] + input_maps = [_expanded_index_charges(index) for index in indices[n_sites:]] + if output_maps != input_maps: + raise ValueError( + "Hamiltonian terms must use matching upper/lower physical charge " + "maps at every site." + ) + if any(site_map != output_maps[0] for site_map in output_maps[1:]): + raise ValueError( + "SymHamiltonian.to_mpo currently requires one physical charge " + "map shared by all sites." + ) + physical_maps = [output_maps[pos] for pos in order] + + term_charge = _normalize_group_charge( + getattr(term, "charge", zero), + symmetry, + ) + if term_charge != zero: + raise ValueError( + "MPO Hamiltonian terms must be neutral under the selected " + f"symmetry; term {term_pos} has charge {term_charge!r}." + ) + + if n_sites == 1: + dense = _dense_numpy(term, dtype=dtype) + for out_i, in_i in product(range(len(physical_maps[0])), repeat=2): + coefficient = dense[out_i, in_i] + if not np.any(coefficient): + continue + local = np.zeros((len(physical_maps[0]), len(physical_maps[0])), dtype=dtype) + local[out_i, in_i] = coefficient + transitions[sites[0]].append((start, done, local)) + return list(physical_maps[0]) + + axes = order + tuple(n_sites + pos for pos in order) + ordered_term = term if order == tuple(range(n_sites)) else term.transpose(axes) + local_term = ordered_term.fuse( + *((pos, n_sites + pos) for pos in range(n_sites)) + ) + + factors = [] + current = local_term + for pos in range(n_sites - 1): + ndim = current.ndim + if pos == 0: + left_group = (0,) + right_group = tuple(range(1, ndim)) + else: + left_group = (0, 1) + right_group = tuple(range(2, ndim)) + left, _, right = current.fuse(left_group, right_group).svd( + absorb="right" + ) + if pos == 0: + factors.append(left.unfuse(0).transpose((2, 0, 1))) + else: + factors.append( + left.unfuse(0).unfuse(1).transpose((0, 3, 1, 2)) + ) + current = right.unfuse(1) + factors.append(current) + + # Each operator-Schmidt bond becomes a family of MPO channels. The + # channel charge is taken directly from the native factor bond index, + # rather than reconstructed from dense matrix elements. + interval_channel_ids = [] + for interval in range(n_sites - 1): + factor = factors[interval] + bond_axis = 0 if interval == 0 else 1 + bond_map = _expanded_index_charges(factor.indices[bond_axis]) + channel_ids = [] + for bond_pos, bond_charge in enumerate(bond_map): + channel_id = ("native", term_pos, interval, bond_pos) + channel_ids.append(channel_id) + for cut in range(sites[interval], sites[interval + 1]): + channels[cut].append((channel_id, bond_charge)) + interval_channel_ids.append(channel_ids) + + first_data = _dense_numpy(factors[0], dtype=dtype) + for bond_pos, channel_id in enumerate(interval_channel_ids[0]): + op = first_data[bond_pos] + if np.any(op): + transitions[sites[0]].append((start, channel_id, op)) + + for interval in range(n_sites - 2): + factor_data = _dense_numpy(factors[interval + 1], dtype=dtype) + left_ids = interval_channel_ids[interval] + right_ids = interval_channel_ids[interval + 1] + for left_pos, left_id in enumerate(left_ids): + for right_pos, right_id in enumerate(right_ids): + op = factor_data[left_pos, right_pos] + if np.any(op): + transitions[sites[interval + 1]].append( + (left_id, right_id, op) + ) + + last_data = _dense_numpy(factors[-1], dtype=dtype) + for bond_pos, channel_id in enumerate(interval_channel_ids[-1]): + op = last_data[bond_pos] + if np.any(op): + transitions[sites[-1]].append((channel_id, done, op)) + + identity = np.eye(len(physical_maps[0]), dtype=dtype) + for interval, channel_ids in enumerate(interval_channel_ids): + for site in range(sites[interval] + 1, sites[interval + 1]): + for channel_id in channel_ids: + transitions[site].append((channel_id, channel_id, identity)) + + return list(physical_maps[0]) + + def _generic_symhamiltonian_to_mpo( hamiltonian, L, @@ -5756,14 +5909,17 @@ def _generic_symhamiltonian_to_mpo( site_tag_id="I{}", to_backend=None, dtype=None, + fermionic=False, ): - """Build a Symmray MPO from explicit one- and two-site terms.""" + """Build a Symmray MPO from explicit local terms.""" _, coo2idx_use, mapped_L = _resolve_mpo_mapping( mapper=mapper, idx2coo=idx2coo, coo2idx=coo2idx, ) raw_wheres = tuple(hamiltonian.terms) + if not raw_wheres: + raise ValueError("At least one Hamiltonian term is required to build an MPO.") coordinate_sites = _term_mapping_uses_coordinate_sites(raw_wheres) wheres = tuple( _as_term_where(where, coordinate_sites=coordinate_sites) @@ -5791,7 +5947,12 @@ def _generic_symhamiltonian_to_mpo( if dtype is None else np.dtype(dtype) ) - zero = _normalize_group_charge(getattr(next(iter(hamiltonian.terms.values())), "charge", 0), hamiltonian.symmetry) + first_term = next(iter(hamiltonian.terms.values())) + first_charge = _normalize_group_charge( + getattr(first_term, "charge", 0), + hamiltonian.symmetry, + ) + zero = _zero_like_charge(first_charge) start = ("start",) done = ("done",) channels = [[(start, zero), (done, zero)] for _ in range(max(L - 1, 0))] @@ -5801,6 +5962,32 @@ def _generic_symhamiltonian_to_mpo( for term_pos, (raw_where, where) in enumerate(zip(raw_wheres, mapped_wheres)): term = hamiltonian.terms[raw_where] term_is_fermionic = _is_fermionic_symmray_array(term) + if fermionic and not term_is_fermionic: + raise TypeError( + "Native fermionic MPO construction requires every Hamiltonian " + "term to be a Symmray FermionicArray." + ) + + if fermionic: + term_phys = _add_native_term_to_mpo( + term, + where, + term_pos=term_pos, + channels=channels, + transitions=transitions, + symmetry=hamiltonian.symmetry, + dtype=dtype, + zero=zero, + ) + if phys_map is None: + phys_map = term_phys + elif phys_map != term_phys: + raise ValueError( + "SymHamiltonian.to_mpo requires one physical charge map " + "shared by all sites." + ) + continue + if len(where) == 1: site = int(where[0]) if not 0 <= site < L: @@ -5816,6 +6003,13 @@ def _generic_symhamiltonian_to_mpo( transitions[site].append((start, done, dense)) continue + if len(where) > 2: + raise NotImplementedError( + "Jordan-Wigner compatibility MPO conversion currently supports " + "one- and two-site terms; use fermionic=True for native " + "multi-site terms." + ) + i, j = (int(where[0]), int(where[1])) if i == j: raise ValueError("Hamiltonian edges must connect distinct sites.") @@ -5830,7 +6024,11 @@ def _generic_symhamiltonian_to_mpo( symmetry=hamiltonian.symmetry, dtype=dtype, reverse=reverse, - fermionic=term_is_fermionic, + # ``_decompose_neutral_two_site_term(..., fermionic=True)`` is the + # legacy conversion from native local data to a bosonic/JW + # site-major matrix. A native graded MPO keeps the raw fermionic + # tensor ordering and lets Symmray supply the Koszul signs. + fermionic=term_is_fermionic and not fermionic, ) if left_phys != right_phys: raise ValueError( @@ -5852,7 +6050,11 @@ def _generic_symhamiltonian_to_mpo( channels[cut].append((channel_id, channel_charge)) transitions[i].append((start, channel_id, left_op)) - if term_is_fermionic and _charged_op_needs_fermion_string(left_charge): + if ( + term_is_fermionic + and not fermionic + and _charged_op_needs_fermion_string(left_charge) + ): string_op = _fermion_parity_operator(phys_map, dtype) else: string_op = np.eye(len(phys_map), dtype=dtype) @@ -5879,6 +6081,7 @@ def _generic_symhamiltonian_to_mpo( lower_ind_id=lower_ind_id, site_tag_id=site_tag_id, to_backend=to_backend, + fermionic=fermionic, ) @@ -5920,8 +6123,8 @@ def from_terms( ): """Build a Hamiltonian container from explicit local operators. - ``terms`` maps a site label or an edge tuple to a native one- or - two-site operator. This preserves the operator locations while + ``terms`` maps a site label or support tuple to a native local + operator. This preserves the operator locations while retaining the model and symmetry metadata required for fermionic MPO conversion. """ @@ -5974,15 +6177,18 @@ def to_mpo( site_tag_id="I{}", to_backend=None, dtype=None, + fermionic=False, ): """Build a symmetry-preserving MPS-chain MPO for this Hamiltonian. Coordinate-lattice edges can be mapped with ``mapper=OneDMap(...)`` or the ``idx2coo, coo2idx`` dictionaries from ``OneDMap(...).build()``. - Fermionic paths include parity strings along non-adjacent mapped - hopping channels. + Fermionic compatibility paths include parity strings along + non-adjacent mapped hopping channels. With ``fermionic=True``, native + Symmray ``FermionicArray`` tensors are built directly from arbitrary + neutral one- or multi-site terms. """ - if self.explicit_terms: + if self.explicit_terms or fermionic: return _generic_symhamiltonian_to_mpo( self, L, @@ -5997,6 +6203,7 @@ def to_mpo( site_tag_id=site_tag_id, to_backend=to_backend, dtype=dtype, + fermionic=fermionic, ) if self.model == "fermi_hubbard_spinless": @@ -7446,7 +7653,7 @@ def lattice_half_filling( site: (1, 0) if (site[0] + site[1]) % 2 == 0 else (0, 1) for site in sites } - if self.symmetry == "U1U1": + if self.symmetry in {"U1U1", "Z2Z2"}: occupations = dict(spin_occupations) else: occupations = { @@ -7773,6 +7980,44 @@ def sz_operator(self): """Alias for :meth:`spin_z_operator`.""" return self.spin_z_operator() + def spin_x_term(self, site, *, field): + """Return ``field * Sx`` on one spinful physical site.""" + self._require_spin_flip_symmetry("Sx terms") + if field is None: + raise TypeError("spin_x_term requires explicit field=... .") + field = _node_parameter(field, site) + return self.operator_term( + [(0.5 * field, ((site, "s_plus"),))], + sites=(site,), + add_hc=True, + ) + + def spin_y_term(self, site, *, field): + """Return ``field * Sy`` on one spinful physical site.""" + self._require_spin_flip_symmetry("Sy terms") + if field is None: + raise TypeError("spin_y_term requires explicit field=... .") + field = _node_parameter(field, site) + return self.operator_term( + [(-0.5j * field, ((site, "s_plus"),))], + sites=(site,), + add_hc=True, + ) + + def spin_z_term(self, site, *, field): + """Return ``field * Sz`` on one spinful physical site.""" + self._require_spinful("Sz terms") + if field is None: + raise TypeError("spin_z_term requires explicit field=... .") + field = _node_parameter(field, site) + return self.operator_term( + [ + (0.5 * field, ((site, "number_up"),)), + (-0.5 * field, ((site, "number_down"),)), + ], + sites=(site,), + ) + def spin_z_correlator(self): """Return the bare native two-site ``Sz_i Sz_j`` operator.""" self._require_spinful("Sz-Sz correlators") @@ -8265,6 +8510,8 @@ def hopping_operator(self, *, spin=None, peierls_angle=0.0): def hopping_term(self, edge, *, spin=None, t, peierls_angle=0.0): """Return ``-t`` times the hopping operator on ``edge``.""" + if t is None: + raise TypeError("hopping_term requires explicit t=... .") try: left, right = tuple(edge) except (TypeError, ValueError) as exc: @@ -8293,6 +8540,13 @@ def interaction_operator(self): def interaction_term(self, site, *, U): """Return ``U n_up n_down`` on one physical site.""" + if not self.spinful: + raise ValueError( + "Spinless fermions have no onsite doublon interaction; use " + "density_term(...) for the nearest-neighbor V interaction." + ) + if U is None: + raise TypeError("interaction_term requires explicit U=... .") U = _node_parameter(U, site) return self.operator_term( [(U, ((site, "double"),))], @@ -8312,6 +8566,8 @@ def chemical_potential_operator(self): def chemical_potential_term(self, site, *, mu): """Return ``-mu n`` on one physical site.""" + if mu is None: + raise TypeError("chemical_potential_term requires explicit mu=... .") if self.spinful: mu = _node_parameter(mu, site) mu_up, mu_down = _as_spin_pair(mu, name="mu") @@ -8339,6 +8595,11 @@ def onsite_term(self, site, *, U=None, mu=0.0): ) ) else: + if U is not None: + raise TypeError( + "onsite_term does not accept U=... for spinless fermions; " + "use V=... for nearest-neighbor density interactions." + ) terms.append((-_node_parameter(mu, site), ((site, "number"),))) return self.operator_term(terms, sites=(site,)) @@ -8357,6 +8618,8 @@ def density_operator(self): def density_term(self, edge, *, V): """Return ``V n_i n_j`` on a physical edge.""" + if V is None: + raise TypeError("density_term requires explicit V=... .") try: left, right = tuple(edge) except (TypeError, ValueError) as exc: @@ -8520,6 +8783,8 @@ def interaction_gate(self, dt, *, site=None, U, imaginary=False): "Spinless fermions have no onsite doublon interaction; use " "density_gate(...) for the nearest-neighbor V interaction." ) + if U is None: + raise TypeError("interaction_gate requires explicit U=... .") U = U if site is None else _node_parameter(U, site) theta = dt * U @@ -8540,6 +8805,12 @@ def onsite_gate(self, dt, *, site=None, U=None, mu=0.0, imaginary=False): fermions and ``-mu n`` for spinless fermions. ``U`` and ``mu`` may be site-dependent mappings or callables when ``site`` is supplied. """ + if not self.spinful and U is not None: + raise TypeError( + "onsite_gate does not accept U=... for spinless fermions; " + "use V=... for nearest-neighbor density interactions." + ) + if site is not None: U = _node_parameter(U, site) mu = _node_parameter(mu, site) @@ -8573,6 +8844,9 @@ def build(): def hopping_gate(self, dt, *, t, peierls_angle=0.0, imaginary=False): """Return a two-site native fermionic hopping gate with Peierls phase.""" + if t is None: + raise TypeError("hopping_gate requires explicit t=... .") + def build(): if not self.spinful: gate = _spinless_hopping_gate( @@ -8600,6 +8874,8 @@ def density_gate(self, dt, *, V, imaginary=False): For spinless fermions this is ``V n_i n_j``. For spinful fermions it is ``V (n_up + n_down)_i (n_up + n_down)_j``. """ + if V is None: + raise TypeError("density_gate requires explicit V=... .") theta = dt * V def build(): @@ -8617,6 +8893,8 @@ def build(): def chemical_potential_gate(self, dt, *, mu, site=None, imaginary=False): """Return the chemical-potential part of an onsite gate.""" + if mu is None: + raise TypeError("chemical_potential_gate requires explicit mu=... .") mu = mu if site is None else _node_parameter(mu, site) if self.spinful: mu_up, mu_down = _as_spin_pair(mu, name="mu") @@ -8639,88 +8917,171 @@ def build(): def gate(self, name, dt, *, site=None, where=None, imaginary=False, **params): """Build a named native gate using the local fermionic conventions.""" + if "edge" in params and where is not None: + raise TypeError( + "Fermion.gate accepts at most one of where=... and edge=... ." + ) edge = params.pop("edge", where) del where # Gate locations belong to the stream entry, not the tensor. name = str(name).lower().replace("-", "_") + + def require(parameter): + try: + return params.pop(parameter) + except KeyError as exc: + raise TypeError( + f"Fermion.gate({name!r}, ...) requires explicit " + f"{parameter}=... ." + ) from exc + + def finish(gate, *, accepts_site=False, accepts_edge=False): + if site is not None and not accepts_site: + raise TypeError( + f"Fermion.gate({name!r}, ...) does not accept site=... ." + ) + if edge is not None and not accepts_edge: + raise TypeError( + f"Fermion.gate({name!r}, ...) does not accept edge=... ." + ) + if params: + names = ", ".join(sorted(params)) + raise TypeError( + f"Unexpected Fermion.gate parameter(s) for {name!r}: {names}." + ) + return gate + if name in {"sx", "spin_x"}: - return self.spin_x_gate(dt, site=site, imaginary=imaginary) + return finish( + self.spin_x_gate(dt, site=site, imaginary=imaginary), + accepts_site=True, + ) if name in {"sy", "spin_y"}: - return self.spin_y_gate(dt, site=site, imaginary=imaginary) + return finish( + self.spin_y_gate(dt, site=site, imaginary=imaginary), + accepts_site=True, + ) if name in {"sz", "spin_z"}: - return self.spin_z_gate(dt, site=site, imaginary=imaginary) + return finish( + self.spin_z_gate(dt, site=site, imaginary=imaginary), + accepts_site=True, + ) if name in {"sxx", "spin_x_x", "sx_sx"}: - return self.spin_x_correlator_gate(dt, edge=edge, imaginary=imaginary) + return finish( + self.spin_x_correlator_gate(dt, edge=edge, imaginary=imaginary), + accepts_edge=True, + ) if name in {"syy", "spin_y_y", "sy_sy"}: - return self.spin_y_correlator_gate(dt, edge=edge, imaginary=imaginary) + return finish( + self.spin_y_correlator_gate(dt, edge=edge, imaginary=imaginary), + accepts_edge=True, + ) if name in {"szz", "spin_z_z", "sz_sz"}: - return self.spin_z_correlator_gate(dt, edge=edge, imaginary=imaginary) + return finish( + self.spin_z_correlator_gate(dt, edge=edge, imaginary=imaginary), + accepts_edge=True, + ) if name in {"xy", "xy_exchange"}: - return self.xy_exchange_gate(dt, edge=edge, imaginary=imaginary) + return finish( + self.xy_exchange_gate(dt, edge=edge, imaginary=imaginary), + accepts_edge=True, + ) if name in {"heisenberg", "heis"}: - return self.heisenberg_gate(dt, edge=edge, imaginary=imaginary) + return finish( + self.heisenberg_gate(dt, edge=edge, imaginary=imaginary), + accepts_edge=True, + ) if name in {"onsite", "hubbard_onsite"}: - return self.onsite_gate( - dt, - site=site, - U=params.pop("U", None), - mu=params.pop("mu", 0.0), - imaginary=imaginary, + return finish( + self.onsite_gate( + dt, + site=site, + U=params.pop("U", None), + mu=params.pop("mu", 0.0), + imaginary=imaginary, + ), + accepts_site=True, ) if name in {"interaction", "onsite_interaction", "doublon"}: - return self.interaction_gate( - dt, - site=site, - U=params.pop("U", None), - imaginary=imaginary, + return finish( + self.interaction_gate( + dt, + site=site, + U=require("U"), + imaginary=imaginary, + ), + accepts_site=True, ) if name in {"hopping", "hop"}: - return self.hopping_gate( - dt, - t=params.pop("t", None), - peierls_angle=params.pop("peierls_angle", 0.0), - imaginary=imaginary, + return finish( + self.hopping_gate( + dt, + t=require("t"), + peierls_angle=params.pop("peierls_angle", 0.0), + imaginary=imaginary, + ), ) if name in {"density", "density_interaction", "nn"}: - return self.density_gate( - dt, - V=params.pop("V", None), - imaginary=imaginary, + return finish( + self.density_gate( + dt, + V=require("V"), + imaginary=imaginary, + ), ) if name in {"chemical", "chemical_potential", "mu"}: - return self.chemical_potential_gate( - dt, - mu=params.pop("mu", None), - site=site, - imaginary=imaginary, + return finish( + self.chemical_potential_gate( + dt, + mu=require("mu"), + site=site, + imaginary=imaginary, + ), + accepts_site=True, ) raise ValueError(f"Unknown fermion gate {name!r}.") def param_gate(self, name, params, *, imaginary=False, **kwargs): """Build a gate from a Quimb-style parameter sequence.""" name = str(name).lower().replace("-", "_") + + def finish(gate): + if kwargs: + names = ", ".join(sorted(kwargs)) + raise TypeError( + f"Unexpected Fermion.param_gate parameter(s) for {name!r}: " + f"{names}." + ) + return gate + if name in {"interaction", "onsite_interaction", "doublon"}: if not self.spinful: raise ValueError("Spinless fermions do not have doublon gates.") - return fermion_interaction_param_gen( - params, - symmetry=self.symmetry, - imaginary=imaginary, + return finish( + fermion_interaction_param_gen( + params, + symmetry=self.symmetry, + imaginary=imaginary, + ) ) if name in {"density", "density_interaction", "nn"}: if self.spinful: raise ValueError("Spinful density gates are not the onsite interaction gate.") - return fermion_density_param_gen( - params, - symmetry=self.symmetry, - imaginary=imaginary, + return finish( + fermion_density_param_gen( + params, + symmetry=self.symmetry, + imaginary=imaginary, + ) ) if name in {"hopping", "hop"}: - return fermion_hopping_param_gen( - params, - spinful=self.spinful, - symmetry=self.symmetry, - imaginary=imaginary, - peierls_angle=kwargs.pop("peierls_angle", 0.0), + return finish( + fermion_hopping_param_gen( + params, + spinful=self.spinful, + symmetry=self.symmetry, + imaginary=imaginary, + peierls_angle=kwargs.pop("peierls_angle", 0.0), + ) ) raise ValueError(f"Unknown parameterized fermion gate {name!r}.") @@ -8863,6 +9224,11 @@ def gate_stream( U=None, V=0.0, mu=0.0, + field_x=0.0, + field_y=0.0, + field_z=0.0, + pairing=0.0, + pairing_phase=0.0, ): """Return a canonical fermion gate stream with explicit couplings.""" if order not in {1, 2}: @@ -8871,6 +9237,11 @@ def gate_stream( raise TypeError("gate_stream requires explicit t=... .") if self.spinful and U is None: raise TypeError("gate_stream requires explicit U=... for spinful fermions.") + if not self.spinful and U is not None: + raise TypeError( + "gate_stream does not accept U=... for spinless fermions; " + "use V=... for nearest-neighbor density interactions." + ) edges = _as_edges(edges) sites = _sites_from_edges(edges, sites) @@ -8885,6 +9256,11 @@ def gate_stream( U=U, V=V, mu=mu, + field_x=field_x, + field_y=field_y, + field_z=field_z, + pairing=pairing, + pairing_phase=pairing_phase, ) entries = [] @@ -8901,7 +9277,23 @@ def gate_stream( ) for site in sites ) - if V != 0 or isinstance(V, Mapping) or callable(V): + for coupling, gate in ( + (field_x, self.spin_x_gate), + (field_y, self.spin_y_gate), + (field_z, self.spin_z_gate), + ): + if _coupling_is_active(coupling): + entries.extend( + ( + gate( + dt * _node_parameter(coupling, site), + imaginary=imaginary, + ), + site, + ) + for site in sites + ) + if _coupling_is_active(V): entries.extend( ( self.density_gate( @@ -8913,6 +9305,20 @@ def gate_stream( ) for left, right in edges ) + if _coupling_is_active(pairing): + entries.extend( + ( + self.pairing_gate( + dt, + edge=(left, right), + coefficient=_edge_parameter(pairing, left, right), + phase=_edge_angle_parameter(pairing_phase, left, right), + imaginary=imaginary, + ), + (left, right), + ) + for left, right in edges + ) entries.extend( ( self.hopping_gate( @@ -8927,7 +9333,19 @@ def gate_stream( ) return SymGateStream( entries, - hamiltonian=self.hamiltonian(edges, t=t, U=U, V=V, mu=mu), + hamiltonian=self.hamiltonian( + edges, + sites=sites, + t=t, + U=U, + V=V, + mu=mu, + field_x=field_x, + field_y=field_y, + field_z=field_z, + pairing=pairing, + pairing_phase=pairing_phase, + ), dt=dt, imaginary=imaginary, order=1, @@ -8945,6 +9363,11 @@ def strang_gate_stream( U=None, V=0.0, mu=0.0, + field_x=0.0, + field_y=0.0, + field_z=0.0, + pairing=0.0, + pairing_phase=0.0, ): """Return an edge-coloured second-order stream with explicit couplings.""" if t is None: @@ -8953,6 +9376,11 @@ def strang_gate_stream( raise TypeError( "strang_gate_stream requires explicit U=... for spinful fermions." ) + if not self.spinful and U is not None: + raise TypeError( + "strang_gate_stream does not accept U=... for spinless fermions; " + "use V=... for nearest-neighbor density interactions." + ) edges = _as_edges(edges) sites = _sites_from_edges(edges, sites) half_dt = dt / 2 @@ -8970,7 +9398,24 @@ def strang_gate_stream( ) for site in sites ] - if V != 0 or isinstance(V, Mapping) or callable(V): + fields = ( + (field_x, self.spin_x_gate), + (field_y, self.spin_y_gate), + (field_z, self.spin_z_gate), + ) + for coupling, gate in fields: + if _coupling_is_active(coupling): + entries.extend( + ( + gate( + half_dt * _node_parameter(coupling, site), + imaginary=imaginary, + ), + site, + ) + for site in sites + ) + if _coupling_is_active(V): entries.extend( ( self.density_gate( @@ -8982,6 +9427,40 @@ def strang_gate_stream( ) for left, right in edges ) + if _coupling_is_active(pairing): + quarter_dt = dt / 4 + for layer in layers: + entries.extend( + ( + self.pairing_gate( + quarter_dt, + edge=(left, right), + coefficient=_edge_parameter(pairing, left, right), + phase=_edge_angle_parameter( + pairing_phase, left, right + ), + imaginary=imaginary, + ), + (left, right), + ) + for left, right in layer + ) + for layer in reversed(layers): + entries.extend( + ( + self.pairing_gate( + quarter_dt, + edge=(left, right), + coefficient=_edge_parameter(pairing, left, right), + phase=_edge_angle_parameter( + pairing_phase, left, right + ), + imaginary=imaginary, + ), + (left, right), + ) + for left, right in reversed(layer) + ) for layer in layers: entries.extend( ( @@ -9008,7 +9487,41 @@ def strang_gate_stream( ) for left, right in layer ) - if V != 0 or isinstance(V, Mapping) or callable(V): + if _coupling_is_active(pairing): + quarter_dt = dt / 4 + for layer in reversed(layers): + entries.extend( + ( + self.pairing_gate( + quarter_dt, + edge=(left, right), + coefficient=_edge_parameter(pairing, left, right), + phase=_edge_angle_parameter( + pairing_phase, left, right + ), + imaginary=imaginary, + ), + (left, right), + ) + for left, right in reversed(layer) + ) + for layer in layers: + entries.extend( + ( + self.pairing_gate( + quarter_dt, + edge=(left, right), + coefficient=_edge_parameter(pairing, left, right), + phase=_edge_angle_parameter( + pairing_phase, left, right + ), + imaginary=imaginary, + ), + (left, right), + ) + for left, right in layer + ) + if _coupling_is_active(V): entries.extend( ( self.density_gate( @@ -9020,6 +9533,18 @@ def strang_gate_stream( ) for left, right in edges ) + for coupling, gate in reversed(fields): + if _coupling_is_active(coupling): + entries.extend( + ( + gate( + half_dt * _node_parameter(coupling, site), + imaginary=imaginary, + ), + site, + ) + for site in sites + ) entries.extend( ( self.onsite_gate( @@ -9035,7 +9560,19 @@ def strang_gate_stream( ) return SymGateStream( entries, - hamiltonian=self.hamiltonian(edges, t=t, U=U, V=V, mu=mu), + hamiltonian=self.hamiltonian( + edges, + sites=sites, + t=t, + U=U, + V=V, + mu=mu, + field_x=field_x, + field_y=field_y, + field_z=field_z, + pairing=pairing, + pairing_phase=pairing_phase, + ), dt=dt, imaginary=imaginary, order=2, @@ -9098,10 +9635,16 @@ def hamiltonian( self, terms_or_edges, *, + sites=None, t=None, U=None, V=0.0, mu=0.0, + field_x=0.0, + field_y=0.0, + field_z=0.0, + pairing=0.0, + pairing_phase=0.0, flat=False, to_backend=None, ): @@ -9115,11 +9658,17 @@ def hamiltonian( stored on :class:`Fermion`. """ to_backend = self.to_backend if to_backend is None else to_backend + extra_couplings = (field_x, field_y, field_z, pairing) if isinstance(terms_or_edges, Mapping): - if any(value is not None for value in (t, U)) or V != 0 or mu != 0: + if ( + any(value is not None for value in (t, U)) + or _coupling_is_active(V) + or _coupling_is_active(mu) + or any(_coupling_is_active(value) for value in extra_couplings) + ): raise TypeError( "When passing explicit terms, put every coupling in the " - "native arrays rather than passing t/U/V/mu again." + "native arrays rather than passing model couplings again." ) terms = _apply_to_hamiltonian_terms(terms_or_edges, to_backend) self._validate_hamiltonian_terms(terms) @@ -9136,20 +9685,196 @@ def hamiltonian( raise TypeError( "hamiltonian(edges, ...) requires explicit U=... for spinful fermions." ) + if not self.spinful and U is not None: + raise TypeError( + "hamiltonian(edges, ...) does not accept U=... for spinless " + "fermions; use V=... for nearest-neighbor density interactions." + ) + edges = _as_edges(terms_or_edges) params = {"t": t, "V": V, "mu": mu} if self.spinful: params["U"] = U hamiltonian = SymHamiltonian.from_edges( self.model, self.symmetry, - terms_or_edges, + edges, flat=flat, to_backend=to_backend, **params, ) + if not any(_coupling_is_active(value) for value in extra_couplings): + self._validate_hamiltonian_terms(hamiltonian.terms) + return hamiltonian + + sites = _sites_from_edges(edges, sites) + terms = dict(hamiltonian.terms) + for coupling, build_term in ( + (field_x, self.spin_x_term), + (field_y, self.spin_y_term), + (field_z, self.spin_z_term), + ): + if _coupling_is_active(coupling): + for site in sites: + where = (site,) + term = build_term(site, field=coupling) + terms[where] = terms[where] + term if where in terms else term + if _coupling_is_active(pairing): + for left, right in edges: + edge = (left, right) + terms[edge] = terms[edge] + self.pairing_operator( + edge, + coefficient=_edge_parameter(pairing, left, right), + phase=_edge_angle_parameter(pairing_phase, left, right), + ) + parameters = { + **params, + "field_x": field_x, + "field_y": field_y, + "field_z": field_z, + "pairing": pairing, + "pairing_phase": pairing_phase, + } + hamiltonian = SymHamiltonian.from_terms( + self.model, + self.symmetry, + terms, + parameters=parameters, + ) self._validate_hamiltonian_terms(hamiltonian.terms) return hamiltonian + def build_mpo( + self, + terms_or_edges, + *, + L=None, + mapper=None, + idx2coo=None, + coo2idx=None, + max_bond=None, + cutoff=1e-12, + compress=True, + upper_ind_id="k{}", + lower_ind_id="b{}", + site_tag_id="I{}", + dtype=None, + fermionic=False, + to_backend=None, + **params, + ): + """Build a Symmray MPO directly from this fermion model. + + This is the model-facing shorthand for + ``fermion.hamiltonian(...).to_mpo(...)``. The resulting MPO uses the + same symmetry as :meth:`SymHamiltonian.to_mpo`. By default this + model-facing helper builds the current bosonic/Jordan-Wigner + compatibility MPO and can be evolved with ``jw_trotter_gates``. + Pass ``fermionic=True`` to select native graded ``FermionicArray`` + construction; this is equivalent to calling :meth:`to_mpo` with the + same model parameters. + ``t``, ``U``/``V``, and ``mu`` remain explicit build parameters and are + forwarded to :meth:`hamiltonian`. + """ + to_backend = self.to_backend if to_backend is None else to_backend + hamiltonian = self.hamiltonian( + terms_or_edges, + to_backend=to_backend, + **params, + ) + return hamiltonian.to_mpo( + L=L, + mapper=mapper, + idx2coo=idx2coo, + coo2idx=coo2idx, + max_bond=max_bond, + cutoff=cutoff, + compress=compress, + upper_ind_id=upper_ind_id, + lower_ind_id=lower_ind_id, + site_tag_id=site_tag_id, + dtype=dtype, + fermionic=fermionic, + to_backend=to_backend, + ) + + def to_mpo( + self, + terms_or_edges=None, + *, + hamiltonian=None, + L=None, + mapper=None, + idx2coo=None, + coo2idx=None, + max_bond=None, + cutoff=1e-12, + compress=True, + upper_ind_id="k{}", + lower_ind_id="b{}", + site_tag_id="I{}", + dtype=None, + fermionic=True, + to_backend=None, + **params, + ): + """Build a native graded MPO from a fermionic term collection. + + ``terms_or_edges`` may be lattice edges for the built-in model or a + mapping such as ``{(0, 2, 4): fermion.operator_term(...)}``. The + latter supports arbitrary neutral term support, including + non-contiguous sites. Pass an existing :class:`SymHamiltonian` with + ``hamiltonian=`` when the terms have already been assembled. + + The native path is selected by default and returns MPO tensors backed + by Symmray ``FermionicArray`` objects. Set ``fermionic=False`` to + request the explicit Jordan--Wigner compatibility MPO instead. + + ``to_backend`` overrides the backend configured on this ``Fermion`` + instance for both the Hamiltonian terms and the returned MPO blocks. + """ + to_backend = self.to_backend if to_backend is None else to_backend + if hamiltonian is not None: + if terms_or_edges is not None: + raise TypeError( + "Pass either terms_or_edges or hamiltonian, not both." + ) + if not isinstance(hamiltonian, SymHamiltonian): + raise TypeError("hamiltonian must be a SymHamiltonian instance.") + target = hamiltonian + elif isinstance(terms_or_edges, SymHamiltonian): + target = terms_or_edges + else: + if terms_or_edges is None: + raise TypeError("to_mpo requires terms_or_edges or hamiltonian.") + target = self.hamiltonian( + terms_or_edges, + to_backend=to_backend, + **params, + ) + params = {} + + if params: + names = ", ".join(sorted(params)) + raise TypeError( + "Model parameters cannot be supplied with an existing " + f"SymHamiltonian: {names}." + ) + return target.to_mpo( + L=L, + mapper=mapper, + idx2coo=idx2coo, + coo2idx=coo2idx, + max_bond=max_bond, + cutoff=cutoff, + compress=compress, + upper_ind_id=upper_ind_id, + lower_ind_id=lower_ind_id, + site_tag_id=site_tag_id, + dtype=dtype, + fermionic=fermionic, + to_backend=to_backend, + ) + def local_terms(self, edges, *, layout="site", **params): """Return native local terms for site or qMERA energy workflows. @@ -9198,10 +9923,22 @@ def majorana_terms(self, geometry, **params): return self.local_terms(geometry, layout="majorana", **params) -# Kept for callers that adopted the initial public names before the helper was -# generalized to the operator/gate-focused ``Fermion`` API. -SpinfulFermion = Fermion -SpinfulFermionHubbard = Fermion +class SpinfulFermion(Fermion): + """Compatibility constructor that always selects the spinful local space.""" + + def __init__(self, *args, spinful=True, **kwargs): + if len(args) > 3: + raise TypeError( + "SpinfulFermion fixes spinful=True; pass at most symmetry, " + "dtype, and to_backend positionally." + ) + if not spinful: + raise TypeError("SpinfulFermion always uses spinful=True.") + super().__init__(*args, spinful=True, **kwargs) + + +# Compatibility spelling for the initial, overly model-specific public name. +SpinfulFermionHubbard = SpinfulFermion class SymMPS(_SymState): diff --git a/tests/test_optimize_mpo.py b/tests/test_optimize_mpo.py index 85036f8..362f5c1 100644 --- a/tests/test_optimize_mpo.py +++ b/tests/test_optimize_mpo.py @@ -453,3 +453,328 @@ def test_mpo_optimizer_mpo_mode_bra_only_gate(): assert out.L == 4 assert not np.allclose(out.to_dense(), mpo0.to_dense()) + + +def _native_u1u1_identity_mpo(L=3): + """Make a small native graded MPO fixture without testing construction.""" + pytest.importorskip("symmray") + import symmray.utils as sr_utils + + phys_map = [(0, 0), (0, 1), (1, 0), (1, 1)] + arrays = [] + for site in range(L): + if site == 0: + data = np.zeros((1, 4, 4), dtype="complex128") + data[0] = np.eye(4) + index_maps = [[(0, 0)], phys_map, phys_map] + duals = (False, False, True) + elif site == L - 1: + data = np.zeros((1, 4, 4), dtype="complex128") + data[0] = np.eye(4) + index_maps = [[(0, 0)], phys_map, phys_map] + duals = (True, False, True) + else: + data = np.zeros((1, 1, 4, 4), dtype="complex128") + data[0, 0] = np.eye(4) + index_maps = [[(0, 0)], [(0, 0)], phys_map, phys_map] + duals = (True, False, False, True) + arrays.append( + sr_utils.from_dense( + data, + symmetry="U1U1", + index_maps=index_maps, + duals=duals, + fermionic=True, + charge=(0, 0), + ) + ) + return qtn.MatrixProductOperator( + arrays, + shape="lrud", + upper_ind_id="k{}", + lower_ind_id="b{}", + site_tag_id="I{}", + ) + + +@pytest.mark.parametrize("mode", ["svd", "mpo", "dmrg"]) +def test_mpo_optimizer_replays_native_graded_mpo_without_dense_fallback(mode): + """Native graded MPO inputs remain FermionicArray-backed through replay.""" + fermion = py.Fermion(spinful=True, symmetry="U1U1") + gates = fermion.strang_gate_stream( + [(0, 1), (1, 2)], + dt=0.01, + t=1.0, + U=2.0, + mu=0.1, + ) + + out = py.MpoOptimizer( + _native_u1u1_identity_mpo(), + gates=gates, + chi=8, + mode=mode, + ).run(progbar=False, cutoff=1e-10, fidelity_samples=0, n_iter=1) + + assert all(type(tensor.data).__name__ == "U1U1FermionicArray" for tensor in out) + + +def test_mpo_optimizer_materializes_native_long_range_split_gates(): + """Long-range native split gates are canonicalizable after replay.""" + fermion = py.Fermion(spinful=True, symmetry="U1U1") + gates = fermion.strang_gate_stream( + [(0, 3)], + dt=0.01, + t=1.0, + U=2.0, + mu=0.1, + ) + + out = py.MpoOptimizer( + _native_u1u1_identity_mpo(L=4), + gates=gates, + chi=8, + mode="svd", + ).run(progbar=False, cutoff=1e-10, fidelity_samples=0) + + assert out.L == 4 + assert all(len(out.tag_map[f"I{site}"]) == 1 for site in range(4)) + assert all(type(tensor.data).__name__ == "U1U1FermionicArray" for tensor in out) + + +def test_mpo_optimizer_adapts_long_range_native_gate_to_jw_symmray_mpo(): + """The current JW MPO path also handles long-range native even gates.""" + fermion = py.Fermion(spinful=True, symmetry="U1U1") + mpo = fermion.build_mpo( + [(0, 3)], + L=4, + t=1.0, + U=2.0, + mu=0.1, + ) + gates = fermion.strang_gate_stream( + [(0, 3)], + dt=0.01, + t=1.0, + U=2.0, + mu=0.1, + ) + + out = py.MpoOptimizer(mpo, gates=gates, chi=8, mode="svd").run( + progbar=False, + cutoff=1e-10, + fidelity_samples=0, + ) + + assert out.L == 4 + assert all(type(tensor.data).__name__ == "U1U1Array" for tensor in out) + + +@pytest.mark.parametrize("mode", ["svd", "mpo", "dmrg"]) +def test_mpo_optimizer_handles_fermion_symmray_mpo_and_native_gate_stream(mode): + """The optimizer adapts native gates onto the current U1U1 MPO path.""" + pytest.importorskip("symmray") + + fermion = py.Fermion(spinful=True, symmetry="U1U1") + edges = [(0, 1), (1, 2)] + mpo = fermion.build_mpo( + edges, + L=3, + t=1.0, + U=2.0, + mu=0.1, + max_bond=16, + cutoff=1e-12, + ) + gates = fermion.strang_gate_stream( + edges, + dt=0.01, + t=1.0, + U=2.0, + mu=0.1, + ) + + optimizer = py.MpoOptimizer(mpo, gates=gates, chi=8, mode=mode) + out = optimizer.run(progbar=False, cutoff=1e-10, fidelity_samples=0, n_iter=1) + + assert out.L == 3 + assert out.max_bond() <= 8 + assert all(type(tensor.data).__name__ == "U1U1Array" for tensor in out) + + +def test_fermion_build_mpo_and_ham_tn_adapter_preserve_symmetry(): + """Both public MPO builders preserve the model's U1U1 symmetry.""" + pytest.importorskip("symmray") + + fermion = py.Fermion(spinful=True, symmetry="U1U1") + edges = [(0, 1), (1, 2)] + builder = py.ham_tn(Lx=3, Ly=1, data_type="complex128") + + direct = fermion.build_mpo(edges, L=3, t=1.0, U=0.0, mu=0.0) + adapted = builder.build_mpo( + fermion=fermion, + edges=edges, + phys_dim=4, + t=1.0, + U=0.0, + mu=0.0, + ) + positional = builder.build_mpo( + fermion, + edges=edges, + t=1.0, + U=0.0, + mu=0.0, + ) + + assert direct.L == adapted.L == positional.L == 3 + assert all(type(tensor.data).__name__ == "U1U1Array" for tensor in direct) + assert all(type(tensor.data).__name__ == "U1U1Array" for tensor in adapted) + assert all(type(tensor.data).__name__ == "U1U1Array" for tensor in positional) + + +def test_build_mpo_can_select_native_fermionic_construction(): + """The model-facing builder exposes the native MPO path explicitly.""" + pytest.importorskip("symmray") + fermion = py.Fermion(spinful=True, symmetry="U1U1") + + native = fermion.build_mpo( + [(0, 1)], + L=2, + t=1.0, + U=0.0, + mu=0.0, + fermionic=True, + compress=False, + ) + direct = fermion.to_mpo( + [(0, 1)], + L=2, + t=1.0, + U=0.0, + mu=0.0, + compress=False, + ) + + assert all( + type(tensor.data).__name__ == "U1U1FermionicArray" + for tensor in native + ) + assert native.to_dense().allclose(direct.to_dense()) + + +def test_mpo_optimizer_explicit_compress_handles_empty_symmray_stream(): + """The optimizer compresses symmetry-preserving Symmray MPOs directly.""" + pytest.importorskip("symmray") + + fermion = py.Fermion(spinful=True, symmetry="U1U1") + mpo = fermion.build_mpo( + [(0, 1), (1, 2)], + L=3, + t=1.0, + U=0.0, + mu=0.0, + compress=False, + ) + raw_bond = mpo.max_bond() + optimizer = py.MpoOptimizer(mpo, gates=[], chi=2, mode="svd") + out = optimizer.compress(cutoff=1e-10) + + # Symmray can retain a small sector-multiplicity overshoot for a requested + # cap, but the compression must still reduce the raw MPO bond. + assert out.max_bond() < raw_bond + assert all(type(tensor.data).__name__ == "U1U1Array" for tensor in out) + + +def test_fermion_to_mpo_builds_native_mpo_for_optimizer_replay(): + """The native Fermion.to_mpo path feeds the MPO optimizer directly.""" + fermion = py.Fermion(spinful=True, symmetry="U1U1") + hopping = fermion.hopping_operator() + two_site_mpo = fermion.to_mpo( + {(0, 1): hopping}, + L=2, + compress=False, + ) + assert two_site_mpo.to_dense().allclose(hopping.fuse((0, 1), (2, 3))) + + hamiltonian = fermion.hamiltonian( + [(0, 1), (1, 2)], + t=1.0, + U=2.0, + mu=0.1, + ) + mpo = fermion.to_mpo( + hamiltonian=hamiltonian, + L=3, + max_bond=16, + cutoff=1e-12, + ) + + assert all(type(tensor.data).__name__ == "U1U1FermionicArray" for tensor in mpo) + + out = py.MpoOptimizer( + mpo, + gates=hamiltonian.trotter_gates(0.01), + chi=8, + mode="svd", + ).run(progbar=False, cutoff=1e-10, fidelity_samples=0) + assert all(type(tensor.data).__name__ == "U1U1FermionicArray" for tensor in out) + + +def test_fermion_to_mpo_preserves_configured_backend(): + """Native MPO conversion applies the Fermion backend to every block.""" + torch = pytest.importorskip("torch") + backend = py.backend_torch(dtype=torch.complex128, device="cpu") + fermion = py.Fermion( + spinful=True, + symmetry="U1U1", + to_backend=backend, + ) + + mpo = fermion.to_mpo( + [(0, 1)], + L=2, + t=1.0, + U=2.0, + mu=0.1, + compress=False, + ) + + for tensor in mpo: + assert tensor.data.backend == "torch" + assert all(isinstance(block, torch.Tensor) for block in tensor.data.blocks.values()) + + +def test_fermion_to_mpo_accepts_arbitrary_neutral_term_support(): + """Native MPO conversion supports non-contiguous multi-site terms.""" + fermion = py.Fermion(spinful=False, symmetry="U1") + term = fermion.operator_term( + [(1.0, ((2, "create"), (0, "number"), (3, "annihilate")))], + sites=(2, 0, 3), + ) + hamiltonian = fermion.hamiltonian({(2, 0, 3): term}) + mpo = fermion.to_mpo( + hamiltonian=hamiltonian, + L=4, + compress=False, + ) + + assert mpo.L == 4 + assert all(type(tensor.data).__name__ == "U1FermionicArray" for tensor in mpo) + embedded_term = fermion.operator_term( + [(1.0, ((2, "create"), (0, "number"), (3, "annihilate")))], + sites=(0, 1, 2, 3), + ) + assert mpo.to_dense().allclose( + embedded_term.fuse((0, 1, 2, 3), (4, 5, 6, 7)) + ) + + +def test_fermion_to_mpo_handles_one_site_native_term(): + """Native MPO construction also handles the no-virtual-bond case.""" + fermion = py.Fermion(spinful=True, symmetry="U1U1") + term = fermion.interaction_operator() + mpo = fermion.to_mpo({(0,): term}, L=1, compress=False) + + assert mpo.to_dense().allclose(term) + assert mpo.pepsy_compression_report["raw_max_bond"] == 1 diff --git a/tests/test_symmetric_tensors.py b/tests/test_symmetric_tensors.py index ab902fe..1e6809c 100644 --- a/tests/test_symmetric_tensors.py +++ b/tests/test_symmetric_tensors.py @@ -111,6 +111,18 @@ def test_spinful_fermion_helper_bundles_symmetry_aware_building_blocks( assert ham.symmetry == symmetry +def test_spinful_fermion_compatibility_constructors_enforce_spinful_space(): + """Legacy spellings must not accidentally construct a spinless helper.""" + assert SpinfulFermionHubbard is SpinfulFermion + assert SpinfulFermion(symmetry="U1").spinful + assert SymmFermions.spinful(symmetry="U1").spinful + + with pytest.raises(TypeError, match="always uses spinful=True"): + SpinfulFermion(spinful=False) + with pytest.raises(TypeError, match="always uses spinful=True"): + SymmFermions.spinful(spinful=False) + + @pytest.mark.parametrize("symmetry", ["U1", "Z2"]) def test_unified_fermion_helper_supports_spinless_native_workflow(symmetry): """The unified helper should expose the spinless native t-V workflow.""" @@ -150,6 +162,109 @@ def test_unified_fermion_helper_supports_spinless_native_workflow(symmetry): )) == {(0, 1), (1, 2)} +def test_spinless_fermion_rejects_ignored_hubbard_couplings(): + """A spinless t-V model must never silently discard a Hubbard U.""" + fermion = Fermion(spinful=False, symmetry="U1") + edges = ((0, 1),) + + with pytest.raises(TypeError, match="spinless fermions"): + fermion.onsite_term(0, U=2.0) + with pytest.raises(TypeError, match="spinless fermions"): + fermion.onsite_gate(0.01, U=2.0) + with pytest.raises(ValueError, match="Spinless.*doublon"): + fermion.interaction_term(0, U=2.0) + with pytest.raises(TypeError, match="spinless fermions"): + fermion.gate_stream(edges, 0.01, t=1.0, U=2.0) + with pytest.raises(TypeError, match="spinless fermions"): + fermion.strang_gate_stream(edges, 0.01, t=1.0, U=2.0) + with pytest.raises(TypeError, match="spinless fermions"): + fermion.hamiltonian(edges, t=1.0, U=2.0) + + +def test_named_fermion_gate_rejects_missing_and_unknown_parameters(): + """The generic gate front door validates every coupling it accepts.""" + fermion = Fermion(spinful=True, symmetry="U1U1") + + with pytest.raises(TypeError, match="requires explicit t"): + fermion.gate("hopping", 0.01) + with pytest.raises(TypeError, match="requires explicit t"): + fermion.hopping_gate(0.01, t=None) + with pytest.raises(TypeError, match="requires explicit U"): + fermion.gate("interaction", 0.01) + with pytest.raises(TypeError, match="requires explicit V"): + fermion.density_gate(0.01, V=None) + with pytest.raises(TypeError, match="Unexpected Fermion.gate parameter"): + fermion.gate("hopping", 0.01, t=1.0, typo=7.0) + with pytest.raises(TypeError, match="does not accept edge"): + fermion.gate("hopping", 0.01, t=1.0, edge=(0, 1)) + with pytest.raises(TypeError, match="at most one"): + fermion.gate("sxx", 0.01, where=(0, 1), edge=(0, 1)) + with pytest.raises(TypeError, match="Unexpected Fermion.param_gate parameter"): + fermion.param_gate("hopping", (0.01,), typo=7.0) + + +def test_fermion_streams_include_explicit_spin_fields_in_matching_hamiltonian(): + """Total-U1 Hubbard streams support longitudinal and transverse fields.""" + fermion = Fermion(spinful=True, symmetry="U1") + edges = ((0, 1),) + stream = fermion.strang_gate_stream( + edges, + 0.01, + t=1.0, + U=2.0, + field_x={0: 0.2, 1: -0.1}, + field_z=0.3, + ) + hamiltonian = stream.hamiltonian + + assert hamiltonian.explicit_terms + expected = fermion.spin_x_term(0, field=0.2) + fermion.spin_z_term( + 0, field=0.3 + ) + np.testing.assert_allclose( + hamiltonian.terms[(0,)].to_dense(), expected.to_dense() + ) + assert hamiltonian.to_mpo(L=2, compress=False).L == 2 + assert len(stream) == 14 + + with pytest.raises(ValueError, match="symmetry='U1' or 'Z2'"): + Fermion(spinful=True, symmetry="U1U1").strang_gate_stream( + edges, 0.01, t=1.0, U=2.0, field_x=0.1 + ) + + +def test_spinless_z2_streams_include_pairing_in_matching_hamiltonian(): + """Parity-preserving pairing belongs to the spinless Z2 stream only.""" + fermion = Fermion(spinful=False, symmetry="Z2") + edges = ((0, 1),) + stream = fermion.strang_gate_stream( + edges, + 0.01, + t=1.0, + pairing=0.2, + pairing_phase=0.3, + ) + reference = fermion.hamiltonian( + edges, + t=1.0, + pairing=0.2, + pairing_phase=0.3, + ) + + assert stream.hamiltonian.explicit_terms + np.testing.assert_allclose( + stream.hamiltonian.terms[(0, 1)].to_dense(), + reference.terms[(0, 1)].to_dense(), + ) + assert reference.to_mpo(L=2, compress=False).L == 2 + assert len(stream) == 10 + + with pytest.raises(NotImplementedError, match="spinful=False"): + Fermion(spinful=True, symmetry="Z2").strang_gate_stream( + edges, 0.01, t=1.0, U=2.0, pairing=0.2 + ) + + def test_unified_spinful_fermion_gate_stream_runs_native_mps(): """The unified spinful model should evolve a charge-conserving MPS.""" fermion = Fermion(spinful=True, symmetry="U1U1") @@ -698,6 +813,20 @@ def test_unified_spinful_fermion_supports_symmray_parity_symmetries(symmetry): assert type(fermion.interaction_gate(0.01, U=2.0)).__name__.endswith("FermionicArray") +def test_z2z2_lattice_half_filling_keeps_flavor_occupations(): + """Z2Z2 site metadata must preserve up/down parity separately.""" + fermion = Fermion(spinful=True, symmetry="Z2Z2") + + setup = fermion.lattice_half_filling(2, 2) + + assert all( + isinstance(occupation, tuple) and len(occupation) == 2 + for occupation in setup.occupations.values() + ) + assert setup.target_charge == (0, 0) + assert setup.site_charge((0, 0)) == (1, 0) + + def test_unified_spinful_fermion_hamiltonian_keeps_explicit_mu_parameter(): """Explicit spinful chemical potentials must reach Symmray terms.""" fermion = Fermion(spinful=True, symmetry="U1U1") @@ -908,6 +1037,29 @@ def test_fermion_spin_gates_preserve_torch_backend(): assert gate_.backend == "torch" +def test_fermion_fields_and_pairing_preserve_jax_backend(): + """The extended stream terms keep their configured JAX block backend.""" + jnp = pytest.importorskip("jax.numpy") + + spinful = Fermion( + spinful=True, + symmetry="U1", + to_backend=pepsy.backend_jax(dtype=jnp.complex64), + ) + spinful_ham = spinful.hamiltonian( + ((0, 1),), t=1.0, U=2.0, field_z=0.2 + ) + assert all(term.backend == "jax" for term in spinful_ham.terms.values()) + + spinless = Fermion( + spinful=False, + symmetry="Z2", + to_backend=pepsy.backend_jax(dtype=jnp.complex64), + ) + pairing_ham = spinless.hamiltonian(((0, 1),), t=1.0, pairing=0.2) + assert all(term.backend == "jax" for term in pairing_ham.terms.values()) + + def test_mps_energy_uses_explicit_fermion_terms_natively(): """An explicit one- plus two-site SymHamiltonian stays on native MPS terms.""" fermion = Fermion(spinful=True, symmetry="U1U1") @@ -2395,6 +2547,175 @@ def test_native_fermionic_mps_rejects_bosonic_mpo_without_opt_in(): ).energy() +def test_native_fermionic_mpo_energy_is_factorized_on_l12(monkeypatch): + """Native MPO energy remains exact without forming a global operator.""" + L = 12 + state = SymMPS.for_model( + "fermi_hubbard_u1u1", + L, + bond_dim=4, + site_charge=site_charge_from_occupations( + [(1, 0) if site % 2 == 0 else (0, 1) for site in range(L)] + ), + seed=222, + dtype="complex128", + ) + fermion = Fermion(spinful=True, symmetry="U1U1") + hamiltonian = fermion.hamiltonian( + [(site, site + 1) for site in range(L - 1)], + t=1.0, + U=2.0, + mu=0.1, + ) + mpo = fermion.to_mpo( + hamiltonian=hamiltonian, + L=L, + compress=False, + ) + monkeypatch.setattr( + mpo, + "to_dense", + lambda: pytest.fail("native MPO energy must remain factorized"), + ) + + mpo_energy = pepsy.MpsEnergyOptimizer( + state, + mpo, + energy_per_site=False, + real=False, + ).energy().energy + term_energy = pepsy.MpsEnergyOptimizer( + state, + hamiltonian.terms, + energy_per_site=False, + real=False, + ).energy().energy + + assert complex(mpo_energy) == pytest.approx(complex(term_energy)) + + +def test_native_mpo_energy_reuses_paths_and_supports_controlled_compression(): + """Native MPO path caching and bounded compression preserve energy accuracy.""" + L = 12 + state = SymMPS.for_model( + "fermi_hubbard_u1u1", + L, + bond_dim=8, + site_charge=site_charge_from_occupations( + [(1, 0) if site % 2 == 0 else (0, 1) for site in range(L)] + ), + seed=123, + dtype="complex128", + ) + fermion = Fermion(spinful=True, symmetry="U1U1") + hamiltonian = fermion.hamiltonian( + [(site, site + 1) for site in range(L - 1)], + t=1.0, + U=2.0, + mu=0.1, + ) + mpo = fermion.to_mpo(hamiltonian=hamiltonian, L=L, compress=False) + optimizer = pepsy.MpsEnergyOptimizer( + state, + mpo, + energy_per_site=False, + real=False, + ) + + exact = optimizer.energy().energy + path_optimizer = optimizer._native_mpo_path_optimizer + assert path_optimizer is not None + assert path_optimizer.last_opt is not None + + repeated = optimizer.energy().energy + assert optimizer._native_mpo_path_optimizer is path_optimizer + assert complex(repeated) == pytest.approx(complex(exact)) + + compressed = optimizer.energy( + native_mpo_compression={ + "max_bond": 64, + "cutoff": 1e-12, + "method": "svd", + } + ).energy + assert complex(compressed) == pytest.approx(complex(exact), abs=1e-10) + compressed_estimate = optimizer.energy( + native_mpo_compression={"max_bond": 64, "cutoff": 1e-12} + ) + assert compressed_estimate.metadata["native_mpo_compression"]["max_bond"] == 64 + + with pytest.raises(ValueError, match="requires an explicit max_bond"): + optimizer.energy(native_mpo_compression={"cutoff": 1e-12}) + with pytest.raises(ValueError, match="requires an explicit max_bond"): + optimizer.energy(native_mpo_compression={}) +@pytest.mark.parametrize( + ("spinful", "symmetry", "model"), + [ + (False, "U1", "fermi_hubbard_spinless"), + (False, "Z2", "fermi_hubbard_spinless"), + (True, "U1", "fermi_hubbard"), + (True, "U1U1", "fermi_hubbard_u1u1"), + ], +) +def test_native_factorized_mpo_matches_terms_and_jw_reference( + spinful, symmetry, model +): + """Native MPO energies agree with terms and the separate JW operator oracle.""" + L = 4 + edges = [(site, site + 1) for site in range(L - 1)] + occupations = ( + [1] * L + if symmetry == "U1" and spinful + else ([(1, 0) if site % 2 == 0 else (0, 1) for site in range(L)] + if spinful + else [1, 0, 1, 0]) + ) + state = SymMPS.for_model( + model, + L, + symmetry=symmetry, + bond_dim=4, + site_charge=site_charge_from_occupations(occupations), + seed=111, + dtype="complex128", + ) + fermion = Fermion(spinful=spinful, symmetry=symmetry) + if spinful: + hamiltonian = fermion.hamiltonian(edges, t=1.0, U=2.0, mu=0.1) + else: + hamiltonian = fermion.hamiltonian(edges, t=1.0, V=0.4, mu=0.1) + + native_mpo = fermion.to_mpo( + hamiltonian=hamiltonian, + L=L, + compress=False, + ) + jw_mpo = hamiltonian.to_mpo( + L=L, + fermionic=False, + compress=False, + ) + native_energy = pepsy.MpsEnergyOptimizer( + state, + native_mpo, + energy_per_site=False, + real=False, + ).energy().energy + term_energy = pepsy.MpsEnergyOptimizer( + state, + hamiltonian.terms, + energy_per_site=False, + real=False, + ).energy().energy + + assert complex(native_energy) == pytest.approx(complex(term_energy)) + np.testing.assert_allclose( + _mpo_to_dense_matrix(native_mpo, L), + _mpo_to_dense_matrix(jw_mpo, L), + atol=1e-10, + ) + + def test_fermi_hubbard_u1u1_mpo_energy_matches_high_bond_fermionic_mps(): """High-bond fermionic FH MPS energy should use the direct MPO.""" state = SymMPS.for_model( From faffa9faef5589e17cffb751c5076561be2d78a3 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Wed, 29 Jul 2026 20:52:12 -0600 Subject: [PATCH 24/70] Add Symmray-aware D2BP support --- .github/skills/belief-propagation/SKILL.md | 23 +- src/pepsy/bp/cluster.py | 46 ++++ src/pepsy/bp/gauges.py | 304 +++++++++++++++++++-- src/pepsy/bp/relay.py | 113 +++++++- tests/test_bp_symmray.py | 188 +++++++++++++ 5 files changed, 646 insertions(+), 28 deletions(-) create mode 100644 tests/test_bp_symmray.py diff --git a/.github/skills/belief-propagation/SKILL.md b/.github/skills/belief-propagation/SKILL.md index 29dd0c3..2994e04 100644 --- a/.github/skills/belief-propagation/SKILL.md +++ b/.github/skills/belief-propagation/SKILL.md @@ -31,6 +31,10 @@ over `quimb.tensor.belief_propagation`; the annotated paper trail lives in - `one_norm_bp(tn, *, method="l1bp"|"hv1bp"|"d1bp", max_iterations, tol, damping, update, diis, init_messages, ...) -> RelayBPResult` — plain 1-norm BP to a fixed point (partition-function / nonnegative contractions, e.g. decoding). +- `two_norm_bp(tn, *, max_iterations, tol, damping, update, diis, + init_messages, ...) -> RelayBPResult` — native D2BP for wavefunction and + norm contractions. Symmray-backed fermionic PEPS retain their native charge + blocks throughout message updates and distance evaluation. - `relay_bp(tn, *, method, num_relays, max_iterations, gamma_range, tol, damping, update, memory_first_leg, init_messages, seed, ...) -> RelayBPResult` — disordered-memory / relay-BP: per-node random memory (incl. negative) applied @@ -63,8 +67,12 @@ over `quimb.tensor.belief_propagation`; the annotated paper trail lives in returned gauge-transformed network. - SU/simple-gauge bridge helpers: `simple_update_messages_from_gauges`, `d1bp_from_simple_update_gauges`, `run_d1bp_from_simple_update_gauges`, - `simple_update_bp_residual`, and `norm1_gloop_expand`. Use these for - scalar 1-norm cluster work with externally supplied simple gauges. + `simple_update_bp_residual`, `d2bp_from_simple_update_gauges`, + `run_d2bp_from_simple_update_gauges`, + `simple_update_core_and_gauges_from_d2bp`, and `norm1_gloop_expand`. Use + these for scalar 1-norm or PEPS 2-norm work with externally supplied simple + gauges. The D2BP bridge supports dense tensors and native Symmray fermionic + `U1`, `U1U1`, and `Z2` block-sparse tensors. - Result dataclasses: `RelayBPResult` (`.bp`, `.converged`, `.iterations`, `.max_mdiff`, `.contract()`, `.messages`, `.snapshot()`), `LoopClusterResult` (`.estimate`, `.bp_converged`, `.bp_iterations`, `.expand()`, `.messages`). @@ -84,6 +92,13 @@ PNE terms are not loop-cluster regions. A PNE residue is the all-`Q` network; retaining it gives the exact projector identity, while dropping it is the approximation whose error should be monitored. +For native Symmray D2BP, directed messages can temporarily omit charge +sectors after an update. Pepsy aligns message charge maps before distance +checks, D2BP normalization, and loop-cluster expansion, while keeping the +messages and tensor data block-sparse. Do not replace this path with a dense +copy: dense eigendecomposition can reorder eigenvectors across charge sectors +and produce invalid Symmray gates. + ## quimb substrate (do not reimplement) `quimb.tensor.belief_propagation`: - BP families: **1-norm** `L1BP` / `HV1BP` / `D1BP` (partition-function / @@ -165,6 +180,10 @@ Loop corrections split by how much they need a **converged BP fixed point**: - `weight_pass` is intentionally restricted to closed pairwise networks; for a D2 calculation, obtain the environment on the appropriate closed double-layer network before supplying projectors to D2 PNE. +- Native Symmray fermionic BP coverage is in `tests/test_bp_symmray.py` and + currently exercises `U1`, `U1U1`, and `Z2` SU↔D2BP round trips plus loop + cluster expansion. Preserve the array class and charge blocks when adding + new BP or gauge paths. - Keep `pepsy.bp` out of the lazy top-level namespace (import `pepsy.bp`); do **not** edit `src/pepsy/__init__.py` for it. - Tests: `tests/test_bp_relay.py` and `tests/test_simple_update_gen.py`. Env: py312 diff --git a/src/pepsy/bp/cluster.py b/src/pepsy/bp/cluster.py index ce6d9be..0f1f87b 100644 --- a/src/pepsy/bp/cluster.py +++ b/src/pepsy/bp/cluster.py @@ -652,6 +652,51 @@ def _expand(bp, norm, gloops, combine, optimize, strip_exponent, progbar): return bp.contract_gloop_expand(**kwargs) +def _align_symmray_d2bp_messages(bp) -> None: + """Pad omitted Symmray charge blocks before quimb D2BP normalization. + + Quimb's D2BP normalization treats the two directed messages on an + internal bond as dense vectors with the same shape. Symmray instead + stores only non-zero charge blocks, so a rank-deficient message can have a + smaller shape than its partner even though both messages live on the same + virtual bond. ``normalize_message_pair`` then attempts a matrix product + and raises a shape-mismatch error. + + Reindex each message onto the native bond charge map and explicitly add + the missing sectors as zero blocks. The dual flags remain those of the + original message, so graded contractions and fermionic phases are left to + Symmray. Dense or non-Symmray messages are deliberately untouched. + """ + for index, tids in bp.tn.ind_map.items(): + if len(tids) != 2: + continue + + messages = tuple((index, tid, bp.messages[index, tid]) for tid in tids) + if not all( + getattr(message.__class__, "__module__", "").startswith("symmray") + and hasattr(message, "indices") + and hasattr(message, "copy_with") + and hasattr(message, "fill_missing_blocks") + for _, _, message in messages + ): + continue + + # The tensor leg is the authoritative charge map. The two messages + # may carry different subsets of these sectors after BP updates. + tid = next(iter(tids)) + tensor = bp.tn.tensor_map[tid] + axis = tensor.inds.index(index) + bond_index = tensor.data.indices[axis] + for message_index, message_tid, message in messages: + target_indices = tuple( + bond_index.copy_with(dual=message_index.dual) + for message_index in message.indices + ) + aligned = message.copy_with(indices=target_indices) + aligned.fill_missing_blocks() + bp.messages[index, message_tid] = aligned + + @dataclass class LoopClusterResult: """Result of a loop cluster expansion contraction. @@ -1321,6 +1366,7 @@ def loop_cluster_expand( autocomplete=autocomplete, ) else: + _align_symmray_d2bp_messages(bp) estimate = _expand( bp, key, gloops, combine, optimize, strip_exponent, progbar ) diff --git a/src/pepsy/bp/gauges.py b/src/pepsy/bp/gauges.py index 9853d4a..4fa29a1 100644 --- a/src/pepsy/bp/gauges.py +++ b/src/pepsy/bp/gauges.py @@ -97,12 +97,139 @@ def _copy_array(x): def _as_numpy(x): + if hasattr(x, "to_dense"): + return np.asarray(x.to_dense()) try: return np.asarray(ar.to_numpy(x)) except Exception: return np.asarray(x) +def _gauge_values_numpy(gauge): + """Materialize a gauge vector for validation, including Symmray vectors.""" + if hasattr(gauge, "to_dense"): + return np.asarray(gauge.to_dense()) + return _as_numpy(gauge) + + +def _is_symmray_array(value) -> bool: + return getattr(value.__class__, "__module__", "").startswith("symmray") + + +def _symmray_dense_matrix(matrix): + """Convert one small Symmray message to a host dense matrix.""" + if hasattr(matrix, "to_dense"): + return np.asarray(matrix.to_dense()) + return np.asarray(matrix) + + +def _symmray_align_message_to_bond(tn, ix, tid, message): + """Pad a Symmray message to the full charge support of a PEPS bond.""" + if not (_is_symmray_array(message) and hasattr(message, "indices")): + return message + + charge_map = _symmray_bond_chargemap(tn, ix) + indices = tuple( + message_index.copy_with(chargemap=charge_map) + for message_index in message.indices + ) + message = message.copy_with(indices=indices) + message.fill_missing_blocks() + return message + + +def _symmray_bond_chargemap(tn, ix): + """Return the union of endpoint charge maps for one PEPS bond.""" + charge_map = {} + for tid in tn.ind_map[ix]: + tensor = tn.tensor_map[tid] + axis = tensor.inds.index(ix) + for charge, size in tensor.data.indices[axis].chargemap.items(): + previous = charge_map.setdefault(charge, int(size)) + if previous != int(size): + raise ValueError( + f"incompatible endpoint charge dimensions on bond {ix!r}" + ) + return dict(sorted(charge_map.items())) + + +def _symmray_block_vector(tn, ix, values, *, tid=None): + """Create a native Symmray block vector in the bond charge order. + + ``tid`` optionally restricts the result to the charge sectors present in + that endpoint's current sparse tensor data. This matters for fermionic + boundary tensors: their index can retain a larger declared charge map than + the sectors with nonzero stored blocks. + """ + import symmray as sr + + charge_map = _symmray_bond_chargemap(tn, ix) + values = np.asarray(values).reshape(-1) + offsets = {} + offset = 0 + for charge, size in charge_map.items(): + offsets[charge] = offset + offset += int(size) + if offset != values.size: + raise ValueError(f"vector size does not match Symmray bond {ix!r}") + + if tid is None: + selected_charges = tuple(charge_map) + else: + tensor = tn.tensor_map[tid] + axis = tensor.inds.index(ix) + selected_charges = tuple(tensor.data.indices[axis].chargemap) + + blocks = {} + for charge in selected_charges: + size = int(charge_map[charge]) + offset = offsets[charge] + size = int(size) + blocks[charge] = values[offset : offset + size].copy() + return sr.BlockVector(blocks) + + +def _symmray_block_matrix(tn, ix, tid, matrix, *, full=False): + """Create a native Symmray matrix from a charge-preserving dense matrix. + + ``matrix`` is indexed in the union charge order of the PEPS bond. By + default only sectors present in ``tid``'s current sparse tensor data are + emitted, while ``full=True`` retains every endpoint-supported sector for a + standalone density matrix or eigendecomposition. + """ + tensor = tn.tensor_map[tid] + data = tensor.data + axis = tensor.inds.index(ix) + bond_index = data.indices[axis] + matrix = np.asarray(matrix) + charge_map = _symmray_bond_chargemap(tn, ix) + selected_charges = ( + tuple(charge_map) + if full + else tuple(bond_index.chargemap) + ) + offsets = {} + offset = 0 + for charge, size in charge_map.items(): + offsets[charge] = offset + offset += int(size) + if offset != matrix.shape[0] or matrix.shape[0] != matrix.shape[1]: + raise ValueError(f"matrix size does not match Symmray bond {ix!r}") + blocks = {} + for charge in selected_charges: + size = int(charge_map[charge]) + offset = offsets[charge] + size = int(size) + blocks[(charge, charge)] = matrix[ + offset : offset + size, offset : offset + size + ].copy() + return type(data).from_blocks( + blocks, + duals=(bond_index.dual, not bond_index.dual), + phases={}, + ) + + def _as_float(x) -> float: return float(np.asarray(_as_numpy(x))) @@ -294,13 +421,15 @@ def _d2bp_messages_from_simple_update_gauges( if ix in gauges: gauge = _copy_array(gauges[ix]) elif missing == "ones": - gauge = _ones_for_index(tn, ix) + gauge = None elif missing == "raise": raise KeyError(f"missing simple-update gauge for index {ix!r}") else: raise ValueError("missing must be 'ones' or 'raise'") - gauge_np = np.real_if_close(_as_numpy(gauge)) + gauge_np = np.ones(tn.ind_size(ix)) if gauge is None else np.real_if_close( + _gauge_values_numpy(gauge) + ) if gauge_np.ndim != 1 or gauge_np.shape[0] != tn.ind_size(ix): raise ValueError( f"SU gauge for {ix!r} must be a length-{tn.ind_size(ix)} vector" @@ -312,20 +441,68 @@ def _d2bp_messages_from_simple_update_gauges( "D2BP SU initialization requires real nonnegative Vidal " f"gauges; bond {ix!r} is invalid" ) - if smudge: + if smudge and gauge is not None: gauge = _smudge_gauge(gauge, smudge) + tida, tidb = tids # In the symmetric PEPS gauge, sqrt(lambda) is absorbed on each # physical-site tensor. D2BP sees both layers, hence its directed # density message is diag(lambda), rather than the D1 sqrt(lambda). - message = ar.do("diag", gauge) - tida, tidb = tids - messages[ix, tida] = _copy_array(message) - messages[ix, tidb] = _copy_array(message) + messages[ix, tida] = _d2bp_diagonal_message( + tn, ix, tida, gauge, smudge=smudge + ) + messages[ix, tidb] = _d2bp_diagonal_message( + tn, ix, tidb, gauge, smudge=smudge + ) return messages +def _d2bp_diagonal_message(tn, ix, tid, gauge, *, smudge=0.0): + """Build one SU ``diag(lambda)`` D2BP message in the native backend.""" + tensor = tn.tensor_map[tid] + data = tensor.data + if gauge is None: + gauge_values = np.ones(tn.ind_size(ix)) + smudge + else: + gauge_values = gauge + + if not ( + getattr(data.__class__, "__module__", "").startswith("symmray") + and hasattr(data, "indices") + ): + return _copy_array(ar.do("diag", gauge_values)) + + axis = tensor.inds.index(ix) + bond_index = data.indices[axis] + gauge_blocks = gauge_values.blocks if hasattr(gauge_values, "blocks") else {} + blocks = {} + for charge, size in bond_index.chargemap.items(): + values = gauge_blocks.get(charge) + if values is None: + fill = 1.0 if gauge is None else 0.0 + values = ar.do( + "ones" if fill else "zeros", + (int(size),), + like=data.get_any_array(), + ) + if fill and smudge: + values = values * (1.0 + smudge) + else: + values = ar.do("reshape", values, (int(size),)) + blocks[(charge, charge)] = ar.do("diag", values) + + # The message axes must match the destination tensor leg and its dual. + # Do not copy the PEPS tensor's dummy modes: these auxiliary density + # messages are not physical fermion legs and must have no dummy mode. + message_cls = type(data) + return message_cls.from_blocks( + blocks, + duals=(bond_index.dual, not bond_index.dual), + phases={}, + ) + + def d2bp_from_simple_update_gauges( tn, gauges=None, @@ -1419,8 +1596,37 @@ def _psd_eigh(matrix, *, smudge: float, label: str): """Diagonalize a PSD message, clipping only numerical null modes.""" if smudge < 0.0: raise ValueError("smudge must be nonnegative") - matrix = _hermitize(matrix) - values, vectors = ar.do("linalg.eigh", matrix) + if _is_symmray_array(matrix): + # A dense eigh is not symmetry aware: even a diagonal matrix can have + # its eigenvectors returned in a charge-permuting order. That is + # harmless for ordinary arrays but dropping the resulting off-sector + # gates back into a Symmray array changes the state. Diagonalize each + # charge block independently and keep the native block order. + charge_map = matrix.indices[0].chargemap + size = sum(int(n) for n in charge_map.values()) + dtype = np.result_type(*matrix.get_all_blocks(), float) + values = np.empty(size, dtype=float) + vectors = np.zeros((size, size), dtype=dtype) + offset = 0 + for charge, block_size in charge_map.items(): + block_size = int(block_size) + block = matrix.blocks.get((charge, charge)) + if block is None: + block = np.zeros((block_size, block_size), dtype=dtype) + block = np.asarray(block) + block = 0.5 * (block + block.conj().T) + block_values, block_vectors = np.linalg.eigh(block) + values[offset : offset + block_size] = np.real_if_close( + block_values + ) + vectors[ + offset : offset + block_size, + offset : offset + block_size, + ] = block_vectors + offset += block_size + else: + matrix = _hermitize(matrix) + values, vectors = ar.do("linalg.eigh", matrix) values_np = np.real_if_close(_as_numpy(values)) if np.iscomplexobj(values_np) or not np.all(np.isfinite(values_np)): raise ValueError(f"D2BP message for {label} has invalid eigenvalues") @@ -1449,8 +1655,12 @@ def _psd_sqrt_and_inverse(matrix, *, smudge: float, label: str): values, vectors = _psd_eigh(matrix, smudge=smudge, label=label) roots = ar.do("sqrt", values) - sqrt = qtn.decomp.rdmul(vectors, roots) @ ar.dag(vectors) - sqrt_inv = qtn.decomp.rdmul(vectors, 1.0 / roots) @ ar.dag(vectors) + if _is_symmray_array(matrix): + sqrt = (vectors * roots) @ vectors.conj().T + sqrt_inv = (vectors * (1.0 / roots)) @ vectors.conj().T + else: + sqrt = qtn.decomp.rdmul(vectors, roots) @ ar.dag(vectors) + sqrt_inv = qtn.decomp.rdmul(vectors, 1.0 / roots) @ ar.dag(vectors) return sqrt, sqrt_inv @@ -1510,14 +1720,32 @@ def simple_update_core_and_gauges_from_d2bp( # solve the simultaneous congruence problem on this bond. m_from_a = bp.messages[ix, tidb] m_from_b = bp.messages[ix, tida] + if _is_symmray_array(m_from_a): + m_from_a = _symmray_align_message_to_bond( + bp.tn, ix, tida, m_from_a + ) + m_from_b = _symmray_align_message_to_bond( + bp.tn, ix, tidb, m_from_b + ) sqrt_a, sqrt_a_inv = _psd_sqrt_and_inverse( m_from_a, smudge=smudge, label=f"bond {ix!r}, source {tida!r}", ) - metric_product = _hermitize( - sqrt_a @ ar.do("transpose", m_from_b) @ sqrt_a - ) + symmray_messages = _is_symmray_array(m_from_a) + if symmray_messages: + m_from_b_dense = _symmray_dense_matrix(m_from_b) + metric_product = sqrt_a @ m_from_b_dense.T @ sqrt_a + metric_product = 0.5 * ( + metric_product + metric_product.conj().T + ) + metric_product = _symmray_block_matrix( + bp.tn, ix, tida, metric_product, full=True + ) + else: + metric_product = _hermitize( + sqrt_a @ ar.do("transpose", m_from_b) @ sqrt_a + ) lambda_squared, vectors = _psd_eigh( metric_product, smudge=smudge, @@ -1532,13 +1760,47 @@ def simple_update_core_and_gauges_from_d2bp( # The transposed first gate and inverse second gate preserve the # single-layer contraction exactly; removing sqrt(Lambda) from both # tensors exposes the external SU gauge. - G = sqrt_a_inv @ qtn.decomp.rdmul(vectors, sqrt_gauge) - G_inv = qtn.decomp.lddiv(sqrt_gauge, ar.dag(vectors)) @ sqrt_a - core.tensor_map[tida].gate_(ar.do("transpose", G), ix) - core.tensor_map[tidb].gate_(G_inv, ix) - core.tensor_map[tida].multiply_index_diagonal_(ix, 1.0 / sqrt_gauge) - core.tensor_map[tidb].multiply_index_diagonal_(ix, 1.0 / sqrt_gauge) - gauges[ix] = _copy_array(gauge) + if symmray_messages: + G = ( + sqrt_a_inv + @ (vectors * sqrt_gauge) + ) + G_inv = ( + np.diag(1.0 / sqrt_gauge) + @ vectors.conj().T + @ sqrt_a + ) + core.tensor_map[tida].gate_( + _symmray_block_matrix(core, ix, tida, G.T), ix + ) + core.tensor_map[tidb].gate_( + _symmray_block_matrix(core, ix, tidb, G_inv), ix + ) + inv_sqrt_gauge = _symmray_block_vector( + core, ix, 1.0 / sqrt_gauge, tid=tida + ) + core.tensor_map[tida].multiply_index_diagonal_( + ix, inv_sqrt_gauge + ) + inv_sqrt_gauge = _symmray_block_vector( + core, ix, 1.0 / sqrt_gauge, tid=tidb + ) + core.tensor_map[tidb].multiply_index_diagonal_( + ix, inv_sqrt_gauge + ) + gauges[ix] = _symmray_block_vector(core, ix, gauge) + else: + G = sqrt_a_inv @ qtn.decomp.rdmul(vectors, sqrt_gauge) + G_inv = qtn.decomp.lddiv(sqrt_gauge, ar.dag(vectors)) @ sqrt_a + core.tensor_map[tida].gate_(ar.do("transpose", G), ix) + core.tensor_map[tidb].gate_(G_inv, ix) + core.tensor_map[tida].multiply_index_diagonal_( + ix, 1.0 / sqrt_gauge + ) + core.tensor_map[tidb].multiply_index_diagonal_( + ix, 1.0 / sqrt_gauge + ) + gauges[ix] = _copy_array(gauge) return core, gauges diff --git a/src/pepsy/bp/relay.py b/src/pepsy/bp/relay.py index 270b99c..b9edb15 100644 --- a/src/pepsy/bp/relay.py +++ b/src/pepsy/bp/relay.py @@ -52,6 +52,56 @@ _RELAY_METHODS = {"l1bp", "d1bp", "d2bp"} +def _is_symmray_array(value) -> bool: + return getattr(value.__class__, "__module__", "").startswith("symmray") + + +def _align_symmray_message_pair(left, right): + """Return charge-support-aligned copies of two Symmray messages.""" + if not ( + _is_symmray_array(left) + and _is_symmray_array(right) + and hasattr(left, "indices") + and hasattr(right, "indices") + ): + return left, right + + left_indices = [] + right_indices = [] + for left_index, right_index in zip(left.indices, right.indices): + left_map = dict(left_index.chargemap) + right_map = dict(right_index.chargemap) + for charge in set(left_map) & set(right_map): + if left_map[charge] != right_map[charge]: + raise ValueError("incompatible Symmray message charge dimensions") + charge_map = {**left_map, **right_map} + left_indices.append( + left_index.copy_with(chargemap=charge_map, dual=left_index.dual) + ) + right_indices.append( + right_index.copy_with(chargemap=charge_map, dual=right_index.dual) + ) + + left = left.copy_with(indices=tuple(left_indices)) + right = right.copy_with(indices=tuple(right_indices)) + left.fill_missing_blocks() + right.fill_missing_blocks() + return left, right + + +def _symmray_message_distance(left, right) -> float: + """L2 distance after aligning omitted sparse charge blocks.""" + left, right = _align_symmray_message_pair(left, right) + if hasattr(left, "to_dense") and hasattr(right, "to_dense"): + left = left.to_dense() + right = right.to_dense() + return float(np.linalg.norm(np.asarray(left) - np.asarray(right))) + + +def _uses_symmray(tn) -> bool: + return any(_is_symmray_array(tensor.data) for tensor in tn.tensors) + + def _method_key(method: str, classes: Mapping[str, str]) -> str: """Validate and normalize a public BP method name against ``classes``.""" key = str(method).lower() @@ -77,12 +127,26 @@ def _message_data(message): Checking for ``modify`` deliberately avoids ``ndarray.data``, which is a memory-view rather than the message array. """ - return message.data if hasattr(message, "modify") else message + return ( + message.data + if hasattr(message, "modify") and hasattr(message, "data") + else message + ) + + +def _copy_message(message): + """Copy a message through its native backend before using Autoray.""" + if hasattr(message, "copy"): + try: + return message.copy() + except Exception: + pass + return ar.do("copy", message) def _set_message(bp, key, message, data) -> None: """Replace one message, supporting Tensor and bare-array BPs.""" - if hasattr(message, "modify"): + if hasattr(message, "modify") and hasattr(message, "data"): message.modify(data=data) else: bp.messages[key] = data @@ -100,7 +164,7 @@ def _snapshot(messages): return {key: _snapshot(value) for key, value in messages.items()} if isinstance(messages, tuple): return tuple(_snapshot(value) for value in messages) - return ar.do("copy", _message_data(messages)) + return _copy_message(_message_data(messages)) def _message_shape(message) -> tuple[int, ...]: @@ -108,6 +172,37 @@ def _message_shape(message) -> tuple[int, ...]: return tuple(ar.do("shape", _message_data(message))) +def _symmray_message_compatible(template, snapshot) -> bool: + """Allow sparse charge support to differ while preserving bond topology.""" + if not ( + _is_symmray_array(template) + and _is_symmray_array(snapshot) + and hasattr(template, "indices") + and hasattr(snapshot, "indices") + ): + return False + if len(template.indices) != len(snapshot.indices): + return False + for template_index, snapshot_index in zip( + template.indices, snapshot.indices + ): + if template_index.dual != snapshot_index.dual: + return False + template_map = template_index.chargemap + snapshot_map = snapshot_index.chargemap + if not ( + set(template_map).issubset(snapshot_map) + or set(snapshot_map).issubset(template_map) + ): + return False + if any( + template_map[charge] != snapshot_map[charge] + for charge in set(template_map) & set(snapshot_map) + ): + return False + return True + + def _validate_message_tree(template, snapshot, path="messages") -> None: """Check that a snapshot exactly matches a BP message layout.""" if isinstance(template, Mapping): @@ -129,7 +224,9 @@ def _validate_message_tree(template, snapshot, path="messages") -> None: _validate_message_tree(value, saved, f"{path}[{i}]") return - if _message_shape(template) != _message_shape(snapshot): + if _message_shape(template) != _message_shape(snapshot) and not ( + _symmray_message_compatible(template, snapshot) + ): raise ValueError( f"{path} has shape {_message_shape(snapshot)}, expected " f"{_message_shape(template)} for this tensor-network topology" @@ -153,7 +250,7 @@ def _set_messages(bp, messages) -> None: return for key, message in bp.messages.items(): - _set_message(bp, key, message, ar.do("copy", messages[key])) + _set_message(bp, key, message, _copy_message(messages[key])) def _bp_constructor_kwargs(method_key: str, damping, update, bp_opts): @@ -566,6 +663,9 @@ def two_norm_bp( """ if not isinstance(max_iterations, (int, np.integer)) or max_iterations < 1: raise ValueError("max_iterations must be a positive integer") + bp_opts = dict(bp_opts) + if _uses_symmray(tn): + bp_opts.setdefault("distance", _symmray_message_distance) bp = _bp_class("d2bp")( tn, **_bp_constructor_kwargs("d2bp", damping, update, bp_opts), @@ -684,6 +784,9 @@ def relay_bp( _validate_d1_graph(tn) + bp_opts = dict(bp_opts) + if method_key == "d2bp" and _uses_symmray(tn): + bp_opts.setdefault("distance", _symmray_message_distance) bp_class = _bp_class(method_key) rng = np.random.default_rng(seed) bp = bp_class( diff --git a/tests/test_bp_symmray.py b/tests/test_bp_symmray.py new file mode 100644 index 0000000..2900873 --- /dev/null +++ b/tests/test_bp_symmray.py @@ -0,0 +1,188 @@ +"""Symmray fermionic coverage for Pepsy's 2-norm BP corrections.""" + +from __future__ import annotations + +import numpy as np +import pytest + +pytest.importorskip("symmray") + +from pepsy.bp import gauge_all, loop_cluster_expand, two_norm_bp # noqa: E402 +from pepsy.tensors import ( # noqa: E402 + SymPEPS, + site_charge_alternating, +) + + +@pytest.mark.parametrize("bond_dim", (2, 3)) +def test_fermionic_u1_two_norm_bp_is_exact_on_a_tree(bond_dim): + """D2BP preserves the native graded contraction on a PEPS tree.""" + state = SymPEPS.random( + 1, + 4, + symmetry="U1", + bond_dim=bond_dim, + phys_dim=2, + fermionic=True, + seed=700 + bond_dim, + dtype="complex128", + ) + + exact = complex(state.norm()) + bp = two_norm_bp( + state.tn, + max_iterations=100, + tol=1e-10, + diis=False, + ) + + assert bp.converged + assert abs(complex(bp.contract()) - exact) <= 1e-10 * max(1.0, abs(exact)) + + corrected = loop_cluster_expand( + state.tn, + gloops=0, + norm="2norm", + max_iterations=100, + tol=1e-10, + diis=False, + ) + assert corrected.bp_converged + assert abs(complex(corrected.estimate) - exact) <= 1e-10 * max( + 1.0, abs(exact) + ) + + +@pytest.mark.parametrize("bond_dim", (2, 3)) +def test_fermionic_u1_loop_cluster_runs_on_3x4_peps(bond_dim): + """D2BP loop clusters pad sparse charge blocks without densifying.""" + state = SymPEPS.random( + 3, + 4, + symmetry="U1", + bond_dim=bond_dim, + phys_dim=2, + fermionic=True, + seed=800 + bond_dim, + dtype="complex128", + ) + + result = loop_cluster_expand( + state.tn, + gloops=4, + norm="2norm", + max_iterations=300, + tol=1e-10, + diis=False, + ) + + assert result.bp_converged + assert result.bp.__class__.__name__ == "D2BP" + assert np.isfinite(float(np.real(result.estimate))) + assert all( + type(message).__name__ == "U1FermionicArray" + for message in result.messages.values() + ) + assert all( + type(tensor.data).__name__ == "U1FermionicArray" + for tensor in state.tn.tensors + ) + + +def _fermionic_symmetry_cases(): + return ( + ( + "U1", + "fermi_hubbard", + {"symmetry": "U1", "phys_dim": 4}, + ), + ( + "U1U1", + "fermi_hubbard_u1u1", + { + "symmetry": "U1U1", + "phys_dim": 4, + "site_charge": site_charge_alternating( + (1, 0), (0, 1) + ), + }, + ), + ( + "Z2", + "fermi_hubbard_spinless", + {"symmetry": "Z2", "phys_dim": 2}, + ), + ) + + +@pytest.mark.parametrize( + ("label", "model", "state_options"), + _fermionic_symmetry_cases(), +) +def test_fermionic_sympeps_su_d2bp_bridge_preserves_symmetry_blocks( + label, model, state_options +): + """SU gauges and BP loop expansion preserve native fermionic blocks.""" + state = SymPEPS.for_model( + model, + 3, + 4, + bond_dim=2, + seed=1200, + dtype="complex128", + **state_options, + ) + exact_norm = state.tn.norm() + + forward = gauge_all( + state.tn, + start="su", + target="bp", + norm="2norm", + su_options={"max_iterations": 8, "tol": 0.0}, + bp_options={ + "run_opts": { + "max_iterations": 300, + "tol": 1e-10, + "diis": False, + } + }, + ) + assert forward.bp.converged, label + reconstructed = forward.core.copy() + reconstructed.gauge_simple_insert(forward.gauges) + np.testing.assert_allclose(reconstructed.norm(), exact_norm, rtol=1e-10) + assert all( + type(message).__name__ == type(state.tn.tensors[0].data).__name__ + for message in forward.messages.values() + ) + + cluster = loop_cluster_expand( + forward.bp.tn, + gloops=4, + norm="2norm", + messages=forward.messages, + run_bp=False, + ) + assert np.isfinite(float(np.real(cluster.estimate))), label + + reverse = gauge_all( + state.tn, + start="bp", + target="su", + norm="2norm", + bp_options={ + "run_opts": { + "max_iterations": 300, + "tol": 1e-10, + "diis": False, + } + }, + conversion_options={"smudge": 1e-12}, + ) + assert reverse.bp.converged, label + reverse_reconstructed = reverse.core.copy() + reverse_reconstructed.gauge_simple_insert(reverse.gauges) + np.testing.assert_allclose( + reverse_reconstructed.norm(), exact_norm, rtol=1e-10 + ) From 0b8d0e20ec54a38e101a8bd35ecd4ac6f256d2f8 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Wed, 29 Jul 2026 21:25:02 -0600 Subject: [PATCH 25/70] Complete Symmray BP compatibility --- .github/skills/belief-propagation/SKILL.md | 29 ++- src/pepsy/bp/_symmray.py | 258 +++++++++++++++++++++ src/pepsy/bp/cluster.py | 68 ++---- src/pepsy/bp/gauges.py | 21 ++ src/pepsy/bp/pne.py | 28 ++- src/pepsy/bp/relay.py | 80 +++---- src/pepsy/bp/series.py | 48 +++- src/pepsy/bp/weights.py | 28 +++ tests/test_bp_symmray.py | 138 ++++++++++- 9 files changed, 583 insertions(+), 115 deletions(-) create mode 100644 src/pepsy/bp/_symmray.py diff --git a/.github/skills/belief-propagation/SKILL.md b/.github/skills/belief-propagation/SKILL.md index 2994e04..54ce30d 100644 --- a/.github/skills/belief-propagation/SKILL.md +++ b/.github/skills/belief-propagation/SKILL.md @@ -34,7 +34,9 @@ over `quimb.tensor.belief_propagation`; the annotated paper trail lives in - `two_norm_bp(tn, *, max_iterations, tol, damping, update, diis, init_messages, ...) -> RelayBPResult` — native D2BP for wavefunction and norm contractions. Symmray-backed fermionic PEPS retain their native charge - blocks throughout message updates and distance evaluation. + blocks throughout message updates and distance evaluation; DIIS + automatically falls back to native sequential updates because Symmray does + not expose Quimb's dense vectorizer concatenation. - `relay_bp(tn, *, method, num_relays, max_iterations, gamma_range, tol, damping, update, memory_first_leg, init_messages, seed, ...) -> RelayBPResult` — disordered-memory / relay-BP: per-node random memory (incl. negative) applied @@ -92,12 +94,20 @@ PNE terms are not loop-cluster regions. A PNE residue is the all-`Q` network; retaining it gives the exact projector identity, while dropping it is the approximation whose error should be monitored. -For native Symmray D2BP, directed messages can temporarily omit charge +For native Symmray D2BP, relay, loop-series, and PNE paths, directed messages +can temporarily omit charge sectors after an update. Pepsy aligns message charge maps before distance -checks, D2BP normalization, and loop-cluster expansion, while keeping the -messages and tensor data block-sparse. Do not replace this path with a dense -copy: dense eigendecomposition can reorder eigenvectors across charge sectors -and produce invalid Symmray gates. +checks, D2BP normalization, loop-cluster/series expansion, and PNE projector +construction, while keeping the messages and tensor data block-sparse. Do not +replace this D2 path with a dense copy: dense eigendecomposition can reorder +eigenvectors across charge sectors and produce invalid Symmray gates. + +For valid closed scalar Symmray networks, the 1-norm APIs (`L1BP`, `HV1BP`, +and `D1BP`), their loop/series/PNE corrections, D1 SU bridge, and +`weight_pass` use a topology-identical dense BP shadow because Quimb's D1 +initializers and weight SVDs require dense scalar operations. A raw fermionic +PEPS with dangling physical legs is not a direct 1-norm input: use D2BP for +its wavefunction norm (or explicitly construct a valid closed scalar network). ## quimb substrate (do not reimplement) `quimb.tensor.belief_propagation`: @@ -181,9 +191,10 @@ Loop corrections split by how much they need a **converged BP fixed point**: D2 calculation, obtain the environment on the appropriate closed double-layer network before supplying projectors to D2 PNE. - Native Symmray fermionic BP coverage is in `tests/test_bp_symmray.py` and - currently exercises `U1`, `U1U1`, and `Z2` SU↔D2BP round trips plus loop - cluster expansion. Preserve the array class and charge blocks when adding - new BP or gauge paths. + exercises `U1`, `U1U1`, and `Z2` SU↔D2BP round trips, D2 relay/loop-series/ + PNE corrections, and closed-scalar 1-norm compatibility. Preserve the array + class and charge blocks when adding new D2 BP or gauge paths; route valid + scalar D1 paths through the documented dense shadow. - Keep `pepsy.bp` out of the lazy top-level namespace (import `pepsy.bp`); do **not** edit `src/pepsy/__init__.py` for it. - Tests: `tests/test_bp_relay.py` and `tests/test_simple_update_gen.py`. Env: py312 diff --git a/src/pepsy/bp/_symmray.py b/src/pepsy/bp/_symmray.py new file mode 100644 index 0000000..f0d82ce --- /dev/null +++ b/src/pepsy/bp/_symmray.py @@ -0,0 +1,258 @@ +"""Shared native Symmray helpers for Pepsy BP corrections. + +Quimb's BP algorithms operate on ordinary dense message shapes in a few +places, while Symmray stores only the charge blocks that are currently +present. The helpers here keep the tensor network and BP messages native, +padding message support or constructing small charge-preserving operators +only where the quimb API requires a common layout. +""" + +from __future__ import annotations + +from collections.abc import Mapping + +import numpy as np + + +def is_symmray_array(value) -> bool: + """Return whether ``value`` is a Symmray block-sparse array.""" + cls = value.__class__ + return ( + getattr(cls, "__module__", "").startswith("symmray") + or (hasattr(value, "blocks") and hasattr(value, "indices")) + ) + + +def uses_symmray(tn) -> bool: + """Return whether any tensor in ``tn`` stores native Symmray data.""" + return any(is_symmray_array(tensor.data) for tensor in tn.tensors) + + +def dense_bp_tn(tn): + """Return a dense BP shadow for a Symmray scalar-network calculation. + + Quimb's D1/L1/HV1 initializers currently call dense-style ``einsum`` + expressions that Symmray cannot represent for arbitrary fermionic index + layouts. The 1-norm APIs are defined for closed scalar/factor networks, + so their compatible fallback is a topology-identical dense shadow. D2BP + never uses this helper and remains native. + """ + if not uses_symmray(tn): + return tn + dense = tn.copy() + for tensor in dense.tensors: + if is_symmray_array(tensor.data): + tensor.modify(data=to_dense(tensor.data)) + return dense + + +def dense_message_tree(messages): + """Materialize a nested BP message snapshot for a dense BP shadow.""" + if isinstance(messages, Mapping): + return { + key: dense_message_tree(value) for key, value in messages.items() + } + if isinstance(messages, tuple): + return tuple(dense_message_tree(value) for value in messages) + return to_dense(messages) if is_symmray_array(messages) else messages + + +def to_dense(value): + """Materialize one small Symmray value for a local operator calculation.""" + if hasattr(value, "to_dense"): + return np.asarray(value.to_dense()) + return np.asarray(value) + + +def dense_index_map(chargemap): + """Expand Symmray's ``{charge: size}`` map to dense-index charges.""" + result = {} + dense_index = 0 + for charge, size in chargemap.items(): + for _ in range(int(size)): + result[dense_index] = charge + dense_index += 1 + return result + + +def zero_charge(chargemap): + """Return the neutral charge matching a scalar or product symmetry.""" + charge = next(iter(chargemap), 0) + return tuple(0 for _ in charge) if isinstance(charge, tuple) else 0 + + +def align_message_pair(left, right): + """Pad two native Symmray messages onto their union of charge sectors.""" + if not ( + is_symmray_array(left) + and is_symmray_array(right) + and hasattr(left, "indices") + and hasattr(right, "indices") + ): + return left, right + + left_indices = [] + right_indices = [] + for left_index, right_index in zip(left.indices, right.indices): + left_map = dict(left_index.chargemap) + right_map = dict(right_index.chargemap) + for charge in set(left_map) & set(right_map): + if left_map[charge] != right_map[charge]: + raise ValueError("incompatible Symmray message charge dimensions") + charge_map = {**left_map, **right_map} + left_indices.append( + left_index.copy_with(chargemap=charge_map, dual=left_index.dual) + ) + right_indices.append( + right_index.copy_with(chargemap=charge_map, dual=right_index.dual) + ) + + left = left.copy_with(indices=tuple(left_indices)) + right = right.copy_with(indices=tuple(right_indices)) + left.fill_missing_blocks() + right.fill_missing_blocks() + return left, right + + +def message_distance(left, right) -> float: + """Return an L2 message distance after aligning sparse charge support.""" + left, right = align_message_pair(left, right) + return float(np.linalg.norm(to_dense(left) - to_dense(right))) + + +def align_d2bp_messages(bp) -> None: + """Pad omitted Symmray charge blocks before D2BP pair operations.""" + for index, tids in bp.tn.ind_map.items(): + if len(tids) != 2: + continue + + messages = tuple((index, tid, bp.messages[index, tid]) for tid in tids) + if not all( + is_symmray_array(message) + and hasattr(message, "indices") + and hasattr(message, "copy_with") + and hasattr(message, "fill_missing_blocks") + for _, _, message in messages + ): + continue + + tid = next(iter(tids)) + tensor = bp.tn.tensor_map[tid] + axis = tensor.inds.index(index) + bond_index = tensor.data.indices[axis] + for message_index, message_tid, message in messages: + target_indices = tuple( + bond_index.copy_with(dual=message_index.dual) + for message_index in message.indices + ) + aligned = message.copy_with(indices=target_indices) + aligned.fill_missing_blocks() + bp.messages[index, message_tid] = aligned + + +def _bond_endpoint_data(tn, index): + """Return endpoint data and their live Symmray bond indices.""" + left, right = tuple(tn.ind_map[index]) + left_tensor = tn.tensor_map[left] + right_tensor = tn.tensor_map[right] + left_axis = left_tensor.inds.index(index) + right_axis = right_tensor.inds.index(index) + return ( + left, + right, + left_tensor.data, + right_tensor.data, + left_tensor.data.indices[left_axis], + right_tensor.data.indices[right_axis], + ) + + +def rank4_operator_from_dense(tn, index, operator, *, layout="pne"): + """Build a rank-four bond operator, native when ``tn`` is Symmray. + + ``layout="pne"`` matches PNE's ``(ket-left, bra-left, ket-right, + bra-right)`` labels. ``layout="series"`` matches the loop-series local + network's ``(bra-left, ket-left, bra-right, ket-right)`` labels. + """ + left, right, left_data, right_data, left_index, right_index = ( + _bond_endpoint_data(tn, index) + ) + dimension = int(left_data.shape[left_data.indices.index(left_index)]) + dense = np.asarray(operator) + if dense.shape == (dimension * dimension, dimension * dimension): + dense = dense.reshape(dimension, dimension, dimension, dimension) + elif dense.shape != (dimension, dimension, dimension, dimension): + raise ValueError( + f"bond operator for {index!r} must have shape " + f"{(dimension * dimension, dimension * dimension)} or " + f"{(dimension, dimension, dimension, dimension)}, got {dense.shape}" + ) + + if not (is_symmray_array(left_data) and is_symmray_array(right_data)): + return dense + + array_cls = type(left_data) + kwargs = {} + if array_cls.__name__ in {"AbelianArray", "FermionicArray"}: + symmetry = getattr(left_data, "symmetry", None) + if symmetry is not None: + kwargs["symmetry"] = symmetry + if layout == "pne": + duals = ( + left_index.dual, + not left_index.dual, + right_index.dual, + not right_index.dual, + ) + elif layout == "series": + duals = ( + not left_index.dual, + left_index.dual, + not right_index.dual, + right_index.dual, + ) + else: + raise ValueError("layout must be 'pne' or 'series'") + + return array_cls.from_dense( + dense, + index_maps=( + dense_index_map(left_index.chargemap), + dense_index_map(left_index.chargemap), + dense_index_map(right_index.chargemap), + dense_index_map(right_index.chargemap), + ), + duals=duals, + charge=zero_charge(left_index.chargemap), + **kwargs, + ) + + +def rank_one_d2_projector( + tn, index, left_message, right_message, *, layout="pne" +): + """Construct a D2 rank-one projector, native when ``tn`` is Symmray.""" + left = to_dense(left_message).reshape(-1) + right = to_dense(right_message).reshape(-1) + return rank4_operator_from_dense( + tn, index, np.outer(left, right), layout=layout + ) + + +def d2_operator(tn, index, operator, *, complement=False, layout="pne"): + """Normalize a D2 projector/operator and preserve native Symmray data.""" + dense = to_dense(operator) + left, _, left_data, _, left_index, _ = _bond_endpoint_data(tn, index) + del left + dimension = int(left_data.shape[left_data.indices.index(left_index)]) + if dense.shape == (dimension, dimension, dimension, dimension): + dense = dense.reshape(dimension * dimension, dimension * dimension) + elif dense.shape != (dimension * dimension, dimension * dimension): + raise ValueError( + f"D2BP projector for {index!r} must have shape " + f"{(dimension * dimension, dimension * dimension)} or " + f"{(dimension, dimension, dimension, dimension)}, got {dense.shape}" + ) + if complement: + dense = np.eye(dimension * dimension, dtype=dense.dtype) - dense + return rank4_operator_from_dense(tn, index, dense, layout=layout) diff --git a/src/pepsy/bp/cluster.py b/src/pepsy/bp/cluster.py index 0f1f87b..02420e7 100644 --- a/src/pepsy/bp/cluster.py +++ b/src/pepsy/bp/cluster.py @@ -74,6 +74,13 @@ import autoray as ar import numpy as np +from ._symmray import ( + align_d2bp_messages as _align_symmray_d2bp_messages, + dense_bp_tn as _dense_bp_tn, + dense_message_tree as _dense_message_tree, + uses_symmray as _uses_symmray, +) + __all__ = [ "BPCandidateScore", "BPCandidateSelection", @@ -137,6 +144,14 @@ def _run_plain_bp( ) -> dict[str, Any]: """Run quimb BP and record strict rather than plateau convergence.""" info: dict[str, Any] = {} + if ( + diis is not False + and bp.__class__.__name__ == "D2BP" + and _uses_symmray(bp.tn) + ): + # Symmray messages cannot be packed by Quimb's dense DIIS + # vectorizer; sequential D2BP remains native and convergent. + diis = False bp.run( max_iterations=max_iterations, tol=tol, @@ -652,51 +667,6 @@ def _expand(bp, norm, gloops, combine, optimize, strip_exponent, progbar): return bp.contract_gloop_expand(**kwargs) -def _align_symmray_d2bp_messages(bp) -> None: - """Pad omitted Symmray charge blocks before quimb D2BP normalization. - - Quimb's D2BP normalization treats the two directed messages on an - internal bond as dense vectors with the same shape. Symmray instead - stores only non-zero charge blocks, so a rank-deficient message can have a - smaller shape than its partner even though both messages live on the same - virtual bond. ``normalize_message_pair`` then attempts a matrix product - and raises a shape-mismatch error. - - Reindex each message onto the native bond charge map and explicitly add - the missing sectors as zero blocks. The dual flags remain those of the - original message, so graded contractions and fermionic phases are left to - Symmray. Dense or non-Symmray messages are deliberately untouched. - """ - for index, tids in bp.tn.ind_map.items(): - if len(tids) != 2: - continue - - messages = tuple((index, tid, bp.messages[index, tid]) for tid in tids) - if not all( - getattr(message.__class__, "__module__", "").startswith("symmray") - and hasattr(message, "indices") - and hasattr(message, "copy_with") - and hasattr(message, "fill_missing_blocks") - for _, _, message in messages - ): - continue - - # The tensor leg is the authoritative charge map. The two messages - # may carry different subsets of these sectors after BP updates. - tid = next(iter(tids)) - tensor = bp.tn.tensor_map[tid] - axis = tensor.inds.index(index) - bond_index = tensor.data.indices[axis] - for message_index, message_tid, message in messages: - target_indices = tuple( - bond_index.copy_with(dual=message_index.dual) - for message_index in message.indices - ) - aligned = message.copy_with(indices=target_indices) - aligned.fill_missing_blocks() - bp.messages[index, message_tid] = aligned - - @dataclass class LoopClusterResult: """Result of a loop cluster expansion contraction. @@ -867,6 +837,10 @@ def linked_cluster_expand( if require_fixed_point is None: require_fixed_point = bool(run_bp) + if _uses_symmray(tn): + tn = _dense_bp_tn(tn) + messages = _dense_message_tree(messages) + from .gauges import _validate_d1_graph _validate_d1_graph(tn) @@ -1216,6 +1190,10 @@ def loop_cluster_expand( LoopClusterResult """ key, bp_cls = _cluster_bp_class(norm) + if key == "1norm" and _uses_symmray(tn): + tn = _dense_bp_tn(tn) + messages = _dense_message_tree(messages) + gauges = _dense_message_tree(gauges) contract_opts = {} if contract_opts is None else dict(contract_opts) if run_bp and ( not isinstance(max_iterations, (int, np.integer)) or max_iterations < 1 diff --git a/src/pepsy/bp/gauges.py b/src/pepsy/bp/gauges.py index 4fa29a1..4604052 100644 --- a/src/pepsy/bp/gauges.py +++ b/src/pepsy/bp/gauges.py @@ -10,6 +10,12 @@ import autoray as ar import numpy as np +from ._symmray import ( + dense_bp_tn as _dense_bp_tn, + dense_message_tree as _dense_message_tree, + uses_symmray as _uses_symmray, +) + __all__ = [ "GaugeResult", "RelayGaugeOptions", @@ -319,6 +325,9 @@ def simple_update_messages_from_gauges( ``sqrt(lambda)``. For a raw, non-gauge-inserted TN use ``message_power=1.0``. """ + if _uses_symmray(tn): + tn = _dense_bp_tn(tn) + gauges = _dense_message_tree(gauges) _validate_d1_graph(tn) gauges = {} if gauges is None else gauges @@ -366,6 +375,9 @@ def d1bp_from_simple_update_gauges( """ from quimb.tensor.belief_propagation import D1BP + if _uses_symmray(tn): + tn = _dense_bp_tn(tn) + gauges = _dense_message_tree(gauges) _validate_d1_graph(tn) work = tn.copy() gauges_copy = copy_gauges(gauges) @@ -1320,6 +1332,15 @@ def gauge_all( if norm_key not in {"1norm", "2norm"}: raise ValueError("norm must be either '1norm' or '2norm'") + if norm_key == "1norm" and _uses_symmray(tn): + # D1BP and the scalar SU bridge require dense scalar-network + # contractions. Keep the topology and values, but avoid asking the + # native Symmray fermionic contraction path to implement D1's + # one-index einsum initialization. + tn = _dense_bp_tn(tn) + su_gauges = _dense_message_tree(su_gauges) + bp_messages = _dense_message_tree(bp_messages) + su_options = {} if su_options is None else dict(su_options) bp_options = {} if bp_options is None else dict(bp_options) conversion_options = ( diff --git a/src/pepsy/bp/pne.py b/src/pepsy/bp/pne.py index a13598f..0cf4183 100644 --- a/src/pepsy/bp/pne.py +++ b/src/pepsy/bp/pne.py @@ -27,6 +27,13 @@ import numpy as np from .series import _build_bp +from ._symmray import ( + align_d2bp_messages as _align_symmray_d2bp_messages, + d2_operator as _symmray_d2_operator, + rank_one_d2_projector as _symmray_rank_one_d2_projector, + to_dense as _symmray_to_dense, + uses_symmray as _uses_symmray, +) __all__ = [ "PNEExpansionTerm", @@ -232,6 +239,8 @@ def pne_projectors(source, indices=None): if indices is None: indices = tuple(bp.tn.inner_inds()) indices = tuple(indices) + if bp.__class__.__name__ == "D2BP": + _align_symmray_d2bp_messages(bp) bp.normalize_message_pairs() return _projector_map( bp, @@ -244,7 +253,7 @@ def pne_projector_diagnostics(projectors): """Return rank, idempotency, and Hermiticity diagnostics for ``P`` maps.""" diagnostics = {} for index, operator in projectors.items(): - matrix = np.asarray(ar.to_numpy(operator)) + matrix = _symmray_to_dense(operator) if matrix.ndim > 2: half = int(np.sqrt(matrix.size)) matrix = matrix.reshape(half, half) @@ -277,6 +286,13 @@ def _rank_one_d1_projector(bp, index): def _rank_one_d2_projector(bp, index): left, right = tuple(bp.tn.ind_map[index]) + if _uses_symmray(bp.tn): + return _symmray_rank_one_d2_projector( + bp.tn, + index, + bp.messages[index, left], + bp.messages[index, right], + ) ml = bp.messages[index, left].reshape(-1) mr = bp.messages[index, right].reshape(-1) return ar.do("einsum", "i,j->ij", ml, mr) @@ -317,6 +333,14 @@ def _operator_for(bp, index, projectors, *, complement): operator = ar.do("eye", expected) - operator return operator + if _uses_symmray(bp.tn): + return _symmray_d2_operator( + bp.tn, + index, + operator, + complement=complement, + ) + left = next(iter(bp.tn.ind_map[index])) dimension = bp.tn.tensor_map[left].ind_size(index) flat_dimension = dimension * dimension @@ -495,6 +519,8 @@ def _partitioned_contract( "multi-index PNE partitions currently require form='combinatorial'" ) if normalize: + if bp.__class__.__name__ == "D2BP": + _align_symmray_d2bp_messages(bp) bp.normalize_message_pairs() if not (bp.__class__.__name__ == "D1BP" and open_inds): bp.normalize_tensors() diff --git a/src/pepsy/bp/relay.py b/src/pepsy/bp/relay.py index b9edb15..3bf0869 100644 --- a/src/pepsy/bp/relay.py +++ b/src/pepsy/bp/relay.py @@ -38,6 +38,15 @@ import autoray as ar import numpy as np +from ._symmray import ( + align_message_pair as _align_symmray_message_pair, + dense_bp_tn as _dense_bp_tn, + dense_message_tree as _dense_message_tree, + is_symmray_array as _is_symmray_array, + message_distance as _symmray_message_distance, + uses_symmray as _uses_symmray, +) + __all__ = [ "BPState", "BPUpdateResult", @@ -52,56 +61,6 @@ _RELAY_METHODS = {"l1bp", "d1bp", "d2bp"} -def _is_symmray_array(value) -> bool: - return getattr(value.__class__, "__module__", "").startswith("symmray") - - -def _align_symmray_message_pair(left, right): - """Return charge-support-aligned copies of two Symmray messages.""" - if not ( - _is_symmray_array(left) - and _is_symmray_array(right) - and hasattr(left, "indices") - and hasattr(right, "indices") - ): - return left, right - - left_indices = [] - right_indices = [] - for left_index, right_index in zip(left.indices, right.indices): - left_map = dict(left_index.chargemap) - right_map = dict(right_index.chargemap) - for charge in set(left_map) & set(right_map): - if left_map[charge] != right_map[charge]: - raise ValueError("incompatible Symmray message charge dimensions") - charge_map = {**left_map, **right_map} - left_indices.append( - left_index.copy_with(chargemap=charge_map, dual=left_index.dual) - ) - right_indices.append( - right_index.copy_with(chargemap=charge_map, dual=right_index.dual) - ) - - left = left.copy_with(indices=tuple(left_indices)) - right = right.copy_with(indices=tuple(right_indices)) - left.fill_missing_blocks() - right.fill_missing_blocks() - return left, right - - -def _symmray_message_distance(left, right) -> float: - """L2 distance after aligning omitted sparse charge blocks.""" - left, right = _align_symmray_message_pair(left, right) - if hasattr(left, "to_dense") and hasattr(right, "to_dense"): - left = left.to_dense() - right = right.to_dense() - return float(np.linalg.norm(np.asarray(left) - np.asarray(right))) - - -def _uses_symmray(tn) -> bool: - return any(_is_symmray_array(tensor.data) for tensor in tn.tensors) - - def _method_key(method: str, classes: Mapping[str, str]) -> str: """Validate and normalize a public BP method name against ``classes``.""" key = str(method).lower() @@ -609,6 +568,9 @@ def one_norm_bp( method_key = _method_key(method, _ONE_NORM_CLASSES) if not isinstance(max_iterations, (int, np.integer)) or max_iterations < 1: raise ValueError("max_iterations must be a positive integer") + if _uses_symmray(tn): + tn = _dense_bp_tn(tn) + init_messages = _dense_message_tree(init_messages) if method_key == "d1bp": from .gauges import _validate_d1_graph @@ -666,6 +628,11 @@ def two_norm_bp( bp_opts = dict(bp_opts) if _uses_symmray(tn): bp_opts.setdefault("distance", _symmray_message_distance) + # Quimb's DIIS vectorizer asks Autoray for ``symmray.concatenate``. + # Symmray intentionally has no such dense tree-vector operation; the + # native sequential D2BP iteration remains fully backend aware. + if diis is not False: + diis = False bp = _bp_class("d2bp")( tn, **_bp_constructor_kwargs("d2bp", damping, update, bp_opts), @@ -784,9 +751,15 @@ def relay_bp( _validate_d1_graph(tn) + if method_key in _ONE_NORM_CLASSES and _uses_symmray(tn): + tn = _dense_bp_tn(tn) + init_messages = _dense_message_tree(init_messages) + bp_opts = dict(bp_opts) if method_key == "d2bp" and _uses_symmray(tn): bp_opts.setdefault("distance", _symmray_message_distance) + if diis is not False: + diis = False bp_class = _bp_class(method_key) rng = np.random.default_rng(seed) bp = bp_class( @@ -841,7 +814,12 @@ def relay_bp( g = gamma[sources[key]] data = _message_data(message) if g != 0.0: - data = g * prev[key] + (1.0 - g) * data + old_data, new_data = prev[key], data + if method_key == "d2bp" and _uses_symmray(tn): + old_data, new_data = _align_symmray_message_pair( + old_data, new_data + ) + data = g * old_data + (1.0 - g) * new_data _set_message(bp, key, message, data) # Use the BP's selected distance metric on the post-memory # messages, rather than silently switching to an L-infinity diff --git a/src/pepsy/bp/series.py b/src/pepsy/bp/series.py index 6d72684..6671636 100644 --- a/src/pepsy/bp/series.py +++ b/src/pepsy/bp/series.py @@ -34,6 +34,14 @@ _filter_gauge_init_only_bp_opts, _run_plain_bp, ) +from ._symmray import ( + align_d2bp_messages as _align_symmray_d2bp_messages, + dense_bp_tn as _dense_bp_tn, + dense_message_tree as _dense_message_tree, + d2_operator as _symmray_d2_operator, + rank_one_d2_projector as _symmray_rank_one_d2_projector, + uses_symmray as _uses_symmray, +) __all__ = [ "LoopSeriesCache", @@ -471,16 +479,36 @@ def _get_d2_edge_excited(bp, term: LoopSeriesTerm): right = tids[tid_right] ml = bp.messages[index, tid_left] mr = bp.messages[index, tid_right] - p0 = ar.do("einsum", "i,j->ij", ml.reshape(-1), mr.reshape(-1)) + if _uses_symmray(bp.tn): + p0 = _symmray_rank_one_d2_projector( + bp.tn, index, ml, mr, layout="series" + ) + else: + p0 = ar.do( + "einsum", + "i,j->ij", + ml.reshape(-1), + mr.reshape(-1), + ) if index in excited_edges: - projector = ar.do("eye", ar.do("shape", p0)[0]) - p0 + if _uses_symmray(bp.tn): + projector = _symmray_d2_operator( + bp.tn, + index, + p0, + complement=True, + layout="series", + ) + else: + projector = ar.do("eye", ar.do("shape", p0)[0]) - p0 else: projector = p0 - projector = ar.do( - "reshape", - projector, - ar.do("shape", ml) + ar.do("shape", mr), - ) + if not _uses_symmray(bp.tn): + projector = ar.do( + "reshape", + projector, + ar.do("shape", ml) + ar.do("shape", mr), + ) inds = (*left, *right) local |= qtn.Tensor(projector, inds) @@ -543,6 +571,8 @@ def _contract_loop_series( normalize, ): if normalize: + if bp.__class__.__name__ == "D2BP": + _align_symmray_d2bp_messages(bp) bp.normalize_message_pairs() bp.normalize_tensors() @@ -585,6 +615,10 @@ def _build_bp( validate_graph=True, ): key, bp_cls = _cluster_bp_class(norm) + if key == "1norm" and _uses_symmray(tn): + tn = _dense_bp_tn(tn) + messages = _dense_message_tree(messages) + gauges = _dense_message_tree(gauges) if bp_runner not in {"plain", "relay"}: raise ValueError("bp_runner must be either 'plain' or 'relay'") if gauges is not None and messages is not None: diff --git a/src/pepsy/bp/weights.py b/src/pepsy/bp/weights.py index 32459fb..fc1604d 100644 --- a/src/pepsy/bp/weights.py +++ b/src/pepsy/bp/weights.py @@ -19,6 +19,10 @@ import numpy as np +from ._symmray import dense_bp_tn as _dense_bp_tn +from ._symmray import to_dense as _symmray_to_dense +from ._symmray import uses_symmray as _uses_symmray + __all__ = ["WeightPassingResult", "weight_pass"] @@ -243,6 +247,30 @@ def weight_pass( if tol < 0 or eps <= 0: raise ValueError("tol must be non-negative and eps must be positive") + if _uses_symmray(tn): + # Weight passing is an SVD-based scalar-network utility. Quimb's + # implementation uses dense reshapes and cannot perform those local + # factorizations on native Symmray blocks. Keep the public API + # compatible by using a topology-identical dense shadow; D2BP itself + # remains native and should be preferred for physical Symmray norms. + dense_weights = ( + None + if weights is None + else { + index: _symmray_to_dense(value) + for index, value in weights.items() + } + ) + return weight_pass( + _dense_bp_tn(tn), + alpha=alpha, + max_iterations=max_iterations, + tol=tol, + weights=dense_weights, + index_order=index_order, + eps=eps, + ) + _validate_network(tn) network = tn.copy() bond_order = ( diff --git a/tests/test_bp_symmray.py b/tests/test_bp_symmray.py index 2900873..c616dfd 100644 --- a/tests/test_bp_symmray.py +++ b/tests/test_bp_symmray.py @@ -5,9 +5,19 @@ import numpy as np import pytest -pytest.importorskip("symmray") +sr = pytest.importorskip("symmray") +qtn = pytest.importorskip("quimb.tensor") -from pepsy.bp import gauge_all, loop_cluster_expand, two_norm_bp # noqa: E402 +from pepsy.bp import ( # noqa: E402 + gauge_all, + loop_cluster_expand, + loop_series_expand, + one_norm_bp, + partitioned_expand, + relay_bp, + two_norm_bp, + weight_pass, +) from pepsy.tensors import ( # noqa: E402 SymPEPS, site_charge_alternating, @@ -186,3 +196,127 @@ def test_fermionic_sympeps_su_d2bp_bridge_preserves_symmetry_blocks( np.testing.assert_allclose( reverse_reconstructed.norm(), exact_norm, rtol=1e-10 ) + + +def test_fermionic_u1_d2bp_corrections_and_relay_use_native_messages(): + """The D2 relay, series, and PNE paths share native Symmray handling.""" + state = SymPEPS.random( + 2, + 2, + symmetry="U1", + bond_dim=2, + phys_dim=2, + fermionic=True, + seed=1300, + dtype="complex128", + ) + bond = next(iter(state.tn.inner_inds())) + + relay = relay_bp( + state.tn, + method="d2bp", + num_relays=2, + max_iterations=200, + tol=1e-10, + ) + assert relay.converged + assert all( + type(message).__name__ == "U1FermionicArray" + for message in relay.messages.values() + ) + + series = loop_series_expand( + state.tn, + 2, + norm="2norm", + max_iterations=200, + tol=1e-10, + ) + assert series.bp_converged + assert np.isfinite(float(np.real(series.estimate))) + + pne = partitioned_expand( + state.tn, + partition_inds=(bond,), + norm="2norm", + max_iterations=200, + tol=1e-10, + ) + assert pne.bp_converged + assert np.isfinite(float(np.real(pne.estimate))) + assert all( + type(message).__name__ == "U1FermionicArray" + for message in pne.messages.values() + ) + + +def _closed_u1_scalar_network(): + """Build a small native Symmray scalar network for the D1 compatibility path.""" + maps = ({0: 0, 1: 1},) * 2 + left = sr.U1Array.from_dense( + np.diag([2.0, 3.0]), + index_maps=maps, + duals=(False, True), + charge=0, + ) + right = sr.U1Array.from_dense( + np.diag([5.0, 7.0]), + index_maps=maps, + duals=(True, False), + charge=0, + ) + return qtn.TensorNetwork( + [ + qtn.Tensor(left, inds=("x", "y")), + qtn.Tensor(right, inds=("x", "y")), + ] + ) + + +def test_native_symmray_closed_scalar_network_works_through_one_norm_apis(): + """Valid closed 1-norm inputs use the dense-compatible BP shadow.""" + tn = _closed_u1_scalar_network() + + result = one_norm_bp( + tn, + method="d1bp", + max_iterations=100, + tol=1e-10, + ) + assert result.converged + assert all(type(tensor.data).__name__ == "ndarray" for tensor in result.bp.tn) + + cluster = loop_cluster_expand( + tn, + 0, + norm="1norm", + max_iterations=100, + tol=1e-10, + ) + assert cluster.bp_converged + assert np.isfinite(float(np.real(cluster.estimate))) + + series = loop_series_expand( + tn, + 0, + norm="1norm", + max_iterations=100, + tol=1e-10, + ) + assert series.bp_converged + assert np.isfinite(float(np.real(series.estimate))) + + pne = partitioned_expand( + tn, + partition_inds=("x",), + norm="1norm", + max_iterations=100, + tol=1e-10, + ) + assert pne.bp_converged + assert np.isfinite(float(np.real(pne.estimate))) + + weights = weight_pass(tn, max_iterations=2) + assert all( + type(tensor.data).__name__ == "ndarray" for tensor in weights.network + ) From a7cb9b11f70f615f9d24a65db41130159aedff79 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Wed, 29 Jul 2026 22:52:44 -0600 Subject: [PATCH 26/70] Fix fermionic D2BP contraction phase metadata --- .github/skills/belief-propagation/SKILL.md | 10 ++++ src/pepsy/bp/_symmray.py | 47 ++++++++++++++++ src/pepsy/bp/cluster.py | 3 + src/pepsy/bp/gauges.py | 3 + src/pepsy/bp/relay.py | 6 ++ src/pepsy/bp/series.py | 3 + tests/test_bp_symmray.py | 64 ++++++++++++++++++++++ 7 files changed, 136 insertions(+) diff --git a/.github/skills/belief-propagation/SKILL.md b/.github/skills/belief-propagation/SKILL.md index 54ce30d..e7793d3 100644 --- a/.github/skills/belief-propagation/SKILL.md +++ b/.github/skills/belief-propagation/SKILL.md @@ -102,6 +102,16 @@ construction, while keeping the messages and tensor data block-sparse. Do not replace this D2 path with a dense copy: dense eigendecomposition can reorder eigenvectors across charge sectors and produce invalid Symmray gates. +Native fermionic arrays also require their global-phase metadata. An odd +Symmray array with a site `label` must carry the corresponding implicit +`dummy_modes`; `FermionicArray.new_with(...)` intentionally drops those modes. +The D2 entry points repair labelled odd arrays on a private network copy before +constructing BP or finite clusters, preserving blocks and lazy phases. This is +what makes the result invariant under Cotengra contraction-tree choices. Do +not treat `optimize="auto-hq"` as a fermionic sign fix: it is a Cotengra path +preset, just like a reusable `PathOptimizer`, and a path-dependent sign means +the input metadata is invalid or incomplete. + For valid closed scalar Symmray networks, the 1-norm APIs (`L1BP`, `HV1BP`, and `D1BP`), their loop/series/PNE corrections, D1 SU bridge, and `weight_pass` use a topology-identical dense BP shadow because Quimb's D1 diff --git a/src/pepsy/bp/_symmray.py b/src/pepsy/bp/_symmray.py index f0d82ce..9d9d54f 100644 --- a/src/pepsy/bp/_symmray.py +++ b/src/pepsy/bp/_symmray.py @@ -28,6 +28,53 @@ def uses_symmray(tn) -> bool: return any(is_symmray_array(tensor.data) for tensor in tn.tensors) +def restore_fermionic_dummy_modes(tn): + """Restore implicit dummy modes on labelled odd Symmray tensors. + + Symmray's ``FermionicArray.new_with`` deliberately drops ``dummy_modes``. + That is a useful low-level default, but it is unsafe when a tensor is + subsequently used in a fermionic contraction: an odd array with a site + label and no dummy mode is treated as phase-neutral, so the result can + depend on the contraction tree. Quimb simple-update factorizations can + create exactly this representation. + + The site label is the canonical mode identity already carried by native + fermionic PEPS tensors. Reconstructing with ``dummy_modes=None`` asks + Symmray to recreate that mode while preserving the blocks, lazy phases, + indices, and charge. The input network is copied before any repair. + Odd arrays without a label are left unchanged because their canonical + fermionic mode cannot be inferred safely. + """ + if not uses_symmray(tn): + return tn + + target = tn.copy() + for tensor in target.tensors: + data = tensor.data + if not ( + is_symmray_array(data) + and hasattr(data, "dummy_modes") + and getattr(data, "parity", 0) + and not data.dummy_modes + and getattr(data, "label", None) is not None + ): + continue + + tensor.modify( + data=type(data).from_blocks( + data.blocks, + duals=data.indices, + charge=getattr(data, "charge", None), + symmetry=getattr(data, "symmetry", None), + phases=getattr(data, "phases", None), + label=data.label, + dummy_modes=None, + ) + ) + + return target + + def dense_bp_tn(tn): """Return a dense BP shadow for a Symmray scalar-network calculation. diff --git a/src/pepsy/bp/cluster.py b/src/pepsy/bp/cluster.py index 02420e7..a36b3d9 100644 --- a/src/pepsy/bp/cluster.py +++ b/src/pepsy/bp/cluster.py @@ -78,6 +78,7 @@ align_d2bp_messages as _align_symmray_d2bp_messages, dense_bp_tn as _dense_bp_tn, dense_message_tree as _dense_message_tree, + restore_fermionic_dummy_modes as _restore_fermionic_dummy_modes, uses_symmray as _uses_symmray, ) @@ -1190,6 +1191,8 @@ def loop_cluster_expand( LoopClusterResult """ key, bp_cls = _cluster_bp_class(norm) + if key == "2norm" and _uses_symmray(tn): + tn = _restore_fermionic_dummy_modes(tn) if key == "1norm" and _uses_symmray(tn): tn = _dense_bp_tn(tn) messages = _dense_message_tree(messages) diff --git a/src/pepsy/bp/gauges.py b/src/pepsy/bp/gauges.py index 4604052..d26fd51 100644 --- a/src/pepsy/bp/gauges.py +++ b/src/pepsy/bp/gauges.py @@ -13,6 +13,7 @@ from ._symmray import ( dense_bp_tn as _dense_bp_tn, dense_message_tree as _dense_message_tree, + restore_fermionic_dummy_modes as _restore_fermionic_dummy_modes, uses_symmray as _uses_symmray, ) @@ -548,6 +549,8 @@ def d2bp_from_simple_update_gauges( """ from quimb.tensor.belief_propagation import D2BP + if _uses_symmray(tn): + tn = _restore_fermionic_dummy_modes(tn) _validate_d2_graph(tn) work = tn.copy() gauges_copy = copy_gauges(gauges) diff --git a/src/pepsy/bp/relay.py b/src/pepsy/bp/relay.py index 3bf0869..e642a45 100644 --- a/src/pepsy/bp/relay.py +++ b/src/pepsy/bp/relay.py @@ -44,6 +44,7 @@ dense_message_tree as _dense_message_tree, is_symmray_array as _is_symmray_array, message_distance as _symmray_message_distance, + restore_fermionic_dummy_modes as _restore_fermionic_dummy_modes, uses_symmray as _uses_symmray, ) @@ -625,6 +626,8 @@ def two_norm_bp( """ if not isinstance(max_iterations, (int, np.integer)) or max_iterations < 1: raise ValueError("max_iterations must be a positive integer") + if _uses_symmray(tn): + tn = _restore_fermionic_dummy_modes(tn) bp_opts = dict(bp_opts) if _uses_symmray(tn): bp_opts.setdefault("distance", _symmray_message_distance) @@ -751,6 +754,9 @@ def relay_bp( _validate_d1_graph(tn) + if method_key == "d2bp" and _uses_symmray(tn): + tn = _restore_fermionic_dummy_modes(tn) + if method_key in _ONE_NORM_CLASSES and _uses_symmray(tn): tn = _dense_bp_tn(tn) init_messages = _dense_message_tree(init_messages) diff --git a/src/pepsy/bp/series.py b/src/pepsy/bp/series.py index 6671636..1dc8123 100644 --- a/src/pepsy/bp/series.py +++ b/src/pepsy/bp/series.py @@ -40,6 +40,7 @@ dense_message_tree as _dense_message_tree, d2_operator as _symmray_d2_operator, rank_one_d2_projector as _symmray_rank_one_d2_projector, + restore_fermionic_dummy_modes as _restore_fermionic_dummy_modes, uses_symmray as _uses_symmray, ) @@ -615,6 +616,8 @@ def _build_bp( validate_graph=True, ): key, bp_cls = _cluster_bp_class(norm) + if key == "2norm" and _uses_symmray(tn): + tn = _restore_fermionic_dummy_modes(tn) if key == "1norm" and _uses_symmray(tn): tn = _dense_bp_tn(tn) messages = _dense_message_tree(messages) diff --git a/tests/test_bp_symmray.py b/tests/test_bp_symmray.py index c616dfd..a421a66 100644 --- a/tests/test_bp_symmray.py +++ b/tests/test_bp_symmray.py @@ -63,6 +63,70 @@ def test_fermionic_u1_two_norm_bp_is_exact_on_a_tree(bond_dim): ) +def test_d2bp_repairs_labelled_odd_arrays_after_new_with(): + """D2BP repairs the phase metadata dropped by Symmray ``new_with``.""" + state = SymPEPS.random( + 2, + 2, + symmetry="U1", + bond_dim=2, + phys_dim=2, + fermionic=True, + seed=750, + dtype="complex128", + ) + broken = state.tn.copy() + for tensor in broken.tensors: + data = tensor.data + if not getattr(data, "parity", 0): + continue + tensor.modify( + data=type(data).from_blocks( + data.blocks, + duals=data.indices, + charge=data.charge, + symmetry=data.symmetry, + phases=data.phases, + label=data.label, + dummy_modes=(), + ) + ) + + assert all( + not tensor.data.dummy_modes + for tensor in broken.tensors + if getattr(tensor.data, "parity", 0) + ) + + result = two_norm_bp( + broken, + max_iterations=100, + tol=1e-10, + diis=False, + ) + assert result.converged + assert all( + tensor.data.dummy_modes + for tensor in result.bp.tn.tensors + if getattr(tensor.data, "parity", 0) + ) + + cluster = loop_cluster_expand( + broken, + gloops=0, + norm="2norm", + max_iterations=100, + tol=1e-10, + diis=False, + ) + assert cluster.bp_converged + assert all( + tensor.data.dummy_modes + for tensor in cluster.bp.tn.tensors + if getattr(tensor.data, "parity", 0) + ) + + @pytest.mark.parametrize("bond_dim", (2, 3)) def test_fermionic_u1_loop_cluster_runs_on_3x4_peps(bond_dim): """D2BP loop clusters pad sparse charge blocks without densifying.""" From 1084add4a532d69aa2770f08f1c740d997a09dc0 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Thu, 30 Jul 2026 07:55:09 -0700 Subject: [PATCH 27/70] backend diagnostics for tensor optimizers --- docs/api/optimizers/mps.md | 11 + docs/api/optimizers/stabilizer_tn.md | 10 + docs/api/optimizers/tree.md | 8 + docs/api/optimizers/tree_stabilizer.md | 10 + docs/api/package.md | 6 + src/pepsy/__init__.py | 2 + src/pepsy/backends/__init__.py | 4 + src/pepsy/backends/convert.py | 152 +++++++++++++ src/pepsy/optimizers/mps/optimizer.py | 207 +++++++++--------- .../stabilizer_tn/mps_stab_optimizer.py | 157 +++++++++++-- src/pepsy/optimizers/tree/optimizer.py | 122 +++++------ .../optimizers/tree_stabilizer/optimizer.py | 33 ++- tests/test_backends.py | 16 ++ tests/test_optimize_mps.py | 104 +++++++++ tests/test_optimize_tree.py | 50 +++++ tests/test_optimize_tree_stabilizer.py | 5 + tests/test_public_api.py | 4 +- tests/test_stabilizer_tn.py | 43 ++++ 18 files changed, 744 insertions(+), 200 deletions(-) diff --git a/docs/api/optimizers/mps.md b/docs/api/optimizers/mps.md index 8efda5e..a9269a5 100644 --- a/docs/api/optimizers/mps.md +++ b/docs/api/optimizers/mps.md @@ -15,6 +15,17 @@ helper `MpsOptimizer.submpo_event(mpo, where)` builds the tuple form. These events are applied with `gate_with_submpo_` and compressed to `chi`; they are only accepted in `mode="mpo"`. +`MpsOptimizer.backend_info()` reports the backend, dtype, and device inferred +from every live MPS tensor; the same values are also available as the +state-derived `backend`, `backend_dtype`, and `backend_device` attributes. +Every gate and every tensor in a sub-MPO is checked against that signature +before replay. Explicit mismatches are converted on an execution copy with a +`UserWarning`; the queued payloads remain unchanged. Native Symmray MPS data +reports `backend="symmray"` and includes `array_backend` for the underlying +NumPy, Torch, or CuPy charge-sector blocks. Dense payloads cannot be promoted +to native Symmray gates because that would lose charge and fermionic metadata; +construct those gates with the matching Symmray convention instead. + Streams may also include control events. `("measure", pauli, where[, outcome])` collapses onto a Pauli eigenvalue and records `(pauli, where, outcome, prob)`. `("reset", where[, basis])` resets each target to the `+1` eigenstate of diff --git a/docs/api/optimizers/stabilizer_tn.md b/docs/api/optimizers/stabilizer_tn.md index 7ba34ab..34e9687 100644 --- a/docs/api/optimizers/stabilizer_tn.md +++ b/docs/api/optimizers/stabilizer_tn.md @@ -312,5 +312,15 @@ small-`n` validation only. The focused STN tests exercise NumPy, Torch, JAX, and CuPy paths; optional JAX/CuPy tests skip only when the dependency or CUDA runtime is unavailable. +When an existing coefficient MPS is supplied, the stabilizer optimizer infers +its common `backend`, `dtype`, and `device` automatically, even when +`to_backend` is omitted. `backend_info()` returns the live mapping and refreshes +the public `backend`, `backend_dtype`, and `backend_device` attributes. Explicit +matrix gates and coefficient-frame sub-MPOs are checked against that signature; +foreign payloads emit one warning per source/target combination. Sub-MPOs are +copied to the live backend without mutating the caller's MPO; dense physical +matrices use a temporary NumPy view for Stim/Pauli classification and their +coefficient contractions remain on the inferred backend. + > API details are maintained as handwritten Markdown in this page. diff --git a/docs/api/optimizers/tree.md b/docs/api/optimizers/tree.md index 9debf35..baaad03 100644 --- a/docs/api/optimizers/tree.md +++ b/docs/api/optimizers/tree.md @@ -620,6 +620,14 @@ Mixed-backend initial states fail immediately because there is no unambiguous safe execution backend. Internal Pauli/projector tensors follow the state backend automatically. +The same diagnostic is reflected by the state-derived `backend`, +`backend_dtype`, `backend_device`, and `array_backend` attributes. The +complete gate stream is checked before replay, including every tensor in each +sub-MPO; mismatches are converted on execution copies and warned about without +mutating queued inputs. Native Symmray states report `backend="symmray"` plus +the underlying NumPy, Torch, or CuPy `array_backend`, preserving U1/U1U1 charge +and fermionic metadata. + ```python import pepsy as py import torch diff --git a/docs/api/optimizers/tree_stabilizer.md b/docs/api/optimizers/tree_stabilizer.md index 5eccca4..a679353 100644 --- a/docs/api/optimizers/tree_stabilizer.md +++ b/docs/api/optimizers/tree_stabilizer.md @@ -157,5 +157,15 @@ measurement offsets and computational bits (`+1 -> 0`, `-1 -> 1`). The action must be one gate entry. Stim `CX/CY/CZ rec[k] q` instructions are lowered to this same form by `compile_stim_circuit`. +TreeStab derives `backend`, `dtype`, and `device` from every live coefficient +TTN tensor, including a caller-supplied Torch, JAX, or CuPy tree when +`to_backend` is omitted. `backend_info()` refreshes the same public +`backend`, `backend_dtype`, `backend_device`, and `array_backend` attributes. +Explicit matrix gates and sub-MPO payloads are checked at the stream boundary; +foreign arrays warn once and sub-MPOs are copied before conversion, preserving +the caller's operator and the TTN's canonical/isometry metadata. Stim gate +classification remains a NumPy-side operation, while TreeOptimizer applies +the coefficient update on the inferred backend. + > API details are maintained as handwritten Markdown in this page. diff --git a/docs/api/package.md b/docs/api/package.md index 91e3aca..5887463 100644 --- a/docs/api/package.md +++ b/docs/api/package.md @@ -33,6 +33,11 @@ from pepsy.sampling import MpsSampler from pepsy.tensors import OneDMap, ps_to_mps, ps_to_peps, tn_norm ``` +For a shared backend contract across tensor-network classes, use +`pepsy.backend_infer(value)`. It accepts an array or an MPS/TTN and returns +`backend`, `dtype`, and `device`; Symmray inputs also report the underlying +`array_backend` used by their charge-sector blocks. + ## Advanced namespaces | Area | Canonical import | Notes | @@ -63,6 +68,7 @@ pepsy.BdyMPS pepsy.rx pepsy.SweepOptimizer pepsy.ps_to_mps +pepsy.backend_infer ``` For new code, prefer the canonical namespace imports above. They make the diff --git a/src/pepsy/__init__.py b/src/pepsy/__init__.py index e4d8a21..bf5ebac 100644 --- a/src/pepsy/__init__.py +++ b/src/pepsy/__init__.py @@ -41,6 +41,7 @@ "peps_norm": ".boundary", "peps_normalize": ".boundary", "backend_cupy": ".backends", + "backend_infer": ".backends", "backend_jax": ".backends", "backend_numpy": ".backends", "backend_torch": ".backends", @@ -281,6 +282,7 @@ def __getattr__(name): from .bp import gauge_all, gauge_all_simple, one_norm_bp # noqa: F401 from . import backends, boundary, bp, experimental, fitting, operators, optimizers, sampling, solvers, tensors, vmc # noqa: F401 from .backends import ( # noqa: F401 + backend_infer, get_default_array_backend, get_default_grad_backend, build_backend, diff --git a/src/pepsy/backends/__init__.py b/src/pepsy/backends/__init__.py index 2f2d208..1d09333 100644 --- a/src/pepsy/backends/__init__.py +++ b/src/pepsy/backends/__init__.py @@ -3,8 +3,10 @@ from importlib import import_module from .convert import ( + backend_infer, dispatch_backend_converter, infer_backend_and_dtype, + infer_backend_signature, infer_backend_converter_from_sample, resolve_backend_sample_data, resolve_backend_sample_data_from_tn, @@ -25,6 +27,7 @@ ) __all__ = [ + "backend_infer", "build_backend", "backend_cupy", "backend_jax", @@ -32,6 +35,7 @@ "backend_torch", "dispatch_backend_converter", "infer_backend_and_dtype", + "infer_backend_signature", "infer_backend_converter_from_sample", "resolve_backend_sample_data", "resolve_backend_sample_data_from_tn", diff --git a/src/pepsy/backends/convert.py b/src/pepsy/backends/convert.py index d89b8bc..af2bd31 100644 --- a/src/pepsy/backends/convert.py +++ b/src/pepsy/backends/convert.py @@ -115,6 +115,135 @@ def resolve_backend_sample_data_from_tn(tn): return None +def _is_symmray_array(value): + """Return whether ``value`` is a Symmray block-sparse array.""" + return ( + type(value).__module__.split(".", 1)[0] == "symmray" + or hasattr(value, "blocks") and hasattr(value, "indices") + ) + + +def _symmray_block_signatures(value): + """Return backend signatures for the raw arrays held by a Symmray value.""" + blocks = getattr(value, "blocks", None) + if not blocks: + return () + signatures = [] + for block in blocks.values(): + backend, dtype = infer_backend_and_dtype(block) + device = getattr(block, "device", None) + signatures.append( + (backend, str(dtype), None if device is None else str(device)) + ) + return tuple(signatures) + + +def infer_backend_signature(sample_data): + """Infer comparable backend, dtype, device, and Symmray-block metadata. + + Dense arrays return the traditional ``(backend, dtype, device)`` tuple. + Symmray arrays return ``(symmray, dtype, device, block_backend)`` where + ``block_backend`` is the backend of their raw charge-sector blocks. The + extra field is essential: ``ar.infer_backend`` intentionally reports the + structured Symmray container rather than the Torch/CuPy backend used by + its blocks. + """ + if sample_data is None: + raise ValueError("Cannot infer backend: sample_data is None.") + + try: + backend, dtype = infer_backend_and_dtype(sample_data) + except (AttributeError, KeyError, TypeError, ValueError): + # Untyped Python sequences are convenience inputs. Treat them as + # non-backend data so callers can materialize them on the state + # backend without emitting a transfer warning. + try: + array = np.asarray(sample_data) + dtype = ar.get_dtype_name(array) + except (TypeError, ValueError, KeyError) as exc: + raise TypeError( + "Could not infer a backend or dtype from the supplied array." + ) from exc + return "builtins", str(dtype), None + device = getattr(sample_data, "device", None) + device = None if device is None else str(device) + if backend != "symmray" and not _is_symmray_array(sample_data): + return backend, str(dtype), device + + block_signatures = _symmray_block_signatures(sample_data) + if block_signatures: + unique = set(block_signatures) + if len(unique) != 1: + raise TypeError( + "Symmray arrays must use one underlying backend, dtype, and " + f"device; found {sorted(unique)!r}." + ) + block_backend, block_dtype, block_device = block_signatures[0] + # The raw block dtype/device is authoritative for a structured array. + dtype = block_dtype + device = block_device + else: + block_backend = str(getattr(sample_data, "backend", "symmray")) + return "symmray", str(dtype), device, str(block_backend) + + +def _backend_data_values(value): + """Return array payloads from an array, tensor, or tensor network.""" + tensor_map = getattr(value, "tensor_map", None) + if tensor_map is not None: + values = tuple( + getattr(tensor, "data", None) for tensor in tensor_map.values() + ) + return tuple(data for data in values if data is not None) + + tensors = getattr(value, "tensors", None) + if tensors is not None and not hasattr(value, "shape"): + values = tuple(getattr(tensor, "data", None) for tensor in tensors) + return tuple(data for data in values if data is not None) + + data = getattr(value, "data", None) + if data is not None and hasattr(data, "shape") and hasattr(data, "dtype"): + return (data,) + return (value,) + + +def backend_infer(value): + """Infer and validate backend metadata from an array or tensor network. + + Parameters + ---------- + value + An array-like payload, Quimb tensor, or tensor network such as an MPS + or :class:`TreeTensorNetwork`. For a tensor network, every tensor is + checked for one common backend, dtype, and device. + + Returns + ------- + dict + The normalized metadata mapping ``backend``, ``dtype``, and + ``device``. Native Symmray arrays additionally include + ``array_backend`` for their underlying NumPy, Torch, or CuPy blocks. + """ + values = _backend_data_values(value) + if not values: + raise ValueError("Cannot infer backend: value contains no tensors.") + + signatures = tuple(infer_backend_signature(data) for data in values) + signature = signatures[0] + mismatched = tuple(candidate for candidate in signatures[1:] if candidate != signature) + if mismatched: + raise TypeError( + "Backend arrays must use one compatible backend, dtype, and " + f"device; found {signature!r} and {mismatched[0]!r}." + ) + + backend, dtype, device = signature[:3] + info = {"backend": backend, "dtype": dtype, "device": device} + if len(signature) > 3: + info["array_backend"] = signature[3] + return info + + def infer_backend_and_dtype(sample_data): """Infer backend name and dtype name from sample tensor data.""" if sample_data is None: @@ -334,6 +463,29 @@ def infer_backend_converter_from_sample( return None backend, dtype_name = infer_backend_and_dtype(sample_data) + if backend == "symmray" or _is_symmray_array(sample_data): + blocks = getattr(sample_data, "blocks", None) or {} + block_sample = next(iter(blocks.values()), None) + if block_sample is None: + return None + block_converter = infer_backend_converter_from_sample( + block_sample, + cast_complex_to_real=cast_complex_to_real, + ) + if block_converter is None: + return None + + def _to_symmray_or_block(value): + # Network-level ``apply_to_arrays`` callbacks may receive raw + # sector blocks, while public payload conversion receives the + # structured Symmray object. Support both call sites. + if _is_symmray_array(value): + converted = value.copy() + converted.apply_to_arrays(block_converter) + return converted + return block_converter(value) + + return _to_symmray_or_block try: return dispatch_backend_converter( backend=backend, diff --git a/src/pepsy/optimizers/mps/optimizer.py b/src/pepsy/optimizers/mps/optimizer.py index 72de4b1..64aa4d1 100644 --- a/src/pepsy/optimizers/mps/optimizer.py +++ b/src/pepsy/optimizers/mps/optimizer.py @@ -62,6 +62,11 @@ import numpy as np import quimb.tensor as qtn +from ...backends import ( + backend_infer, + infer_backend_converter_from_sample, + infer_backend_signature, +) from ...fitting.local import FIT from ...operators.gates import ( _normalize_gate_entries, @@ -87,6 +92,11 @@ _NORM_INCLUDES_EXPONENT_CACHE = {} +def _array_backend_signature(array): + """Return comparable backend / dtype / device metadata for an array.""" + return infer_backend_signature(array) + + def _normalize_event_name(name): """Normalize a stream event name for matching.""" return str(name).replace("-", "_").strip().lower() @@ -1090,7 +1100,12 @@ def __init__( # pylint: disable=too-many-arguments,too-many-positional-argument self._unitary_initial_norm = None self._unitary_previous_norm = None self._unitary_global_norm_tracking = False - self._backend_mismatch_warned = False + self._backend_conversion_warnings = set() + self.backend = None + self.backend_dtype = None + self.backend_device = None + self.array_backend = None + self.backend_info() self._init_canonicalization() def _info_for_state(self, p, info=None): @@ -1317,7 +1332,11 @@ def _prepare_mix_dmrg_state(self): def set_p(self, p): """Assign a new state and reset canonicalization metadata.""" - self.p = self._install_represented_norm(p if self.inplace else p.copy()) + new_p = self._install_represented_norm(p if self.inplace else p.copy()) + # Validate before replacing the live state so a mixed-backend input + # cannot leave this optimizer half-updated after a failed assignment. + self._state_backend_info_for(new_p) + self.p = new_p self.qubits = list(range(int(getattr(self.p, "L", 0)))) self.logical_order = list(self.qubits) self._persistent_layout_plan = None @@ -1328,7 +1347,8 @@ def set_p(self, p): self._su_gauges_state = None self._su_force_regauge = self.mode == "su" self.p_ungauged = None - self._backend_mismatch_warned = False + self._backend_conversion_warnings = set() + self.backend_info() self._init_canonicalization() def normalize(self, eps=1e-15, insert=None): @@ -1397,7 +1417,9 @@ def copy(self) -> "MpsOptimizer": copied._unitary_initial_norm = self._unitary_initial_norm copied._unitary_previous_norm = self._unitary_previous_norm copied._unitary_global_norm_tracking = self._unitary_global_norm_tracking - copied._backend_mismatch_warned = self._backend_mismatch_warned + copied._backend_conversion_warnings = set( + self._backend_conversion_warnings + ) copied._su_gauges_supplied = True copied._su_gauges_ready = self._su_gauges_ready copied._su_gauges_state = copied.p if self._su_gauges_ready else None @@ -2726,38 +2748,59 @@ def _state_backend_like(self): return tensor.data return None + @staticmethod + def _state_backend_info_for(state): + """Validate and describe the common backend of an MPS-like state.""" + return backend_infer(state) + + def backend_info(self): + """Return the state-derived backend, dtype, and device diagnostics.""" + info = self._state_backend_info_for(self.p) + self.backend = info["backend"] + self.backend_dtype = info["dtype"] + self.backend_device = info["device"] + self.array_backend = info.get("array_backend", info["backend"]) + return info + + def _warn_backend_conversion(self, source_signature, target_signature, *, kind): + """Warn once for one explicit stream source/target conversion.""" + warning_key = (kind, source_signature, target_signature) + if ( + source_signature[0] != "builtins" + and warning_key not in self._backend_conversion_warnings + ): + self._backend_conversion_warnings.add(warning_key) + warnings.warn( + f"MpsOptimizer converted a {kind} payload from " + f"backend/dtype/device {source_signature!r} to the live MPS " + f"backend/dtype/device {target_signature!r}; provide matching " + f"{kind} payloads to avoid this conversion.", + UserWarning, + stacklevel=3, + ) + def _to_state_backend(self, array): - """Return ``array`` cast to ``self.p``'s backend and complex dtype.""" + """Return ``array`` cast to the backend and dtype owned by ``self.p``.""" like = self._state_backend_like() if like is None: return np.asarray(array, dtype=complex) - # Avoid any Autoray conversion for an already-compatible payload. This - # is important for Symmray, whose backend intentionally does not expose - # a generic ``array`` constructor, and keeps the common matching-gate - # path allocation-free for every backend. - try: - if ( - ar.infer_backend(array) == ar.infer_backend(like) - and getattr(array, "dtype", None) == getattr(like, "dtype", None) - ): - return array - except (AttributeError, TypeError, ValueError): - pass - dtype = getattr(like, "dtype", complex) - if "complex" not in str(dtype): - dtype = getattr( - ar.do("array", np.asarray(1.0j), like=like), "dtype", complex - ) - if ( - getattr(array, "device", None) == getattr(like, "device", None) - and getattr(array, "dtype", None) == dtype - ): + target_signature = _array_backend_signature(like) + source_signature = _array_backend_signature(array) + if source_signature == target_signature: return array - # ``np.asarray`` would try to materialize a CUDA/Torch gate on the CPU. - # Let autoray move/cast any foreign array directly to the MPS backend. - # The final astype also preserves the complex dtype of the live state. - arr = ar.do("array", array, like=like) - return ar.do("astype", arr, dtype) + if target_signature[0] == "symmray" and source_signature[0] != "symmray": + raise TypeError( + "Cannot convert a dense gate/operator payload into a native " + "Symmray MPS without charge and fermionic metadata. Build the " + "payload as a Symmray array on the target U1/U1U1 backend." + ) + converter = infer_backend_converter_from_sample(like) + if converter is not None: + return converter(array) + if target_signature[0] == "numpy": + return ar.to_numpy(array) + # Keep the old Autoray fallback for optional/custom dense backends. + return ar.do("array", array, like=like) def to_backend(self, array): """Return ``array`` on the backend currently owned by ``self.p``. @@ -2772,94 +2815,42 @@ def _prepare_gate_stream_backend(self, gates, event_types): """Prepare gate and sub-MPO payloads for the live MPS backend lazily. Gate streams are commonly authored as NumPy arrays even when the live - MPS uses Torch, JAX, CuPy, or another Autoray backend. The fast path in - :meth:`_to_state_backend` returns an already-compatible payload - unchanged, so matching streams incur no array copy. One representative - gate is used for the stream-level backend decision. Explicit sub-MPO - payloads are copied and converted with ``apply_to_arrays`` when needed, - preserving their tensor labels and operator bonds. + MPS uses Torch, JAX, CuPy, or a Symmray block backend. Every ordinary + gate and every tensor in every sub-MPO is checked; matching payloads + are returned by identity, while foreign payloads are copied or cast + without mutating the public queue. """ if not gates: return gates like = self._state_backend_like() - like_backend = None - like_dtype = None - if like is not None: - try: - like_backend = ar.infer_backend(like) - except (AttributeError, TypeError, ValueError): - pass - like_dtype = getattr(like, "dtype", None) - - # Gate streams are expected to be backend-homogeneous. Inspect one - # ordinary gate, then apply that decision to the whole executable - # segment so matching streams are left entirely untouched. - gate_needs_conversion = like_backend is None - gate_backend = None - if like_backend is not None: - for candidate, event_type in zip(gates, event_types): - if event_type != "gate": - continue - try: - gate_backend = ar.infer_backend(candidate) - gate_needs_conversion = ( - gate_backend != like_backend - or getattr(candidate, "dtype", None) != like_dtype - ) - except (AttributeError, TypeError, ValueError): - gate_needs_conversion = True - break - - if ( - gate_needs_conversion - and gate_backend is not None - and gate_backend != like_backend - and not self._backend_mismatch_warned - ): - warnings.warn( - "MpsOptimizer converted a gate payload from backend " - f"{gate_backend!r} to the live MPS backend " - f"{like_backend!r}; provide matching gate payloads to " - "avoid this conversion.", - UserWarning, - stacklevel=3, - ) - self._backend_mismatch_warned = True - + if like is None: + return gates + target_signature = _array_backend_signature(like) prepared = [] + stream_converter = infer_backend_converter_from_sample(like) for gate, event_type in zip(gates, event_types): if event_type == "gate": - if gate_needs_conversion: + source_signature = _array_backend_signature(gate) + if source_signature != target_signature: + self._warn_backend_conversion( + source_signature, target_signature, kind="gate" + ) gate = self.to_backend(gate) - elif event_type == "submpo" and like_backend is not None: + elif event_type == "submpo": # ``apply_to_arrays`` changes only the raw tensor payloads, # unlike rebuilding an MPO, which can lose custom labels or # operator bonds. Keep the caller's stream immutable by # applying it to a shallow network copy. - # Tensor-network payloads are expected to use one backend and - # dtype throughout, so inspect one representative tensor only. - tensor = next(iter(getattr(gate, "tensors", ())), None) - if tensor is None: - needs_conversion = False - else: - array = tensor.data - try: - needs_conversion = ( - ar.infer_backend(array) != like_backend - or getattr(array, "dtype", None) != like_dtype - ) - except (AttributeError, TypeError, ValueError): - needs_conversion = True - if needs_conversion: - if not self._backend_mismatch_warned: - warnings.warn( - "MpsOptimizer converted a sub-MPO payload to the " - f"live MPS backend {like_backend!r}; provide matching " - "sub-MPO payloads to avoid this conversion.", - UserWarning, - stacklevel=3, - ) - self._backend_mismatch_warned = True + tensors = tuple(getattr(gate, "tensors", ())) + source_signatures = { + _array_backend_signature(tensor.data) for tensor in tensors + } + if source_signatures and source_signatures != {target_signature}: + for source_signature in source_signatures: + if source_signature != target_signature: + self._warn_backend_conversion( + source_signature, target_signature, kind="sub-MPO" + ) gate = gate.copy() apply_to_arrays = getattr(gate, "apply_to_arrays", None) if not callable(apply_to_arrays): @@ -2867,7 +2858,7 @@ def _prepare_gate_stream_backend(self, gates, event_types): "sub-MPO payloads must provide apply_to_arrays() " "for backend conversion." ) - apply_to_arrays(self.to_backend) + apply_to_arrays(stream_converter or self.to_backend) prepared.append(gate) return prepared diff --git a/src/pepsy/optimizers/stabilizer_tn/mps_stab_optimizer.py b/src/pepsy/optimizers/stabilizer_tn/mps_stab_optimizer.py index d4b6e80..61ff136 100644 --- a/src/pepsy/optimizers/stabilizer_tn/mps_stab_optimizer.py +++ b/src/pepsy/optimizers/stabilizer_tn/mps_stab_optimizer.py @@ -49,6 +49,7 @@ import math import time +import warnings from collections.abc import Mapping from numbers import Integral from typing import List, Optional @@ -57,6 +58,11 @@ import numpy as np import quimb.tensor as qtn +from ...backends import ( + backend_infer, + infer_backend_converter_from_sample, + infer_backend_signature, +) from ..mps.layout import MpsGateStreamLayoutFinder from ..mps.optimizer import ( _resolve_conditional, @@ -625,13 +631,18 @@ def __init__( self.last_layout_plan = None self.to_backend = to_backend + self._explicit_backend_converter = to_backend self._bk_cache: dict = {} + self._backend_conversion_warnings = set() + self._backend_signature = None + self._backend_converter = None self._clifford_rot_cache: dict = {} self._localizer_cache: dict = {} if to_backend is not None: # Place the coefficient MPS |nu> on the requested backend; gate/MPO # arrays are converted on the fly by the _bk* helpers below. self.state.p.apply_to_arrays(to_backend) + self.backend_info() self._queue: List[object] = [] self.infidelities: List[float] = [] @@ -2948,26 +2959,135 @@ def _disentangle_event(self, params) -> list[dict]: # ------------------------------------------------------------------ # # Backend helpers (place |nu> gates/MPOs on the configured backend) # ------------------------------------------------------------------ # - def _bk(self, mat) -> np.ndarray: - """Backend copy of a (possibly parametrized) gate matrix (dtype-cast).""" + def _state_backend_like(self): + """Return a representative live coefficient-MPS array.""" + for tensor in getattr(self.state.p, "tensors", ()): + return tensor.data + return None + + def backend_info(self): + """Return and cache the live coefficient-MPS backend diagnostics.""" + info = backend_infer(self.state.p) + signature = infer_backend_signature(self._state_backend_like()) + if signature != self._backend_signature: + self._bk_cache.clear() + self._backend_signature = signature + self._backend_converter = ( + self._explicit_backend_converter + or infer_backend_converter_from_sample(self._state_backend_like()) + ) + self.backend = info["backend"] + self.backend_dtype = info["dtype"] + self.backend_device = info["device"] + self.array_backend = info.get("array_backend", info["backend"]) + # The live array dtype is authoritative when a caller supplies an + # existing MPS. This keeps generated coefficient operators aligned + # with the state rather than with STNState's constructor default. + self.dtype = info["dtype"] + self.state.dtype = self.dtype + return info + + def _warn_backend_conversion(self, source_signature, target_signature, *, kind): + """Warn once for an explicit stream payload conversion.""" + warning_key = (kind, source_signature, target_signature) + if ( + source_signature[0] != "builtins" + and warning_key not in self._backend_conversion_warnings + ): + self._backend_conversion_warnings.add(warning_key) + warnings.warn( + f"MpsStabOptimizer is converting a {kind} payload from " + f"backend/dtype/device {source_signature!r} to the live " + f"coefficient-MPS state {target_signature!r}; provide matching " + f"{kind} payloads to avoid this transfer or cast.", + UserWarning, + stacklevel=3, + ) + + def _to_state_backend(self, array, *, warn=False, kind="operator"): + """Return an array converted to the live coefficient-MPS signature.""" + like = self._state_backend_like() + if like is None: + return np.asarray(array, dtype=self.dtype) + # The full diagnostic validates every MPS tensor and is intentionally + # exposed through ``backend_info``. Internal rotations use the cached + # live signature so backend checks do not become an O(n) scan inside + # every gate/MPO contraction. + if self._backend_signature is None: + self.backend_info() + target_signature = self._backend_signature + source_signature = infer_backend_signature(array) + if source_signature == target_signature: + return array + if warn: + self._warn_backend_conversion(source_signature, target_signature, kind=kind) + if target_signature[0] == "symmray" and source_signature[0] != "symmray": + raise TypeError( + "Cannot convert a dense gate/operator payload into a native " + "Symmray MPS without charge and fermionic metadata. Build the " + "payload as a Symmray array on the target U1/U1U1 backend." + ) + converter = self._backend_converter + if converter is not None: + return converter(array) + return ar.do("array", array, like=like) + + def _diagnose_gate_backend(self, gate): + """Warn if an explicit stream matrix is foreign to the live MPS.""" + self.backend_info() + source_signature = infer_backend_signature(gate) + if source_signature != self._backend_signature: + self._warn_backend_conversion( + source_signature, self._backend_signature, kind="gate" + ) + + def _bk(self, mat): + """Backend copy of an internally generated gate matrix.""" arr = np.asarray(mat, dtype=self.dtype) - return self.to_backend(arr) if self.to_backend is not None else arr + return self._to_state_backend(arr) def _bk_const(self, tag: str, mat): """Backend copy of a *constant* gate matrix, cached by ``tag``.""" - if self.to_backend is None: - return np.asarray(mat, dtype=self.dtype) cached = self._bk_cache.get(tag) if cached is None: - cached = self.to_backend(np.asarray(mat, dtype=self.dtype)) + cached = self._to_state_backend(np.asarray(mat, dtype=self.dtype)) self._bk_cache[tag] = cached return cached - def _bk_mpo(self, mpo): - """Place a sub-MPO's arrays on the configured backend (in place).""" - if self.to_backend is not None: - mpo.apply_to_arrays(self.to_backend) - return mpo + def _bk_mpo(self, mpo, *, warn=True): + """Return a sub-MPO on the live backend without mutating its source.""" + tensors = tuple(getattr(mpo, "tensors", ())) + if not tensors: + return mpo + if self._backend_signature is None: + self.backend_info() + target_signature = self._backend_signature + source_signatures = { + infer_backend_signature(tensor.data) for tensor in tensors + } + if source_signatures == {target_signature}: + return mpo + for source_signature in source_signatures: + if source_signature != target_signature and warn: + self._warn_backend_conversion( + source_signature, target_signature, kind="sub-MPO" + ) + if target_signature[0] == "symmray": + for source_signature in source_signatures: + if source_signature[0] != "symmray": + raise TypeError( + "Cannot convert a dense sub-MPO into a native Symmray " + "MPS without charge and fermionic metadata." + ) + converted = mpo.copy() + apply_to_arrays = getattr(converted, "apply_to_arrays", None) + if not callable(apply_to_arrays): + raise TypeError( + "sub-MPO payloads must provide apply_to_arrays() for backend " + "conversion." + ) + apply_to_arrays(self._backend_converter or self._to_state_backend) + return converted @staticmethod def _to_scalar(x) -> complex: @@ -3416,6 +3536,7 @@ def _apply_entry(self, entry) -> None: raise ValueError(f"Unsupported gate stream entry: {entry!r}.") # matrix form: (gate_tensor, where) gate, where = entry + self._diagnose_gate_backend(gate) self._apply_matrix(self._gate_to_numpy(gate), where) return @@ -3661,7 +3782,7 @@ def _apply_rotation(self, name, params) -> None: coef = -1j * sign * np.sin(theta / 2) mps_terms = self._mps_terms(terms) mpo, where = pauli_combo_submpo(c, coef, mps_terms, self.n, dtype=self.dtype) - self._record(self._evolve_p(self._bk_mpo(mpo), where, unitary=True)) + self._record(self._evolve_p(self._bk_mpo(mpo, warn=False), where, unitary=True)) def _evolve_p( self, @@ -3953,14 +4074,16 @@ def cap(self, where, vec, *, absorb="left") -> "MpsStabOptimizer": dims=[2] * (n - 1), **split_opts, ) - if self.to_backend is not None: - p.apply_to_arrays(self.to_backend) + converter = self.to_backend or self._backend_converter + if converter is not None: + p.apply_to_arrays(converter) import stim tableau = stim.TableauSimulator() tableau.set_num_qubits(n - 1) self.state = STNState.from_tableau_and_state(tableau, p, dtype=self.dtype) + self.backend_info() self._localizer_cache.clear() self._invalidate_norm_infidelity() self._record() @@ -4122,7 +4245,7 @@ def _apply_projector( mps_terms = self._mps_terms(terms) mpo, where = pauli_combo_submpo(0.5, coef, mps_terms, self.n, dtype=self.dtype) self._evolve_p( - self._bk_mpo(mpo), + self._bk_mpo(mpo, warn=False), where, renormalize=True, norm_event=norm_event, @@ -5019,8 +5142,8 @@ def _apply_pauli_sum_submpo( ) mpo, where = pauli_sum_submpo(mapped, self.n, dtype=self.dtype) if unitary: - return self._evolve_p(self._bk_mpo(mpo), where, unitary=True) - self._evolve_p(self._bk_mpo(mpo), where) + return self._evolve_p(self._bk_mpo(mpo, warn=False), where, unitary=True) + self._evolve_p(self._bk_mpo(mpo, warn=False), where) if target_norm is None: self._invalidate_norm_infidelity() return None diff --git a/src/pepsy/optimizers/tree/optimizer.py b/src/pepsy/optimizers/tree/optimizer.py index 99cbb41..a655ac8 100644 --- a/src/pepsy/optimizers/tree/optimizer.py +++ b/src/pepsy/optimizers/tree/optimizer.py @@ -47,7 +47,12 @@ import numpy as np import quimb.tensor as qtn -from ...backends import infer_backend_converter_from_sample, to_float +from ...backends import ( + backend_infer, + infer_backend_converter_from_sample, + infer_backend_signature, + to_float, +) from ...operators.gates import _normalize_gate_entries from ..mps.optimizer import ( _control_event_parts as _mps_control_event_parts, @@ -149,13 +154,7 @@ def _is_symmray_array(array): def _array_backend_signature(array): """Return comparable backend / dtype / device metadata for an array.""" - backend = ar.infer_backend(array) - try: - dtype = ar.get_dtype_name(array) - except (AttributeError, KeyError, TypeError, ValueError): - dtype = str(getattr(array, "dtype", None)) - device = getattr(array, "device", None) - return backend, str(dtype), None if device is None else str(device) + return infer_backend_signature(array) def _operator_schmidt_rank(op, where, left_where): @@ -648,6 +647,7 @@ def __init__(self, gates=None, n=None, *, chi=64, else: self._install_tn(tn) self._thread_ind = None + self.backend_info() if run and self.G: if ( @@ -900,29 +900,19 @@ def _invalidate_state_norm_cache(self): @staticmethod def _state_backend_info(state): """Validate and describe the common backend of every state tensor.""" - tensor_map = getattr(state, "tensor_map", None) - if not tensor_map: - raise ValueError("initial state contains no tensors.") - data = [tensor.data for tensor in tensor_map.values()] - signature = _array_backend_signature(data[0]) - mismatched = [] - for array in data[1:]: - candidate = _array_backend_signature(array) - if candidate != signature: - mismatched.append(candidate) - if mismatched: - raise TypeError( - "Initial state tensors must use one compatible backend, dtype, " - "and device. Convert every tensor with the same backend " - "converter before constructing TreeOptimizer; found " - f"{signature!r} and {mismatched[0]!r}." - ) - backend, dtype, device = signature - return {"backend": backend, "dtype": dtype, "device": device} + return backend_infer(state) def backend_info(self): """Return the common backend, dtype, and device of the live TTN.""" - return self._state_backend_info(self.tn) + info = self._state_backend_info(self.tn) + # Keep a state-derived public diagnostic in addition to the detailed + # ``backend_info`` mapping. It is refreshed on every query so direct + # caller mutations cannot leave a stale optimizer backend label. + self.backend = info["backend"] + self.backend_dtype = info["dtype"] + self.backend_device = info["device"] + self.array_backend = info.get("array_backend", info["backend"]) + return info def _warn_backend_conversion(self, source_signature, target_signature): """Warn once for one explicit source/target backend conversion.""" @@ -958,9 +948,7 @@ def _as_state_backend(self, array, *, warn=True): """ like = self._state_like() state_info = self.backend_info() - target_signature = ( - state_info["backend"], state_info["dtype"], state_info["device"] - ) + target_signature = _array_backend_signature(like) source_signature = _array_backend_signature(array) if source_signature == target_signature: return array @@ -969,6 +957,12 @@ def _as_state_backend(self, array, *, warn=True): # array backends/dtypes still receive the transfer/cast warning. if warn: self._warn_backend_conversion(source_signature, target_signature) + if state_info["backend"] == "symmray" and source_signature[0] != "symmray": + raise TypeError( + "Cannot convert a dense gate/operator payload into a native " + "Symmray TTN without charge and fermionic metadata. Build the " + "payload as a Symmray array on the target U1/U1U1 backend." + ) if state_info["backend"] == "numpy": return ar.to_numpy(array) return self._backend_converter(like)(array) @@ -976,67 +970,53 @@ def _as_state_backend(self, array, *, warn=True): def _prepare_gate_stream_backend(self, payloads, event_types): """Prepare one executable gate/sub-MPO stream for the live backend. - Ordinary gates are expected to be backend-homogeneous. One - representative gate decides whether the whole gate stream needs - conversion; matching payloads are returned by identity. Sub-MPOs use - one representative tensor and ``apply_to_arrays`` on a copied network, - preserving the caller's labels and bonds. + Every ordinary gate and every tensor in every sub-MPO is checked; + matching payloads are returned by identity. Foreign sub-MPOs use + ``apply_to_arrays`` on a copied network, preserving the caller's + labels and operator bonds. """ if not payloads: return payloads like = self._state_like() - state_info = self.backend_info() - target_signature = ( - state_info["backend"], state_info["dtype"], state_info["device"] - ) + self.backend_info() + target_signature = _array_backend_signature(like) converter = None prepared = list(payloads) - gate_index = None - gate_signature = None for index, (payload, event_type) in enumerate( zip(payloads, event_types) ): if event_type != "gate": continue - gate_index = index try: - gate_signature = _array_backend_signature(payload) - except (AttributeError, TypeError, ValueError): - gate_signature = None - break - - if gate_index is not None: - gate_needs_conversion = gate_signature != target_signature - if gate_needs_conversion: - if gate_signature is not None: - self._warn_backend_conversion( - gate_signature, target_signature - ) - converter = self._backend_converter(like) - for index, event_type in enumerate(event_types): - if event_type == "gate": - prepared[index] = converter(payloads[index]) + source_signature = _array_backend_signature(payload) + except (AttributeError, KeyError, TypeError, ValueError): + source_signature = None + if source_signature == target_signature: + continue + if source_signature is not None: + self._warn_backend_conversion(source_signature, target_signature) + prepared[index] = self._as_state_backend(payload) for index, (payload, event_type) in enumerate( zip(payloads, event_types) ): if event_type != "submpo": continue - tensor = next(iter(getattr(payload, "tensors", ())), None) - if tensor is None: + tensors = tuple(getattr(payload, "tensors", ())) + if not tensors: continue - try: - source_signature = _array_backend_signature(tensor.data) - except (AttributeError, TypeError, ValueError): - source_signature = None - if source_signature == target_signature: + source_signatures = { + _array_backend_signature(tensor.data) for tensor in tensors + } + if source_signatures == {target_signature}: continue - if source_signature is not None: - self._warn_backend_conversion( - source_signature, target_signature - ) + for source_signature in source_signatures: + if source_signature != target_signature: + self._warn_backend_conversion( + source_signature, target_signature + ) if converter is None: converter = self._backend_converter(like) copied = payload.copy() @@ -1050,7 +1030,7 @@ def _prepare_gate_stream_backend(self, payloads, event_types): # when their dense operator is materialized. continue for op_tensor in tensor_map.values(): - op_tensor.modify(data=converter(op_tensor.data)) + op_tensor.modify(data=self._as_state_backend(op_tensor.data)) prepared[index] = copied return prepared diff --git a/src/pepsy/optimizers/tree_stabilizer/optimizer.py b/src/pepsy/optimizers/tree_stabilizer/optimizer.py index 5e7306a..2fa6b6c 100644 --- a/src/pepsy/optimizers/tree_stabilizer/optimizer.py +++ b/src/pepsy/optimizers/tree_stabilizer/optimizer.py @@ -22,6 +22,7 @@ import numpy as np import quimb.tensor as qtn +from ...backends import infer_backend_signature from ..stabilizer_tn.mps_stab_optimizer import ( DeferredInjectionRecord, DeferredInjectionReport, @@ -872,6 +873,7 @@ def __init__( self.frame_layout_events = tuple(frame_events) self.stim_plan = None self.stim_sample = None + self.backend_info() @classmethod def from_bits(cls, bits, **kwargs): @@ -1556,6 +1558,13 @@ def _apply_entry(self, entry): submpo_parts = submpo_event_parts(entry, normalize_where=True) if submpo_parts is not None: mpo, where = submpo_parts + # Convert every operator tensor once at the stream boundary. The + # ordinary TreeOptimizer helper copies foreign MPOs, preserving + # caller ownership and avoiding repeated per-tensor conversions in + # the native subtree path. + mpo = self._tree._prepare_gate_stream_backend( + [mpo], ["submpo"] + )[0] self._tree.apply_submpo(mpo, where) return if isinstance(entry, (tuple, list)) and entry: @@ -1648,10 +1657,24 @@ def _apply_entry(self, entry): raise ValueError(f"Unknown gate name {head!r} in stream entry {entry!r}.") if len(entry) != 2: raise ValueError(f"Unsupported gate stream entry: {entry!r}.") - self._apply_matrix(entry[0], entry[1]) + gate = entry[0] + self._diagnose_gate_backend(gate) + self._apply_matrix(gate, entry[1]) return raise ValueError(f"Unsupported gate stream entry: {entry!r}.") + def _diagnose_gate_backend(self, gate): + """Warn when an explicit matrix is foreign to the coefficient TTN.""" + target = self._tree._state_like() + if target is None: + return + source_signature = infer_backend_signature(gate) + target_signature = infer_backend_signature(target) + if source_signature != target_signature: + self._tree._warn_backend_conversion( + source_signature, target_signature + ) + def _apply_matrix(self, gate, where): where = _normalize_sites(where) gate = np.asarray(ar.to_numpy(gate), dtype=complex) @@ -3554,7 +3577,13 @@ def norm(self): def backend_info(self): """Return the coefficient TTN backend, dtype, and device.""" - return self._tree.backend_info() + info = self._tree.backend_info() + self.backend = info["backend"] + self.backend_dtype = info["dtype"] + self.backend_device = info["device"] + self.array_backend = info.get("array_backend", info["backend"]) + self.dtype = info["dtype"] + return info def normalize(self): """Normalize the coefficient TTN and return ``self``. diff --git a/tests/test_backends.py b/tests/test_backends.py index fbd28cd..8c03acd 100644 --- a/tests/test_backends.py +++ b/tests/test_backends.py @@ -2,6 +2,7 @@ import numpy as np import pytest +import quimb.tensor as qtn import pepsy @@ -10,6 +11,21 @@ def test_to_float_is_public_backend_helper(): assert pepsy.to_float is pepsy.backends.to_float +def test_backend_infer_is_available_at_the_high_level_and_for_mps(): + """The shared backend contract accepts arrays and tensor networks.""" + assert pepsy.backend_infer is pepsy.backends.backend_infer + + array_info = pepsy.backend_infer(np.ones(2, dtype=np.complex128)) + assert array_info["backend"] == "numpy" + assert array_info["dtype"] == "complex128" + + mps_info = pepsy.backend_infer( + qtn.MPS_computational_state("00", dtype="complex128") + ) + assert mps_info["backend"] == "numpy" + assert mps_info["dtype"] == "complex128" + + def test_to_float_handles_backend_scalar_without_numpy_coercion(): class BackendScalar: shape = () diff --git a/tests/test_optimize_mps.py b/tests/test_optimize_mps.py index 256a00b..77de9e9 100644 --- a/tests/test_optimize_mps.py +++ b/tests/test_optimize_mps.py @@ -241,6 +241,110 @@ def test_mps_optimizer_casts_submpo_stream_arrays_to_torch_state_backend(): assert all(isinstance(tensor.data, torch.Tensor) for tensor in optimizer.p.tensors) +def test_mps_optimizer_backend_diagnostics_and_late_gate_conversion(): + """Backend checks inspect every gate, not only the first stream payload.""" + torch = pytest.importorskip("torch") + + state = qtn.MPS_computational_state("00", dtype="complex128") + to_backend = py.backend_torch(dtype=torch.complex128, device="cpu") + state.apply_to_arrays(to_backend) + matching = to_backend(np.eye(2, dtype=complex)) + foreign = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex) + optimizer = py.MpsOptimizer( + state, + gates=[], + chi=2, + mode="mpo", + inplace=True, + ) + + assert optimizer.backend_info() == { + "backend": "torch", + "dtype": "complex128", + "device": "cpu", + } + assert optimizer.backend == "torch" + with pytest.warns(UserWarning, match="converted a gate payload"): + prepared = optimizer._prepare_gate_stream_backend( + [matching, foreign], ["gate", "gate"] + ) + assert prepared[0] is matching + assert isinstance(prepared[1], torch.Tensor) + + +def test_mps_optimizer_rejects_mixed_state_backends(): + """All live MPS tensors must agree on backend, dtype, and device.""" + torch = pytest.importorskip("torch") + + state = qtn.MPS_computational_state("00", dtype="complex128") + state[0].modify(data=torch.as_tensor(state[0].data, dtype=torch.complex128)) + with pytest.raises(TypeError, match="one compatible backend"): + py.MpsOptimizer(state, gates=[], chi=2, mode="mpo") + + +def test_mps_optimizer_reports_symmray_block_backend(): + """Symmray diagnostics retain the underlying Torch block backend.""" + pytest.importorskip("symmray") + torch = pytest.importorskip("torch") + + fermion = py.Fermion( + spinful=True, + symmetry="U1U1", + dtype="complex128", + ) + state = py.ps_to_mps( + 3, + fermion=fermion, + occupations=((1, 0), (0, 1), (1, 0)), + seed=1, + dtype="complex128", + ) + state.apply_to_arrays(py.backend_torch(dtype=torch.complex128, device="cpu")) + optimizer = py.MpsOptimizer(state, gates=[], chi=2, mode="mpo") + + assert optimizer.backend_info() == { + "backend": "symmray", + "dtype": "complex128", + "device": "cpu", + "array_backend": "torch", + } + + +def test_mps_optimizer_converts_symmray_submpo_blocks_to_state_backend(): + """Symmray sub-MPO copies preserve charge metadata while changing blocks.""" + pytest.importorskip("symmray") + torch = pytest.importorskip("torch") + + fermion = py.Fermion( + spinful=True, + symmetry="U1U1", + dtype="complex128", + ) + state = py.ps_to_mps( + 3, + fermion=fermion, + occupations=((1, 0), (0, 1), (1, 0)), + seed=1, + dtype="complex128", + ) + state.apply_to_arrays(py.backend_torch(dtype=torch.complex128, device="cpu")) + submpo = qtn.MatrixProductOperator.from_dense( + fermion.hopping_gate(0.01, t=1.0), + dims=(4, 4), + sites=(0, 1), + L=3, + ) + optimizer = py.MpsOptimizer(state, gates=[], chi=2, mode="mpo") + + with pytest.warns(UserWarning, match="converted a sub-MPO payload"): + converted = optimizer._prepare_gate_stream_backend( + [submpo], ["submpo"] + )[0] + assert converted is not submpo + assert all(tensor.data.backend == "torch" for tensor in converted.tensors) + assert all(tensor.data.backend == "numpy" for tensor in submpo.tensors) + + def test_mps_optimizer_accepts_perm_mode(): """Perm mode should expose an identity logical-to-physical ordering initially.""" p0 = qtn.MPS_computational_state("0000", dtype="complex128") diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index 095e934..ad4fd56 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -2666,6 +2666,56 @@ def test_tree_gate_stream_backend_preparation_is_stream_level(): assert untouched[1] is matching[1] +def test_tree_gate_stream_backend_preparation_checks_late_payloads(): + """A matching first gate cannot hide a later backend mismatch.""" + torch = pytest.importorskip("torch") + to_backend = pepsy.backend_torch(device="cpu", dtype=torch.complex128) + plan = TreePlan.from_order(range(2), structure="balanced") + state = TreeTensorNetwork.from_plan(plan) + state.apply_to_arrays(to_backend) + matching = to_backend(np.eye(2, dtype=complex)) + foreign = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex) + opt = TreeOptimizer(None, state=state, tree=plan, run=False) + + with pytest.warns(UserWarning, match="converting a gate/operator payload"): + prepared = opt._prepare_gate_stream_backend( + [matching, foreign], ["gate", "gate"] + ) + assert prepared[0] is matching + assert isinstance(prepared[1], torch.Tensor) + + +def test_tree_optimizer_reports_symmray_block_backend(): + """Native fermionic TTNs report their underlying block backend.""" + pytest.importorskip("symmray") + torch = pytest.importorskip("torch") + + fermion = pepsy.Fermion( + spinful=True, + symmetry="U1U1", + dtype="complex128", + ) + plan = TreePlan.from_order(range(3), structure="balanced") + state = pepsy.ps_to_ttn( + 3, + tree=plan, + fermion=fermion, + occupations=((1, 0), (0, 1), (1, 0)), + dtype="complex128", + ) + state.apply_to_arrays( + pepsy.backend_torch(device="cpu", dtype=torch.complex128) + ) + opt = TreeOptimizer(None, state=state, tree=plan, run=False) + + assert opt.backend_info() == { + "backend": "symmray", + "dtype": "complex128", + "device": "cpu", + "array_backend": "torch", + } + + def test_tree_submpo_stream_backend_preparation_preserves_input(): """A mismatched stream sub-MPO is copied and converted by its arrays.""" torch = pytest.importorskip("torch") diff --git a/tests/test_optimize_tree_stabilizer.py b/tests/test_optimize_tree_stabilizer.py index 050597a..eccb6fa 100644 --- a/tests/test_optimize_tree_stabilizer.py +++ b/tests/test_optimize_tree_stabilizer.py @@ -424,6 +424,11 @@ def test_tree_stab_torch_backend_matches_numpy(): inherited = pepsy.TreeStabOptimizer( native_state, max_dense_cap_qubits=4 ) + assert inherited.backend == "torch" + assert inherited.backend_dtype == "complex128" + assert inherited.backend_device == "cpu" + with pytest.warns(UserWarning, match="gate/operator payload"): + inherited.apply([(np.diag([1.0, np.exp(0.1j)]), 0)]) inherited.cap(1, [1.0, 0.0]) assert inherited.backend_info()["backend"] == "torch" assert inherited.validate_isometry_metadata() is inherited diff --git a/tests/test_public_api.py b/tests/test_public_api.py index 8e4adb9..b1cfb7b 100644 --- a/tests/test_public_api.py +++ b/tests/test_public_api.py @@ -52,7 +52,7 @@ def test_tree_optimizers_are_available_from_high_level_api(): "cnot", "cx", "cy", "cz", "swap", "iswap", "phase", "u1", "u2", "cphase", "crx", "cry", "crz", "cu1", "cu2", "cu3", "rx", "ry", "rz", "rxx", "ryy", "rzz", "u3", "su4", "fsim", "fsimg", "haar_random_state", "hrs_to_mps", "hrs_to_peps", "hrs_to_ttn", "ps_to_peps", "ps_to_3dpeps", "expec_mpo", - "id_to_mpo", "id_to_pepo", "ps_to_pepo", "ps_to_mpo", "ps_to_ttn", "make_numpy_array_caster", "to_float", "SweepOptimizer", + "id_to_mpo", "id_to_pepo", "ps_to_pepo", "ps_to_mpo", "ps_to_ttn", "make_numpy_array_caster", "backend_infer", "to_float", "SweepOptimizer", "FDSolver", "MpsEnergyOptimizer", "MpsOptimizer", "MpoOptimizer", "PepsEnergyOptimizer", "PepsOptimizer", "SimpleUpdateGen", "SymDMRG2", "PEPSSampleResult", "PepsBpSampler", "MpsSampler", "FermionConfigurationEncoding", "MpsDiagonalEstimate", "MpsBatchSampleResult", "MpsSampleResult", "VecSampler", "gate", "gauge_all", "gauge_all_simple", "one_norm_bp", "tn_fidelity", "tn_norm", "TreeSampler", "TreeBatchSampleResult", "TreeSampleResult", @@ -119,7 +119,7 @@ def test_internal_symbols_not_exported(): _CALLABLE_EXPORTS = [ "contract_boundary", "contract_flat", "build_bra_ket", "normalize", "peps_normalize", "boundary_norm", "peps_norm", "infidelity", "peps_infidelity", "peps_fidelity", - "to_float", "gauge_all", "gauge_all_simple", "one_norm_bp", + "backend_infer", "to_float", "gauge_all", "gauge_all_simple", "one_norm_bp", "GlobalOptimizer", "FIT", "tns_align", "measure_obs", "build_pepo_from_gates", "build_mpo_from_gates", "pauli", "x", "y", "z", "s", "sdg", "t", "tdg", "h", "hadamard", diff --git a/tests/test_stabilizer_tn.py b/tests/test_stabilizer_tn.py index 40ec9d2..3f7ccf1 100644 --- a/tests/test_stabilizer_tn.py +++ b/tests/test_stabilizer_tn.py @@ -2490,6 +2490,49 @@ def test_torch_backend_matrix_gate_input(): assert _fidelity(sim.to_statevector(), ref) == pytest.approx(1.0, abs=1e-6) +def test_native_mps_backend_is_inferred_and_foreign_payloads_are_diagnosed(): + torch = pytest.importorskip("torch") + backend = _torch_backend() + p = qtn.MatrixProductState.from_dense( + np.array([1, 0, 0, 0], dtype=complex), dims=[2, 2] + ) + p.apply_to_arrays(backend) + + sim = MpsStabOptimizer.from_mps(p) + assert sim.backend_info() == { + "backend": "torch", + "dtype": "complex128", + "device": "cpu", + } + assert sim.backend == "torch" + assert sim.backend_dtype == "complex128" + assert sim.backend_device == "cpu" + + with pytest.warns(UserWarning, match="gate payload"): + sim.apply([(np.diag([1.0, np.exp(0.1j)]), 0)]) + assert isinstance(sim.p[0].data, torch.Tensor) + sim.cap(0, [1.0, 0.0]) + assert sim.backend_info()["backend"] == "torch" + assert isinstance(sim.p[0].data, torch.Tensor) + + +def test_native_mps_submpo_conversion_does_not_mutate_source(): + backend = _torch_backend() + p = qtn.MatrixProductState.from_dense( + np.array([1, 0, 0, 0], dtype=complex), dims=[2, 2] + ) + p.apply_to_arrays(backend) + sim = MpsStabOptimizer.from_mps(p) + mpo = pauli_rotation_mpo(0.2, ["X", "Z"]) + source_types = tuple(type(tensor.data) for tensor in mpo.tensors) + + with pytest.warns(UserWarning, match="sub-MPO payload"): + sim.apply([("submpo", mpo, (0, 1))]) + + assert tuple(type(tensor.data) for tensor in mpo.tensors) == source_types + assert "torch" in type(sim.p[0].data).__module__ + + def test_torch_backend_injection_and_sampling(): tb = _torch_backend() # injection on the torch backend reproduces T From a54049c52027e343e1002e3e5239ecb3b19481e7 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Thu, 30 Jul 2026 10:32:03 -0700 Subject: [PATCH 28/70] stabilize backend linalg registration policy --- docs/api/tensors/core.md | 29 +- docs/development/modules/tensors.md | 10 +- src/pepsy/__init__.py | 8 + src/pepsy/backends/__init__.py | 4 + src/pepsy/backends/config.py | 139 +++++++- src/pepsy/backends/linalg.py | 3 + src/pepsy/backends/linalg_jax.py | 37 +- src/pepsy/backends/linalg_torch.py | 188 +++++++++-- src/pepsy/tensors/__init__.py | 4 + src/pepsy/tensors/core.py | 8 +- tests/test_backends.py | 500 ++++++++++++++++++++++++++++ tests/test_package_layout.py | 8 + tests/test_public_api.py | 18 +- 13 files changed, 894 insertions(+), 62 deletions(-) diff --git a/docs/api/tensors/core.md b/docs/api/tensors/core.md index 74b79c3..3c2d449 100644 --- a/docs/api/tensors/core.md +++ b/docs/api/tensors/core.md @@ -1,16 +1,33 @@ # `pepsy.tensors.core` -`reg_rel_svd_torch` is the preferred torch SVD registration for tensor-network -autodiff. It installs the relative-regularized SVD backward rule used by -`reg_complex_svd_torch`, and its CPU forward path falls back to SciPy `gesvd` -if `torch.linalg.svd` fails. +`register_torch_linalg()` keeps native Torch SVD/QR as the default. Pass +`stabilized=True` to opt into Pepsy's relative-regularized SVD and validated +real-QR rules for tensor-network autodiff. The explicit +`reg_rel_svd_torch()` helper remains available when only the stabilized SVD +rule is wanted; its CPU forward path falls back to SciPy `gesvd` if +`torch.linalg.svd` fails. +Use `register_jax_linalg()` for the same native-versus-stabilized choice on +JAX, or call `reg_native_svd_torch()` / `reg_native_svd_jax()` directly. The SVD/QR registration helpers are also available directly from `pepsy`, e.g. `import pepsy as py; py.reg_rel_svd_torch()`. Torch exposes -`reg_rel_svd_torch()`, `reg_real_svd_torch()`, `reg_complex_svd_torch()`, +`reg_native_svd_torch()`, `reg_rel_svd_torch()`, `reg_real_svd_torch()`, +`reg_complex_svd_torch()`, `reg_real_qr_torch()`, and `reg_complex_qr_torch()`. JAX exposes SVD aliases -`reg_rel_svd_jax()`, `reg_real_svd_jax()`, and `reg_complex_svd_jax()` for a +`reg_native_svd_jax()`, `reg_rel_svd_jax()`, `reg_real_svd_jax()`, and +`reg_complex_svd_jax()` for a thin-SVD custom VJP that preserves JAX's native derivative while safely restoring cotangents from Quimb fixed-rank truncation. +Registration helpers are idempotent: repeated calls do not re-register the +same Autoray implementation, while switching Torch between native/stabilized +or real/complex modes, and JAX between native/stabilized modes, intentionally +updates the active implementation. Stabilized real +QR supports square, tall, wide, and batched reduced QR. Its rank policy can be +`warn`, `native`, or `error`, and the tolerance is configurable. Complex mode +keeps native `torch.linalg.qr`; the explicit complex QR compatibility wrapper +uses the same conjugate-aware native VJP but is not registered because it +recomputes QR during backward. +Calling `reset_linalg_registrations()` restores native Torch/JAX mappings and +clears Pepsy's registration caches. > API details are maintained as handwritten Markdown in this page. diff --git a/docs/development/modules/tensors.md b/docs/development/modules/tensors.md index f8b72c8..5c9934c 100644 --- a/docs/development/modules/tensors.md +++ b/docs/development/modules/tensors.md @@ -52,10 +52,12 @@ Backend helpers manage package-wide defaults and optional linalg shims: - `set_default_array_backend(...)` / `get_default_array_backend()` - `set_default_grad_backend(...)` / `get_default_grad_backend()` - `reset_default_backends()` -- torch and JAX linalg/stop-gradient registrations. For torch SVD, - `reg_rel_svd_torch()` is the preferred full-SVD autodiff shim; it installs - the relative-regularized backward rule also used by `reg_complex_svd_torch()` - and falls back to SciPy `gesvd` on CPU forward-driver failures. +- torch and JAX linalg/stop-gradient registrations. Native thin SVD/QR is the + default through `register_torch_linalg()` and native thin SVD is the default + through `register_jax_linalg()`. The explicit `stabilized=True` mode, or + `reg_rel_svd_torch()` / `reg_rel_svd_jax()`, installs the truncation-safe, + relative-regularized SVD rules for workflows that need them. The stabilized + Torch SVD falls back to SciPy `gesvd` on CPU forward-driver failures. ## Tag and index conventions diff --git a/src/pepsy/__init__.py b/src/pepsy/__init__.py index bf5ebac..4241f91 100644 --- a/src/pepsy/__init__.py +++ b/src/pepsy/__init__.py @@ -48,7 +48,9 @@ "build_backend": ".backends", "get_default_array_backend": ".backends", "get_default_grad_backend": ".backends", + "register_jax_linalg": ".backends", "register_torch_linalg": ".backends", + "reset_linalg_registrations": ".backends", "reset_default_backends": ".backends", "set_default_array_backend": ".backends", "set_default_grad_backend": ".backends", @@ -238,6 +240,8 @@ "ps_to_pepo": ".tensors", "ps_to_peps": ".tensors", "random_haar_qubit": ".tensors", + "reg_native_svd_jax": ".tensors", + "reg_native_svd_torch": ".tensors", "reg_complex_qr_torch": ".tensors", "reg_complex_svd_jax": ".tensors", "reg_complex_svd_torch": ".tensors", @@ -286,7 +290,9 @@ def __getattr__(name): get_default_array_backend, get_default_grad_backend, build_backend, + register_jax_linalg, register_torch_linalg, + reset_linalg_registrations, reset_default_backends, set_default_array_backend, set_default_grad_backend, @@ -401,6 +407,8 @@ def __getattr__(name): ps_to_pepo, ps_to_peps, random_haar_qubit, + reg_native_svd_jax, + reg_native_svd_torch, reg_complex_qr_torch, reg_complex_svd_jax, reg_complex_svd_torch, diff --git a/src/pepsy/backends/__init__.py b/src/pepsy/backends/__init__.py index 1d09333..25dcc67 100644 --- a/src/pepsy/backends/__init__.py +++ b/src/pepsy/backends/__init__.py @@ -20,7 +20,9 @@ build_backend, get_default_array_backend, get_default_grad_backend, + register_jax_linalg, register_torch_linalg, + reset_linalg_registrations, reset_default_backends, set_default_array_backend, set_default_grad_backend, @@ -42,7 +44,9 @@ "to_float", "get_default_array_backend", "get_default_grad_backend", + "register_jax_linalg", "register_torch_linalg", + "reset_linalg_registrations", "reset_default_backends", "set_default_array_backend", "set_default_grad_backend", diff --git a/src/pepsy/backends/config.py b/src/pepsy/backends/config.py index 1856918..54f809d 100644 --- a/src/pepsy/backends/config.py +++ b/src/pepsy/backends/config.py @@ -12,9 +12,11 @@ __all__ = [ "build_backend", "backend_torch", "backend_numpy", "backend_cupy", "backend_jax", - "register_torch_linalg", "reg_rel_svd_torch", "reg_real_svd_torch", + "register_torch_linalg", "register_jax_linalg", "reg_native_svd_torch", + "reg_native_svd_jax", "reg_rel_svd_torch", "reg_real_svd_torch", "reg_complex_svd_torch", "reg_real_qr_torch", "reg_complex_qr_torch", "reg_rel_svd_jax", "reg_real_svd_jax", "reg_complex_svd_jax", + "reset_linalg_registrations", "reg_stop_gradient_torch", "stop_grad", "set_default_array_backend", "get_default_array_backend", "set_default_grad_backend", "get_default_grad_backend", "reset_default_backends", @@ -316,13 +318,26 @@ def cast_array(x, device=target_device, dtype=target_dtype): return cast_array -def register_torch_linalg(mode="complex"): - """Register custom torch linalg gradients in autoray. +def register_torch_linalg( + mode="complex", + *, + stabilized=False, + qr_rank_policy="warn", + qr_rank_tol_factor=1.0, +): + """Register Torch linalg rules in Autoray. Parameters ---------- mode : {"complex", "real"}, default="complex" - Which SVD/QR registrations to install. + Select the real or complex stabilized rule when ``stabilized=True``. + stabilized : bool, default=False + Keep native Torch SVD/QR by default. Set this to ``True`` to install + Pepsy's relative-regularized SVD and validated real-QR rules. + qr_rank_policy : {"warn", "native", "error"}, default="warn" + Response to rank-deficient inputs when stabilized real QR is active. + qr_rank_tol_factor : float, default=1.0 + Multiplier for the scale-aware real-QR rank threshold. """ if torch is None: # pragma: no cover - exercised in no-torch CI raise ImportError( @@ -332,16 +347,115 @@ def register_torch_linalg(mode="complex"): from ..backends import linalg_torch as lr # pylint: disable=import-outside-toplevel if mode == "complex": - lr.reg_rel_svd_torch() + if stabilized: + lr.reg_rel_svd_torch() + else: + lr.reg_native_svd_torch() lr.reg_complex_qr_torch() return if mode == "real": - lr.reg_real_svd_torch() - lr.reg_real_qr_torch() + if stabilized: + lr.reg_real_svd_torch() + lr.reg_real_qr_torch( + rank_policy=qr_rank_policy, + rank_tol_factor=qr_rank_tol_factor, + ) + else: + lr.reg_native_svd_torch() + lr.reg_complex_qr_torch() return raise ValueError("mode must be 'complex' or 'real'") +def register_jax_linalg(*, stabilized=False): + """Register native or truncation-safe JAX SVD in Autoray. + + Parameters + ---------- + stabilized : bool, default=False + Keep native thin SVD by default. Set this to ``True`` to install the + custom VJP that restores cotangents from Quimb fixed-rank truncation. + """ + try: + __import__("jax") + except ImportError as exc: # pragma: no cover - optional dependency + raise ImportError( + "register_jax_linalg requires optional dependency 'jax'. " + "Install it with: pip install jax jaxlib." + ) from exc + from ..backends import linalg_jax as lr # pylint: disable=import-outside-toplevel + + if stabilized: + lr.reg_complex_svd_jax() + else: + lr.reg_native_svd_jax() + + +def reset_linalg_registrations(backend="all"): + """Restore native linalg mappings and clear Pepsy registration caches. + + Parameters + ---------- + backend : {"torch", "jax", "all"}, default="all" + Which optional backend registration cache to reset. ``"all"`` skips + optional backends that are not installed. + """ + if backend not in {"torch", "jax", "all"}: + raise ValueError("backend must be one of: all, jax, torch") + + if backend in {"torch", "all"}: + if torch is None: + if backend == "torch": + raise ImportError( + "reset_linalg_registrations(backend='torch') requires " + "optional dependency 'torch'." + ) + else: + from ..backends import linalg_torch as lr # pylint: disable=import-outside-toplevel + + lr.reset_torch_linalg_registrations() + + if backend in {"jax", "all"}: + try: + __import__("jax") + except ImportError: + if backend == "jax": + raise ImportError( + "reset_linalg_registrations(backend='jax') requires " + "optional dependency 'jax'." + ) + else: + from ..backends import linalg_jax as lr # pylint: disable=import-outside-toplevel + + lr.reset_jax_linalg_registrations() + + +def reg_native_svd_torch(): + """Register native Torch thin SVD in autoray.""" + if torch is None: # pragma: no cover - exercised in no-torch CI + raise ImportError( + "reg_native_svd_torch requires optional dependency 'torch'. " + "Install it with: pip install pepsy[torch] (or pip install torch)." + ) + from ..backends import linalg_torch as lr # pylint: disable=import-outside-toplevel + + lr.reg_native_svd_torch() + + +def reg_native_svd_jax(): + """Register native JAX thin SVD in autoray.""" + try: + __import__("jax") + except ImportError as exc: # pragma: no cover - optional dependency + raise ImportError( + "reg_native_svd_jax requires optional dependency 'jax'. " + "Install it with: pip install jax jaxlib." + ) from exc + from ..backends import linalg_jax as lr # pylint: disable=import-outside-toplevel + + lr.reg_native_svd_jax() + + def reg_rel_svd_torch(): """Register torch SVD with a stable relative-regularized backward rule. @@ -396,7 +510,7 @@ def reg_real_svd_torch(): def reg_complex_qr_torch(): - """Register complex torch QR autograd rule in autoray.""" + """Register native Torch QR for complex inputs in autoray.""" if torch is None: # pragma: no cover - exercised in no-torch CI raise ImportError( "reg_complex_qr_torch requires optional dependency 'torch'. " @@ -407,8 +521,8 @@ def reg_complex_qr_torch(): lr.reg_complex_qr_torch() -def reg_real_qr_torch(): - """Register real torch QR autograd rule in autoray.""" +def reg_real_qr_torch(*, rank_policy="warn", rank_tol_factor=1.0): + """Register real Torch QR with a rank-deficiency policy.""" if torch is None: # pragma: no cover - exercised in no-torch CI raise ImportError( "reg_real_qr_torch requires optional dependency 'torch'. " @@ -416,7 +530,10 @@ def reg_real_qr_torch(): ) from ..backends import linalg_torch as lr # pylint: disable=import-outside-toplevel - lr.reg_real_qr_torch() + lr.reg_real_qr_torch( + rank_policy=rank_policy, + rank_tol_factor=rank_tol_factor, + ) def reg_complex_svd_jax(): diff --git a/src/pepsy/backends/linalg.py b/src/pepsy/backends/linalg.py index 664265e..5c391bb 100644 --- a/src/pepsy/backends/linalg.py +++ b/src/pepsy/backends/linalg.py @@ -15,11 +15,13 @@ "SVD_real", "QR_real", "QR_complex", + "reg_native_svd_torch", "reg_rel_svd_torch", "reg_complex_svd_torch", "reg_real_svd_torch", "reg_real_qr_torch", "reg_complex_qr_torch", + "reset_torch_linalg_registrations", "reg_stop_gradient_torch", "stop_grad", } @@ -28,6 +30,7 @@ "h", "jaxsvd_fwd", "jaxsvd_bwd", + "reg_native_svd_jax", "reg_complex_svd_jax", "reg_rel_svd_jax", "reg_real_svd_jax", diff --git a/src/pepsy/backends/linalg_jax.py b/src/pepsy/backends/linalg_jax.py index c989409..fa9684e 100644 --- a/src/pepsy/backends/linalg_jax.py +++ b/src/pepsy/backends/linalg_jax.py @@ -6,6 +6,10 @@ from jax import custom_vjp +_SVD_REGISTERED = False +_SVD_REGISTERED_FUNCTION = None + + @custom_vjp def svd_jax(A): """Thin JAX SVD with a Quimb-truncation-safe backward rule. @@ -86,9 +90,40 @@ def jaxsvd_bwd(residual, tangents): svd_jax.defvjp(jaxsvd_fwd, jaxsvd_bwd) +def _native_svd_jax(A, *args, **kwargs): + """Use native JAX SVD with Pepsy's thin-factor default.""" + kwargs.setdefault("full_matrices", False) + return jnp.linalg.svd(A, *args, **kwargs) + + +def _register_svd_jax(function): + """Register one JAX SVD implementation and remember the active rule.""" + global _SVD_REGISTERED # pylint: disable=global-statement + global _SVD_REGISTERED_FUNCTION # pylint: disable=global-statement + if _SVD_REGISTERED and _SVD_REGISTERED_FUNCTION is function: + return + ar.register_function("jax", "linalg.svd", function) + _SVD_REGISTERED = True + _SVD_REGISTERED_FUNCTION = function + + +def reg_native_svd_jax(): + """Register native JAX thin SVD in autoray.""" + _register_svd_jax(_native_svd_jax) + + +def reset_jax_linalg_registrations(): + """Restore native JAX SVD and clear Pepsy's registration cache.""" + global _SVD_REGISTERED # pylint: disable=global-statement + global _SVD_REGISTERED_FUNCTION # pylint: disable=global-statement + _SVD_REGISTERED = False + _SVD_REGISTERED_FUNCTION = None + reg_native_svd_jax() + + def reg_complex_svd_jax(): """Register the truncation-safe JAX thin-SVD implementation in autoray.""" - ar.register_function("jax", "linalg.svd", svd_jax) + _register_svd_jax(svd_jax) def reg_rel_svd_jax(): diff --git a/src/pepsy/backends/linalg_torch.py b/src/pepsy/backends/linalg_torch.py index 6750ee1..803209c 100644 --- a/src/pepsy/backends/linalg_torch.py +++ b/src/pepsy/backends/linalg_torch.py @@ -1,5 +1,7 @@ """Torch-side linalg registrations with stabilized autodiff rules.""" +import warnings + import autoray as ar import numpy as np import torch @@ -12,6 +14,51 @@ # pylint: disable=abstract-method,arguments-differ,bad-staticmethod-argument,bare-except,line-too-long,multiple-statements,not-callable,superfluous-parens,too-many-branches,too-many-locals,too-many-statements,unnecessary-semicolon,unused-variable,using-constant-test _SVD_EPS_REL = 1.0e-6 +_REGISTERED_FUNCTIONS = {} +_QR_RANK_POLICIES = {"warn", "native", "error"} +_QR_RANK_POLICY = "warn" +_QR_RANK_TOL_FACTOR = 1.0 + + +def _same_callable(left, right): + """Compare plain functions and class-bound autograd methods robustly.""" + if left is right: + return True + left_func = getattr(left, "__func__", None) + right_func = getattr(right, "__func__", None) + return ( + left_func is not None + and right_func is not None + and left_func is right_func + and getattr(left, "__self__", None) + is getattr(right, "__self__", None) + ) + + +def _register_once(name, function): + """Register one Torch autoray function once per active implementation.""" + if _same_callable(_REGISTERED_FUNCTIONS.get(name), function): + return + ar.register_function("torch", name, function) + _REGISTERED_FUNCTIONS[name] = function + + +def _configure_qr_rank_policy(policy="warn", rank_tol_factor=1.0): + """Configure the real-QR response to detected rank deficiency.""" + if policy not in _QR_RANK_POLICIES: + choices = ", ".join(sorted(_QR_RANK_POLICIES)) + raise ValueError(f"rank_policy must be one of: {choices}") + try: + rank_tol_factor = float(rank_tol_factor) + except (TypeError, ValueError) as exc: + raise TypeError("rank_tol_factor must be a positive finite number") from exc + if not np.isfinite(rank_tol_factor) or rank_tol_factor <= 0.0: + raise ValueError("rank_tol_factor must be a positive finite number") + + global _QR_RANK_POLICY # pylint: disable=global-statement + global _QR_RANK_TOL_FACTOR # pylint: disable=global-statement + _QR_RANK_POLICY = policy + _QR_RANK_TOL_FACTOR = rank_tol_factor def safe_inverse(x, eps_abs=1.0e-12, *, eps_rel=0.0, eps_scale=None): @@ -328,27 +375,52 @@ def backward(ctx, gu, gsigma, gvh): class QR_real(torch.autograd.Function): - """Real-valued QR autograd function using a custom backward pass.""" + """Real QR with a custom full-rank backward and native rank fallback.""" @staticmethod def forward(self, A): Q, R = torch.linalg.qr(A) - self.save_for_backward(A, Q, R) + diagonal = torch.diagonal(R, dim1=-2, dim2=-1).abs() + scale = R.abs().amax(dim=(-2, -1)) + tolerance = ( + _QR_RANK_TOL_FACTOR + * torch.finfo(A.dtype).eps + * max(A.shape[-2:]) + * scale + ) + rank_deficient = (diagonal <= tolerance.unsqueeze(-1)).any(dim=-1) + if bool(rank_deficient.any().item()): + message = ( + "QR_real detected a rank-deficient input; native Torch QR " + "backward may be ill-conditioned." + ) + if _QR_RANK_POLICY == "error": + raise RuntimeError(message) + if _QR_RANK_POLICY == "warn": + warnings.warn(message, RuntimeWarning, stacklevel=2) + self.save_for_backward(A, Q, R, rank_deficient) return Q, R @staticmethod def backward(self, dq, dr): - A, q, r = self.saved_tensors - if r.shape[0] == r.shape[1]: + A, q, r, rank_deficient = self.saved_tensors + if bool(rank_deficient.any().item()): + return _native_qr_backward(A, dq, dr) + m, _n = r.shape[-2:] + if m == _n: return _simple_qr_backward(q, r, dq, dr) - M, _N = r.shape - B = A[:, M:] - dU = dr[:, :M] - dD = dr[:, M:] - U = r[:, :M] - da = _simple_qr_backward(q, U, dq + B @ dD.t(), dU) + B = A[..., :, m:] + dU = dr[..., :, :m] + dD = dr[..., :, m:] + U = r[..., :, :m] + da = _simple_qr_backward( + q, + U, + dq + B @ dD.transpose(-2, -1), + dU, + ) db = q @ dD - return torch.cat([da, db], 1) + return torch.cat([da, db], dim=-1) def _simple_qr_backward(q, r, dq, dr): @@ -359,49 +431,77 @@ def _simple_qr_backward(q, r, dq, dr): "or full_matrices is true and ncols != nrows." ) - qdq = q.t() @ dq - qdq_ = qdq - qdq.t() - rdr = r @ dr.t() - rdr_ = rdr - rdr.t() + qdq = q.transpose(-2, -1) @ dq + qdq_ = qdq - qdq.transpose(-2, -1) + rdr = r @ dr.transpose(-2, -1) + rdr_ = rdr - rdr.transpose(-2, -1) tril = torch.tril(qdq_ + rdr_) def _triangular_solve(x, tri): - return torch.linalg.solve_triangular(tri.T, x.T, upper=True).T + # Solve with the upper-triangular R factor. Using R.T here gives the + # wrong reverse-mode rule and does not generalize to batched inputs. + return torch.linalg.solve_triangular( + tri, + x.transpose(-2, -1), + upper=True, + ).transpose(-2, -1) grad_a = q @ (dr + _triangular_solve(tril, r)) grad_b = _triangular_solve(dq - q @ qdq, r) return grad_a + grad_b +def _native_qr_backward(A, dq, dr): + """Recompute native Torch QR to obtain a reliable VJP fallback.""" + with torch.enable_grad(): + replay = A.detach().requires_grad_(True) + q, r = torch.linalg.qr(replay) + dq = torch.zeros_like(q) if dq is None else dq + dr = torch.zeros_like(r) if dr is None else dr + return torch.autograd.grad( + (q, r), + replay, + (dq, dr), + )[0] + + class QR_complex(torch.autograd.Function): - """Complex-valued QR autograd function using Hermitian symmetrization.""" + """Complex-valued QR wrapper using native Torch's conjugate-aware VJP.""" @staticmethod def forward(ctx, A): Q, R = torch.linalg.qr(A) - ctx.save_for_backward(A, Q, R) + ctx.save_for_backward(A) return Q, R @staticmethod def backward(ctx, dQ, dR): - _A, Q, R = ctx.saved_tensors + (A,) = ctx.saved_tensors + return _native_qr_backward(A, dQ, dR) - Qh = Q.conj().transpose(-2, -1) - Rh = R.conj().transpose(-2, -1) - M = R @ dR.conj().transpose(-2, -1) - Qh @ dQ - sym_h_M = 0.5 * (M + M.conj().transpose(-2, -1)) - R_inv_h = torch.linalg.solve( - Rh, - torch.eye(Rh.size(-1), dtype=R.dtype, device=R.device), - ) - dA = (dQ + Q @ sym_h_M) @ R_inv_h - return dA +def _native_svd(A, *args, **kwargs): + """Use native Torch SVD with the thin-factor default expected by Pepsy.""" + kwargs.setdefault("full_matrices", False) + return torch.linalg.svd(A, *args, **kwargs) + + +def reg_native_svd_torch(): + """Register native Torch SVD with Pepsy's thin-factor default.""" + _register_once("linalg.svd", _native_svd) + + +def reset_torch_linalg_registrations(): + """Restore native Torch SVD/QR mappings and clear Pepsy's cache.""" + _REGISTERED_FUNCTIONS.clear() + _configure_qr_rank_policy() + reg_native_svd_torch() + reg_complex_qr_torch() def reg_rel_svd_torch(): """Register the relative-regularized torch SVD rule in autoray.""" - ar.register_function("torch", "linalg.svd", SVD.apply) + _register_once("linalg.svd", SVD.apply) def reg_complex_svd_torch(): @@ -415,22 +515,36 @@ def reg_complex_svd_torch(): def reg_real_svd_torch(): """Register the real torch SVD autograd implementation in autoray.""" - ar.register_function("torch", "linalg.svd", SVD_real.apply) + _register_once("linalg.svd", SVD_real.apply) -def reg_real_qr_torch(): - """Register the real torch QR autograd implementation in autoray.""" - ar.register_function("torch", "linalg.qr", QR_real.apply) +def reg_real_qr_torch(*, rank_policy="warn", rank_tol_factor=1.0): + """Register real QR with a configurable rank-deficiency policy.""" + _configure_qr_rank_policy(rank_policy, rank_tol_factor) + _register_once("linalg.qr", QR_real.apply) def reg_complex_qr_torch(): - """Register the complex torch QR autograd implementation in autoray.""" - ar.register_function("torch", "linalg.qr", QR_complex.apply) + """Register native Torch QR for complex inputs. + + The explicit :class:`QR_complex` compatibility wrapper is safe but + recomputes native QR during its backward pass, so Autoray uses native QR + directly for the faster default path. + """ + _register_once("linalg.qr", torch.linalg.qr) + + +def _stop_gradient_torch(x): + """Return a detached, independent Torch tensor.""" + return x.detach().clone() def reg_stop_gradient_torch(): """Register torch stop-gradient helper in autoray.""" - ar.register_function("torch", "stop_gradient", lambda x: x.detach().clone()) + _register_once( + "stop_gradient", + _stop_gradient_torch, + ) def stop_grad(x): diff --git a/src/pepsy/tensors/__init__.py b/src/pepsy/tensors/__init__.py index 783b853..c315895 100644 --- a/src/pepsy/tensors/__init__.py +++ b/src/pepsy/tensors/__init__.py @@ -102,11 +102,15 @@ def _register(module, *names): "build_backend", "get_default_array_backend", "get_default_grad_backend", + "register_jax_linalg", "register_torch_linalg", + "reset_linalg_registrations", "reset_default_backends", "set_default_array_backend", "set_default_grad_backend", "reg_complex_qr_torch", + "reg_native_svd_jax", + "reg_native_svd_torch", "reg_complex_svd_jax", "reg_complex_svd_torch", "reg_real_qr_torch", diff --git a/src/pepsy/tensors/core.py b/src/pepsy/tensors/core.py index 6369953..13a8a7f 100644 --- a/src/pepsy/tensors/core.py +++ b/src/pepsy/tensors/core.py @@ -47,6 +47,9 @@ build_backend, get_default_array_backend, get_default_grad_backend, + register_jax_linalg, + reg_native_svd_jax, + reg_native_svd_torch, reg_complex_qr_torch, reg_complex_svd_jax, reg_complex_svd_torch, @@ -56,6 +59,7 @@ reg_rel_svd_jax, reg_rel_svd_torch, reg_stop_gradient_torch, + reset_linalg_registrations, register_torch_linalg, reset_default_backends, set_default_array_backend, @@ -118,9 +122,11 @@ def contract_hypercompressed_tn(*args, **kwargs): __all__ = [ "OneDMap", "build_backend", "backend_torch", "backend_numpy", "backend_cupy", "backend_jax", - "register_torch_linalg", "reg_rel_svd_torch", "reg_real_svd_torch", + "register_torch_linalg", "register_jax_linalg", "reg_native_svd_torch", + "reg_native_svd_jax", "reg_rel_svd_torch", "reg_real_svd_torch", "reg_complex_svd_torch", "reg_real_qr_torch", "reg_complex_qr_torch", "reg_rel_svd_jax", "reg_real_svd_jax", "reg_complex_svd_jax", + "reset_linalg_registrations", "reg_stop_gradient_torch", "stop_grad", "set_default_array_backend", "get_default_array_backend", "set_default_grad_backend", "get_default_grad_backend", "reset_default_backends", "build_contraction", "build_optimizer", "build_compressed_optimizer", diff --git a/tests/test_backends.py b/tests/test_backends.py index 8c03acd..caafd34 100644 --- a/tests/test_backends.py +++ b/tests/test_backends.py @@ -1,5 +1,7 @@ """Tests for public backend conversion helpers.""" +import warnings + import numpy as np import pytest import quimb.tensor as qtn @@ -7,6 +9,22 @@ import pepsy +def _available_torch_devices(): + """Return Torch devices available to backend integration tests.""" + try: + import torch + except ImportError: + return ["cpu"] + + devices = ["cpu"] + if torch.cuda.is_available(): + devices.append("cuda") + if getattr(torch.backends, "mps", None) is not None: + if torch.backends.mps.is_available(): + devices.append("mps") + return devices + + def test_to_float_is_public_backend_helper(): assert pepsy.to_float is pepsy.backends.to_float @@ -26,6 +44,34 @@ def test_backend_infer_is_available_at_the_high_level_and_for_mps(): assert mps_info["dtype"] == "complex128" +@pytest.mark.parametrize("device", _available_torch_devices()) +def test_torch_backend_and_linalg_on_available_devices(device): + """Backend inference and native linalg agree on CPU/CUDA/MPS devices.""" + torch = pytest.importorskip("torch") + + dtype = torch.float64 if device == "cpu" else torch.float32 + to_backend = pepsy.backend_torch(device=device, dtype=dtype) + matrix = to_backend(np.arange(12, dtype=np.float64).reshape(4, 3)) + info = pepsy.backend_infer(matrix) + assert info == { + "backend": "torch", + "dtype": str(dtype).removeprefix("torch."), + "device": str(matrix.device), + } + + try: + q, r = torch.linalg.qr(matrix) + _, sigma, _ = torch.linalg.svd(matrix, full_matrices=False) + except (RuntimeError, NotImplementedError) as exc: + if device == "cpu": + raise + pytest.skip(f"Torch {device} linalg is unavailable: {exc}") + + assert q.device == matrix.device + assert r.device == matrix.device + assert sigma.device == matrix.device + + def test_to_float_handles_backend_scalar_without_numpy_coercion(): class BackendScalar: shape = () @@ -63,6 +109,460 @@ def test_register_torch_svd_for_autoray(): assert vh.shape == (2, 2) +def test_torch_linalg_registration_is_idempotent(monkeypatch): + """Repeated public/backend registration does not re-patch Autoray.""" + pytest.importorskip("torch") + import autoray as ar + from pepsy.backends import linalg_torch + + calls = [] + original_registered = dict(linalg_torch._REGISTERED_FUNCTIONS) + linalg_torch._REGISTERED_FUNCTIONS.pop("linalg.svd", None) + monkeypatch.setattr( + ar, + "register_function", + lambda *args, **kwargs: calls.append((args, kwargs)), + ) + try: + linalg_torch.reg_real_svd_torch() + linalg_torch.reg_real_svd_torch() + finally: + linalg_torch._REGISTERED_FUNCTIONS.clear() + linalg_torch._REGISTERED_FUNCTIONS.update(original_registered) + + assert len(calls) == 1 + + +def test_torch_linalg_registration_can_switch_svd_modes(monkeypatch): + """Real and relative SVD modes can intentionally replace one another.""" + pytest.importorskip("torch") + import autoray as ar + from pepsy.backends import linalg_torch + + calls = [] + original_registered = dict(linalg_torch._REGISTERED_FUNCTIONS) + linalg_torch._REGISTERED_FUNCTIONS.pop("linalg.svd", None) + monkeypatch.setattr( + ar, + "register_function", + lambda *args, **kwargs: calls.append((args, kwargs)), + ) + try: + linalg_torch.reg_real_svd_torch() + linalg_torch.reg_rel_svd_torch() + linalg_torch.reg_rel_svd_torch() + finally: + linalg_torch._REGISTERED_FUNCTIONS.clear() + linalg_torch._REGISTERED_FUNCTIONS.update(original_registered) + + assert len(calls) == 2 + + +@pytest.mark.parametrize("shape", ((4, 4), (5, 3), (3, 5), (2, 5, 3))) +def test_torch_real_qr_backward_matches_native(shape): + """The validated real QR rule matches Torch for all reduced shapes.""" + torch = pytest.importorskip("torch") + from pepsy.backends.linalg_torch import QR_real + + torch.manual_seed(100 + sum(shape)) + matrix = torch.randn(*shape, dtype=torch.float64, requires_grad=True) + q, r = QR_real.apply(matrix) + dq = torch.randn_like(q) + dr = torch.randn_like(r) + actual = torch.autograd.grad((q * dq).sum() + (r * dr).sum(), matrix)[0] + + native_matrix = matrix.detach().clone().requires_grad_() + native_q, native_r = torch.linalg.qr(native_matrix) + expected = torch.autograd.grad( + (native_q * dq).sum() + (native_r * dr).sum(), + native_matrix, + )[0] + + torch.testing.assert_close(actual, expected, rtol=1e-9, atol=1e-10) + assert torch.isfinite(actual).all() + + +@pytest.mark.parametrize("shape", ((4, 4), (5, 3), (3, 5), (2, 5, 3))) +def test_torch_complex_qr_backward_matches_native(shape): + """The explicit complex QR wrapper preserves Torch's conjugate VJP.""" + torch = pytest.importorskip("torch") + from pepsy.backends.linalg_torch import QR_complex + + torch.manual_seed(200 + sum(shape)) + matrix = torch.randn(*shape, dtype=torch.complex128) + matrix = matrix + 1j * torch.randn(*shape, dtype=torch.complex128) + matrix.requires_grad_() + q, r = QR_complex.apply(matrix) + dq = torch.randn_like(q) + 1j * torch.randn_like(q) + dr = torch.randn_like(r) + 1j * torch.randn_like(r) + actual = torch.autograd.grad( + (q.conj() * dq).real.sum() + (r.conj() * dr).real.sum(), + matrix, + )[0] + + native_matrix = matrix.detach().clone().requires_grad_() + native_q, native_r = torch.linalg.qr(native_matrix) + expected = torch.autograd.grad( + (native_q.conj() * dq).real.sum() + (native_r.conj() * dr).real.sum(), + native_matrix, + )[0] + + torch.testing.assert_close(actual, expected, rtol=1e-9, atol=1e-10) + assert torch.isfinite(actual).all() + + +def test_torch_complex_qr_wrapper_passes_gradcheck(): + """Complex QR remains locally differentiable away from rank loss.""" + torch = pytest.importorskip("torch") + from pepsy.backends.linalg_torch import QR_complex + + torch.manual_seed(250) + matrix = torch.randn(3, 2, dtype=torch.complex128) + matrix = (matrix + 1j * torch.randn_like(matrix)).requires_grad_() + dq = torch.randn(3, 2, dtype=torch.complex128) + dr = torch.randn(2, 2, dtype=torch.complex128) + + def loss(value): + q, r = QR_complex.apply(value) + return (q.conj() * dq).real.sum() + (r.conj() * dr).real.sum() + + assert torch.autograd.gradcheck( + loss, + (matrix,), + eps=1.0e-6, + atol=1.0e-5, + rtol=1.0e-4, + ) + + +def test_torch_real_qr_rank_deficient_falls_back_to_native(): + """Rank-deficient real QR warns and follows native Torch backward.""" + torch = pytest.importorskip("torch") + from pepsy.backends.linalg_torch import QR_real + + torch.manual_seed(275) + matrix = torch.randn(4, 3, dtype=torch.float64) + matrix[:, 1] = matrix[:, 0] + matrix.requires_grad_() + + with pytest.warns(RuntimeWarning, match="rank-deficient"): + q, r = QR_real.apply(matrix) + dq = torch.randn_like(q) + dr = torch.randn_like(r) + actual = torch.autograd.grad((q * dq).sum() + (r * dr).sum(), matrix)[0] + + native_matrix = matrix.detach().clone().requires_grad_() + native_q, native_r = torch.linalg.qr(native_matrix) + expected = torch.autograd.grad( + (native_q * dq).sum() + (native_r * dr).sum(), + native_matrix, + )[0] + + torch.testing.assert_close(actual, expected, rtol=1e-9, atol=1e-10) + assert torch.isfinite(actual).all() + + +@pytest.mark.parametrize("complex_input", (False, True)) +@pytest.mark.parametrize("case", ("zero", "repeated", "rank_deficient")) +def test_torch_svd_degenerate_inputs_have_finite_gradients(case, complex_input): + """Regularized SVD gradients stay finite for singular edge cases.""" + torch = pytest.importorskip("torch") + from pepsy.backends.linalg_torch import SVD, SVD_real + + real_dtype = torch.float64 + if case == "zero": + matrix = torch.zeros(3, 3, dtype=real_dtype) + elif case == "repeated": + matrix = torch.diag(torch.tensor((2.0, 2.0, 0.5), dtype=real_dtype)) + else: + matrix = torch.tensor( + ((1.0, 2.0, 3.0), (2.0, 4.0, 6.0), (0.0, 1.0, 1.0)), + dtype=real_dtype, + ) + if complex_input: + matrix = matrix.to(torch.complex128) + matrix = matrix + 0.1j * torch.eye(3, dtype=torch.complex128) + matrix.requires_grad_() + + torch.manual_seed(300 + len(case) + int(complex_input)) + svd = SVD if complex_input else SVD_real + u, sigma, vh = svd.apply(matrix) + gu = torch.randn_like(u) + gsigma = torch.randn_like(sigma) + gvh = torch.randn_like(vh) + loss = ( + (u.conj() * gu).real.sum() + + (sigma * gsigma).real.sum() + + (vh.conj() * gvh).real.sum() + ) + gradient = torch.autograd.grad(loss, matrix)[0] + + assert torch.isfinite(gradient).all() + + +def test_torch_complex_qr_registration_uses_native_fallback(monkeypatch): + """The analytical complex QR rule is not registered through Autoray.""" + torch = pytest.importorskip("torch") + import autoray as ar + from pepsy.backends import linalg_torch + + calls = [] + original_registered = dict(linalg_torch._REGISTERED_FUNCTIONS) + linalg_torch._REGISTERED_FUNCTIONS.pop("linalg.qr", None) + monkeypatch.setattr( + ar, + "register_function", + lambda *args, **kwargs: calls.append((args, kwargs)), + ) + try: + linalg_torch.reg_complex_qr_torch() + finally: + linalg_torch._REGISTERED_FUNCTIONS.clear() + linalg_torch._REGISTERED_FUNCTIONS.update(original_registered) + + assert len(calls) == 1 + assert calls[0][0][0:2] == ("torch", "linalg.qr") + assert linalg_torch._same_callable(calls[0][0][2], torch.linalg.qr) + + +def test_register_torch_linalg_complex_uses_native_defaults(monkeypatch): + """The default complex umbrella registration keeps native linalg.""" + torch = pytest.importorskip("torch") + import autoray as ar + from pepsy.backends import config, linalg_torch + + calls = [] + original_registered = dict(linalg_torch._REGISTERED_FUNCTIONS) + linalg_torch._REGISTERED_FUNCTIONS.pop("linalg.svd", None) + linalg_torch._REGISTERED_FUNCTIONS.pop("linalg.qr", None) + monkeypatch.setattr( + ar, + "register_function", + lambda *args, **kwargs: calls.append((args, kwargs)), + ) + try: + config.register_torch_linalg(mode="complex") + finally: + linalg_torch._REGISTERED_FUNCTIONS.clear() + linalg_torch._REGISTERED_FUNCTIONS.update(original_registered) + + registered = {args[1]: args[2] for args, _kwargs in calls} + assert linalg_torch._same_callable(registered["linalg.qr"], torch.linalg.qr) + assert linalg_torch._same_callable( + registered["linalg.svd"], + linalg_torch._native_svd, + ) + + +def test_register_torch_linalg_stabilized_real_is_opt_in(monkeypatch): + """Stabilized real SVD/QR rules require an explicit opt-in.""" + pytest.importorskip("torch") + import autoray as ar + from pepsy.backends import config, linalg_torch + + calls = [] + original_registered = dict(linalg_torch._REGISTERED_FUNCTIONS) + original_policy = linalg_torch._QR_RANK_POLICY + original_factor = linalg_torch._QR_RANK_TOL_FACTOR + linalg_torch._REGISTERED_FUNCTIONS.pop("linalg.svd", None) + linalg_torch._REGISTERED_FUNCTIONS.pop("linalg.qr", None) + monkeypatch.setattr( + ar, + "register_function", + lambda *args, **kwargs: calls.append((args, kwargs)), + ) + try: + config.register_torch_linalg( + mode="real", + stabilized=True, + qr_rank_policy="error", + qr_rank_tol_factor=2.0, + ) + registered = {args[1]: args[2] for args, _kwargs in calls} + assert linalg_torch._same_callable( + registered["linalg.svd"], + linalg_torch.SVD_real.apply, + ) + assert linalg_torch._same_callable( + registered["linalg.qr"], + linalg_torch.QR_real.apply, + ) + assert linalg_torch._QR_RANK_POLICY == "error" + assert linalg_torch._QR_RANK_TOL_FACTOR == 2.0 + finally: + linalg_torch._QR_RANK_POLICY = original_policy + linalg_torch._QR_RANK_TOL_FACTOR = original_factor + linalg_torch._REGISTERED_FUNCTIONS.clear() + linalg_torch._REGISTERED_FUNCTIONS.update(original_registered) + + +def test_reset_linalg_registrations_restores_native_torch(monkeypatch): + """The public reset helper restores native Torch mappings.""" + torch = pytest.importorskip("torch") + import autoray as ar + from pepsy.backends import config, linalg_torch + + calls = [] + original_registered = dict(linalg_torch._REGISTERED_FUNCTIONS) + linalg_torch._REGISTERED_FUNCTIONS.clear() + monkeypatch.setattr( + ar, + "register_function", + lambda *args, **kwargs: calls.append((args, kwargs)), + ) + try: + config.reset_linalg_registrations(backend="torch") + finally: + linalg_torch._REGISTERED_FUNCTIONS.clear() + linalg_torch._REGISTERED_FUNCTIONS.update(original_registered) + + registered = {args[1]: args[2] for args, _kwargs in calls} + assert linalg_torch._same_callable( + registered["linalg.svd"], + linalg_torch._native_svd, + ) + assert linalg_torch._same_callable(registered["linalg.qr"], torch.linalg.qr) + + +def test_torch_real_qr_rank_policy_error_is_strict(): + """The strict QR policy rejects rank-deficient inputs before backward.""" + torch = pytest.importorskip("torch") + from pepsy.backends import linalg_torch + + original_policy = linalg_torch._QR_RANK_POLICY + original_factor = linalg_torch._QR_RANK_TOL_FACTOR + matrix = torch.randn(4, 3, dtype=torch.float64) + matrix[:, 1] = matrix[:, 0] + try: + linalg_torch._configure_qr_rank_policy("error") + with pytest.raises(RuntimeError, match="rank-deficient"): + linalg_torch.QR_real.apply(matrix) + finally: + linalg_torch._QR_RANK_POLICY = original_policy + linalg_torch._QR_RANK_TOL_FACTOR = original_factor + + +def test_torch_real_qr_rank_policy_native_is_silent(): + """The native rank policy falls back without emitting a warning.""" + torch = pytest.importorskip("torch") + from pepsy.backends import linalg_torch + + original_policy = linalg_torch._QR_RANK_POLICY + original_factor = linalg_torch._QR_RANK_TOL_FACTOR + matrix = torch.randn(4, 3, dtype=torch.float64) + matrix[:, 1] = matrix[:, 0] + matrix.requires_grad_() + try: + linalg_torch._configure_qr_rank_policy("native") + with warnings.catch_warnings(record=True) as caught: + q, r = linalg_torch.QR_real.apply(matrix) + gradient = torch.autograd.grad(q.sum() + r.sum(), matrix)[0] + assert not caught + assert torch.isfinite(gradient).all() + finally: + linalg_torch._QR_RANK_POLICY = original_policy + linalg_torch._QR_RANK_TOL_FACTOR = original_factor + + +def test_jax_linalg_registration_aliases_are_idempotent(): + """JAX real/relative compatibility aliases share one registration.""" + pytest.importorskip("jax") + from pepsy.backends import linalg_jax + + original_registered = linalg_jax._SVD_REGISTERED + original_function = linalg_jax._SVD_REGISTERED_FUNCTION + try: + linalg_jax._SVD_REGISTERED = False + linalg_jax._SVD_REGISTERED_FUNCTION = None + linalg_jax.reg_rel_svd_jax() + assert linalg_jax._SVD_REGISTERED is True + linalg_jax.reg_real_svd_jax() + finally: + linalg_jax._SVD_REGISTERED = original_registered + linalg_jax._SVD_REGISTERED_FUNCTION = original_function + + +def test_jax_linalg_registration_switches_native_and_stabilized(monkeypatch): + """JAX can explicitly switch between native and truncation-safe SVD.""" + pytest.importorskip("jax") + import autoray as ar + from pepsy.backends import linalg_jax + + calls = [] + original_registered = linalg_jax._SVD_REGISTERED + original_function = linalg_jax._SVD_REGISTERED_FUNCTION + linalg_jax._SVD_REGISTERED = False + linalg_jax._SVD_REGISTERED_FUNCTION = None + monkeypatch.setattr( + ar, + "register_function", + lambda *args, **kwargs: calls.append((args, kwargs)), + ) + try: + linalg_jax.reg_native_svd_jax() + linalg_jax.reg_rel_svd_jax() + finally: + linalg_jax._SVD_REGISTERED = original_registered + linalg_jax._SVD_REGISTERED_FUNCTION = original_function + + assert len(calls) == 2 + assert calls[0][0][0:2] == ("jax", "linalg.svd") + assert calls[0][0][2] is linalg_jax._native_svd_jax + assert calls[1][0][2] is linalg_jax.svd_jax + + +def test_register_jax_linalg_defaults_to_native(monkeypatch): + """The JAX umbrella registration defaults to native thin SVD.""" + pytest.importorskip("jax") + import autoray as ar + from pepsy.backends import config, linalg_jax + + calls = [] + original_registered = linalg_jax._SVD_REGISTERED + original_function = linalg_jax._SVD_REGISTERED_FUNCTION + linalg_jax._SVD_REGISTERED = False + linalg_jax._SVD_REGISTERED_FUNCTION = None + monkeypatch.setattr( + ar, + "register_function", + lambda *args, **kwargs: calls.append((args, kwargs)), + ) + try: + config.register_jax_linalg() + finally: + linalg_jax._SVD_REGISTERED = original_registered + linalg_jax._SVD_REGISTERED_FUNCTION = original_function + + assert len(calls) == 1 + assert calls[0][0][2] is linalg_jax._native_svd_jax + + +def test_reset_linalg_registrations_restores_native_jax(monkeypatch): + """The public reset helper restores native JAX thin SVD.""" + pytest.importorskip("jax") + import autoray as ar + from pepsy.backends import config, linalg_jax + + calls = [] + original_registered = linalg_jax._SVD_REGISTERED + original_function = linalg_jax._SVD_REGISTERED_FUNCTION + linalg_jax._SVD_REGISTERED = False + linalg_jax._SVD_REGISTERED_FUNCTION = None + monkeypatch.setattr( + ar, + "register_function", + lambda *args, **kwargs: calls.append((args, kwargs)), + ) + try: + config.reset_linalg_registrations(backend="jax") + finally: + linalg_jax._SVD_REGISTERED = original_registered + linalg_jax._SVD_REGISTERED_FUNCTION = original_function + + assert len(calls) == 1 + assert calls[0][0][2] is linalg_jax._native_svd_jax + + def test_to_float_rejects_non_scalar_backend_array_before_numpy_coercion(): class BackendVector: shape = (2,) diff --git a/tests/test_package_layout.py b/tests/test_package_layout.py index 530bea7..9454cc2 100644 --- a/tests/test_package_layout.py +++ b/tests/test_package_layout.py @@ -60,11 +60,15 @@ reg_complex_qr_torch, reg_complex_svd_jax, reg_complex_svd_torch, + reg_native_svd_jax, + reg_native_svd_torch, reg_real_qr_torch, reg_real_svd_jax, reg_real_svd_torch, reg_rel_svd_jax, reg_rel_svd_torch, + register_jax_linalg, + reset_linalg_registrations, site_charge_from_occupations, ) from pepsy.vmc import ( @@ -158,11 +162,15 @@ def test_new_namespace_imports_resolve(): assert callable(reg_rel_svd_torch) assert callable(reg_real_svd_torch) assert callable(reg_complex_svd_torch) + assert callable(reg_native_svd_jax) + assert callable(reg_native_svd_torch) assert callable(reg_real_qr_torch) assert callable(reg_complex_qr_torch) assert callable(reg_rel_svd_jax) assert callable(reg_real_svd_jax) assert callable(reg_complex_svd_jax) + assert callable(register_jax_linalg) + assert callable(reset_linalg_registrations) assert callable(site_charge_from_occupations) assert FermionSiteEncoding is not None assert ContractionConfig is not None diff --git a/tests/test_public_api.py b/tests/test_public_api.py index b1cfb7b..a31c897 100644 --- a/tests/test_public_api.py +++ b/tests/test_public_api.py @@ -84,9 +84,11 @@ def test_tree_optimizers_are_available_from_high_level_api(): "site_charge_alternating", "site_charge_from_map", "site_charge_from_occupations", "site_charge_uniform", "symmray_block_summary", "symmray_mps_summary", "symmray_mpo_summary", "symmray_peps_summary", "symm_operator_from_dense", - "reg_rel_svd_torch", "reg_real_svd_torch", "reg_complex_svd_torch", + "reg_native_svd_torch", "reg_native_svd_jax", "reg_rel_svd_torch", + "reg_real_svd_torch", "reg_complex_svd_torch", "reg_real_qr_torch", "reg_complex_qr_torch", "reg_rel_svd_jax", "reg_real_svd_jax", "reg_complex_svd_jax", + "register_jax_linalg", "reset_linalg_registrations", ] _EXPECTED_NOT_IN_ALL = [ @@ -151,9 +153,11 @@ def test_internal_symbols_not_exported(): "site_charge_alternating", "site_charge_from_map", "site_charge_from_occupations", "site_charge_uniform", "symmray_block_summary", "symmray_mps_summary", "symmray_mpo_summary", "symmray_peps_summary", "symm_operator_from_dense", - "reg_rel_svd_torch", "reg_real_svd_torch", "reg_complex_svd_torch", + "reg_native_svd_torch", "reg_native_svd_jax", "reg_rel_svd_torch", + "reg_real_svd_torch", "reg_complex_svd_torch", "reg_real_qr_torch", "reg_complex_qr_torch", "reg_rel_svd_jax", "reg_real_svd_jax", "reg_complex_svd_jax", + "register_jax_linalg", "reset_linalg_registrations", ] _BLOCKED_NAMES = _EXPECTED_NOT_IN_ALL @@ -215,16 +219,22 @@ def test_optional_linalg_registrations_resolve(): assert pepsy.reg_rel_svd_torch is pepsy.tensors.reg_rel_svd_torch assert pepsy.reg_real_svd_torch is pepsy.tensors.reg_real_svd_torch assert pepsy.reg_complex_svd_torch is pepsy.tensors.reg_complex_svd_torch + assert pepsy.reg_native_svd_torch is pepsy.tensors.reg_native_svd_torch + assert pepsy.reg_native_svd_jax is pepsy.tensors.reg_native_svd_jax assert pepsy.reg_real_qr_torch is pepsy.tensors.reg_real_qr_torch assert pepsy.reg_complex_qr_torch is pepsy.tensors.reg_complex_qr_torch assert pepsy.reg_rel_svd_jax is pepsy.tensors.reg_rel_svd_jax assert pepsy.reg_real_svd_jax is pepsy.tensors.reg_real_svd_jax assert pepsy.reg_complex_svd_jax is pepsy.tensors.reg_complex_svd_jax + assert pepsy.register_jax_linalg is pepsy.backends.register_jax_linalg + assert pepsy.reset_linalg_registrations is pepsy.backends.reset_linalg_registrations if has_torch: import torch assert callable(pepsy.tensors.core.reg_rel_svd_torch) assert callable(pepsy.tensors.reg_rel_svd_torch) + assert callable(pepsy.tensors.core.reg_native_svd_torch) + assert callable(pepsy.tensors.reg_native_svd_torch) assert callable(pepsy.tensors.core.reg_real_svd_torch) assert callable(pepsy.tensors.reg_real_svd_torch) assert callable(pepsy.tensors.core.reg_complex_svd_torch) @@ -245,3 +255,7 @@ def test_optional_linalg_registrations_resolve(): assert callable(pepsy.tensors.reg_real_svd_jax) assert callable(pepsy.tensors.core.reg_complex_svd_jax) assert callable(pepsy.tensors.reg_complex_svd_jax) + assert callable(pepsy.tensors.core.reg_native_svd_jax) + assert callable(pepsy.tensors.reg_native_svd_jax) + assert callable(pepsy.tensors.core.register_jax_linalg) + assert callable(pepsy.tensors.register_jax_linalg) From 116e8890863458507b2276d69516de06d00f2396 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Thu, 30 Jul 2026 12:00:54 -0600 Subject: [PATCH 29/70] Add edge-resolved local BP loop series --- .github/skills/belief-propagation/SKILL.md | 49 + .../references/belief_propagation.md | 3 + src/pepsy/bp/__init__.py | 16 + src/pepsy/bp/series.py | 2126 +++++++++++++++-- 4 files changed, 1954 insertions(+), 240 deletions(-) diff --git a/.github/skills/belief-propagation/SKILL.md b/.github/skills/belief-propagation/SKILL.md index e7793d3..f07b8d8 100644 --- a/.github/skills/belief-propagation/SKILL.md +++ b/.github/skills/belief-propagation/SKILL.md @@ -50,6 +50,23 @@ over `quimb.tensor.belief_propagation`; the annotated paper trail lives in - `loop_series_expand(tn, gloops, *, norm="2norm"|"1norm", ...)` — edge-resolved loop **series** with explicit excited-bond terms. Its integer cutoff is a maximum excited-bond degree, not a tensor-region size. +- Local D2BP observables: `partial_trace_loop_series_expand` and + `partial_trace_loop_cluster_expand` return local reduced density matrices; + `compute_local_expectation_loop_series` and + `compute_local_expectation_loop_cluster` accept Quimb-style + ``{site_or_sites: operator}`` mappings and return scalar expectations. Share + one D2BP solve across terms. Use the scalar APIs for fermionic observables. +- Explicit-edge local D2BP observables: + `partial_trace_edge_loop_series_expand` and + `compute_local_expectation_edge_loop_series` use the same canonical + `LoopSeriesTerm` Q-edge sets and edge-degree cutoff as + `loop_series_expand`, rather than Quimb's local-region cutoff. The scalar + path inserts the gate before forming the graded bra and is the supported + fermionic expectation route. Terms that put Q on a bond wholly internal to + `where` are currently rejected explicitly. Nonzero-Q fermionic scalar + corrections are currently restricted to one-site gates; reject multi-site + graded-Q gates clearly rather than treating a dense rho trace as a sign + oracle. - `partitioned_expand(tn, partition_inds=... | partitions=..., *, norm="2norm"|"1norm", form="linear"|"combinatorial", ...)` — PNE from Evenbly, Gray, and Chan (arXiv:2512.10910). It inserts complementary @@ -90,6 +107,16 @@ another: | `cluster` | tensor regions with counting numbers | `gloops` = tensor-region size or explicit regions | | `pne` | selected index partitions into `P` and `Q` subspaces | `partition_inds` or factorized `partitions` | +The original *local-RDM* APIs deliberately follow Quimb's `get_local_gloops` region +convention: their integer `gloops` is a local generalized-loop region cutoff, +not the edge-degree cutoff of `loop_series_expand`. Local loop series inserts +`Q` on the eligible internal bonds of each selected region; local loop cluster +contracts BP-closed regions with inclusion--exclusion counts. Do not compare +either term-for-term with a brute-force edge-subset enumeration. For that +case, use `partial_trace_edge_loop_series_expand` or +`compute_local_expectation_edge_loop_series`; do not silently reinterpret a +region cutoff. + PNE terms are not loop-cluster regions. A PNE residue is the all-`Q` network; retaining it gives the exact projector identity, while dropping it is the approximation whose error should be monitored. @@ -112,6 +139,28 @@ not treat `optimize="auto-hq"` as a fermionic sign fix: it is a Cotengra path preset, just like a reusable `PathOptimizer`, and a path-dependent sign means the input metadata is invalid or incomplete. +### Fermionic local observables: required construction and oracle + +Do **not** evaluate a native fermionic observable as `trace(rho @ gate)`, even +when a partial-trace API returns a native Symmray rho: opening and fusing the +physical legs loses the graded gate-routing convention. Treat those rhos as +diagnostics (charge support, trace, Hermiticity), not as an observable oracle. + +For scalar local observables, insert the native gate into the ket with +`tensor_network_gate_inds(..., contract=False)` *before* forming the double +layer. Build the bra from D2BP's `tensor_dual_map` and attach messages using +`index_dual_map`; do not replace this with `tensor.conj()` carrying ket virtual +labels. Normalize with the directly contracted gate-free BP network, not a +trace of the fused rho. This works for both adjacent and separated supports. + +Before trusting a new fermionic BP observable path, add a native-Symmray tree +test where D2BP is exact: compare a parity-even two-site gate (hopping, +pairing, or eta-pairing) against `compute_local_expectation_exact`, then repeat +with reversed site order and a correspondingly transposed native gate. A dense +or trace-only test does not establish fermionic sign correctness. Use a +brute-force edge-subset implementation only as a loop-enumeration oracle; if +it is dense/non-graded, it is not by itself a fermionic-sign oracle. + For valid closed scalar Symmray networks, the 1-norm APIs (`L1BP`, `HV1BP`, and `D1BP`), their loop/series/PNE corrections, D1 SU bridge, and `weight_pass` use a topology-identical dense BP shadow because Quimb's D1 diff --git a/docs/development/references/belief_propagation.md b/docs/development/references/belief_propagation.md index 15d7fa4..6edf4f1 100644 --- a/docs/development/references/belief_propagation.md +++ b/docs/development/references/belief_propagation.md @@ -21,6 +21,9 @@ side · **[roots]** foundational / prior art. | `d2bp_from_simple_update_gauges`, `run_d2bp_from_simple_update_gauges`, `simple_update_core_and_gauges_from_d2bp`, `gauge_all(norm="2norm")` | physical PEPS simple-update/Vidal ↔ D2BP density-message bridge: `diag(lambda)` initialization and PSD metric gauge conversion | Tindall & Fishman (2023); Gray et al. 2510.05647 | | `loop_cluster_expand`, `LoopClusterResult` | loop **cluster** expansion (thin wrapper over quimb `contract_gloop_expand`) | Gray et al. 2510.05647 | | `loop_series_expand`, `LoopSeriesTerm`, `LoopSeriesCache` | edge-resolved `P + Q` loop series for D1BP and D2BP; retains excited-bond degree and distinct embeddings/chord subsets | Evenbly et al. 2409.03108 | +| `partial_trace_loop_series_expand`, `compute_local_expectation_loop_series` | D2BP local reduced-density-matrix and scalar `P + Q` loop series; keeps physical output legs open and uses native Symmray virtual projectors | Evenbly et al. 2409.03108; quimb local loop-series API | +| `partial_trace_edge_loop_series_expand`, `compute_local_expectation_edge_loop_series` | D2BP local RDM and graded scalar observable expansion over canonical explicit Q-edge terms; does not reinterpret Quimb's local-region cutoff | Evenbly et al. 2409.03108; Pepsy API | +| `partial_trace_loop_cluster_expand`, `compute_local_expectation_loop_cluster` | D2BP local reduced-density-matrix and scalar generalized-loop cluster expansion; combines BP-closed regions with inclusion--exclusion counts | Gray et al. 2510.05647; quimb local cluster API | | `loop_expand` | explicit selector between the edge loop series and region loop-cluster expansion; preserves each method's cutoff and result metadata | Pepsy API | | `partitioned_expand`, `pne_expand`, `PNEExpansionResult` | linear and combinatorial partitioned network expansions for D1BP/D2BP, with optional residue, explicit projectors, open outputs, and fixed recursive schedules | Evenbly, Gray & Chan 2512.10910 | | `weight_pass`, `WeightPassingResult` | Appendix-C positive-weight passing on closed pairwise networks and rank-`r` projectors in the returned gauge | Evenbly, Gray & Chan 2512.10910 | diff --git a/src/pepsy/bp/__init__.py b/src/pepsy/bp/__init__.py index 40b3014..98b0185 100644 --- a/src/pepsy/bp/__init__.py +++ b/src/pepsy/bp/__init__.py @@ -10,6 +10,10 @@ * :func:`loop_cluster_expand` -- the loop cluster expansion (arXiv:2510.05647), * :func:`loop_series_expand` -- the edge-resolved ``P + Q`` loop series (arXiv:2409.03108), +* :func:`partial_trace_loop_cluster_expand` -- a D2BP loop-cluster reduced + density matrix and its scalar expectation companion, +* :func:`partial_trace_loop_series_expand` -- a D2BP reduced-density-matrix + loop series and scalar local-observable companion, * :func:`loop_expand` -- an explicit selector between the correction families. * :func:`partitioned_expand` -- the partitioned network expansion (PNE, arXiv:2512.10910), with :func:`recursive_partitioned_expand` for fixed @@ -38,6 +42,12 @@ LoopSeriesCache, LoopSeriesResult, LoopSeriesTerm, + compute_local_expectation_edge_loop_series, + compute_local_expectation_loop_cluster, + compute_local_expectation_loop_series, + partial_trace_edge_loop_series_expand, + partial_trace_loop_cluster_expand, + partial_trace_loop_series_expand, loop_series_expand, ) from .expansion import loop_expand @@ -116,6 +126,12 @@ "LoopSeriesCache", "LoopSeriesResult", "LoopSeriesTerm", + "compute_local_expectation_edge_loop_series", + "compute_local_expectation_loop_cluster", + "compute_local_expectation_loop_series", + "partial_trace_edge_loop_series_expand", + "partial_trace_loop_cluster_expand", + "partial_trace_loop_series_expand", "RelayGaugeOptions", "ReducedALSSolution", "ReducedBondPair", diff --git a/src/pepsy/bp/series.py b/src/pepsy/bp/series.py index 1dc8123..cb69730 100644 --- a/src/pepsy/bp/series.py +++ b/src/pepsy/bp/series.py @@ -23,7 +23,9 @@ from collections import deque from dataclasses import dataclass, field +import functools from itertools import combinations +import operator from typing import Any, ClassVar import autoray as ar @@ -39,8 +41,10 @@ dense_bp_tn as _dense_bp_tn, dense_message_tree as _dense_message_tree, d2_operator as _symmray_d2_operator, + is_symmray_array as _is_symmray_array, rank_one_d2_projector as _symmray_rank_one_d2_projector, restore_fermionic_dummy_modes as _restore_fermionic_dummy_modes, + to_dense as _symmray_to_dense, uses_symmray as _uses_symmray, ) @@ -48,6 +52,12 @@ "LoopSeriesCache", "LoopSeriesResult", "LoopSeriesTerm", + "compute_local_expectation_edge_loop_series", + "compute_local_expectation_loop_cluster", + "partial_trace_loop_cluster_expand", + "partial_trace_edge_loop_series_expand", + "compute_local_expectation_loop_series", + "partial_trace_loop_series_expand", "loop_series_expand", ] @@ -522,135 +532,1135 @@ def _get_edge_excited(bp, term): return _get_d2_edge_excited(bp, term) -def _process_weights( - weights, +def _get_d2_partial_trace_excited( + bp, + tids, *, - mantissa, - exponent, - multi_excitation_correct, - tol_correction, - maxiter_correction, - strip_exponent, + partial_trace_map, + exclude=(), + gate=None, + gate_inds=(), ): - """Use Quimb's loop-series resummation with edge-degree keys.""" - from quimb.tensor.belief_propagation.bp_common import ( - process_loop_series_expansion_weights, - ) + """Build a D2 excited cluster with selected physical legs left open. - suppression = process_loop_series_expansion_weights( - weights, - multi_excitation_correct=multi_excitation_correct, - tol_correction=tol_correction, - maxiter_correction=maxiter_correction, - return_all=True, - ) - correction = -sum( - weight * suppression[edges] for edges, weight in weights.items() - ) - estimate = process_loop_series_expansion_weights( - weights, - mantissa=mantissa, - exponent=exponent, - multi_excitation_correct=multi_excitation_correct, - tol_correction=tol_correction, - maxiter_correction=maxiter_correction, - strip_exponent=strip_exponent, - ) - return estimate, correction, suppression + This is the partial-trace counterpart of :func:`_get_d2_edge_excited`. + It follows Quimb's ``D2BP.get_cluster_excited`` convention, but builds + native Symmray virtual projectors when the BP network is block sparse. + ``exclude`` contains bonds inside the base observable region; those bonds + are traced normally rather than receiving a loop excitation projector. + """ + import quimb.tensor as qtn + stn = bp.tn._select_tids(tids) + exclude = set(exclude) + kixmaps = {tid: {} for tid in stn.tensor_map} + bixmaps = {tid: {} for tid in stn.tensor_map} + excitation_inds = {} + boundary_inds = [] -def _contract_loop_series( + for index, region_tids in stn.ind_map.items(): + region_tids = tuple(region_tids) + if index in bp.output_inds: + if index in partial_trace_map: + (tid,) = region_tids + bixmaps[tid][index] = partial_trace_map[index] + elif index in exclude: + # Trace this bond in the bra layer without inserting P or Q. + bix = qtn.rand_uuid() + for tid in region_tids: + bixmaps[tid][index] = bix + elif index in stn._inner_inds: + for tid in region_tids: + kix = qtn.rand_uuid() + bix = qtn.rand_uuid() + kixmaps[tid][index] = kix + bixmaps[tid][index] = bix + excitation_inds.setdefault(index, {})[tid] = (bix, kix) + else: + (tid,) = region_tids + kix = qtn.rand_uuid() + bix = qtn.rand_uuid() + kixmaps[tid][index] = kix + bixmaps[tid][index] = bix + boundary_inds.append((index, tid)) + + local = qtn.TensorNetwork() + if gate is None: + ket_tn = stn + else: + # Fermionic gates must be inserted into the ket before the bra is + # formed. Contracting a gate with an already-open fermionic rho misses + # the graded swaps associated with the physical legs. + ket_tn = qtn.tensor_network_gate_inds( + stn, + gate, + gate_inds, + contract=False, + tags=[], + info=None, + inplace=False, + ) + + for tid in stn.tensor_map: + tensor = ket_tn.tensor_map[tid] + local |= tensor.reindex(kixmaps[tid]) + # D2BP owns the graded bra tensors and their virtual dual-index map. + # ``tensor.conj()`` has the same numerical blocks but retains the ket + # virtual labels; that misses the fermionic bra ordering when boundary + # messages are attached. + bra_reindex = { + bp.index_dual_map.get(index, index): new_index + for index, new_index in bixmaps[tid].items() + } + local |= bp.tensor_dual_map[tid].reindex(bra_reindex) + + # ``contract=False`` works for both adjacent and separated supports. The + # added gate tensor carries the graded physical routing between the ket and + # bra; original site tensors retain their ids above. + for tid, tensor in ket_tn.tensor_map.items(): + if tid not in stn.tensor_map: + local |= tensor + + for index, tid in boundary_inds: + data = bp.messages[index, tid] + local |= qtn.Tensor( + data, + inds=(bixmaps[tid][index], kixmaps[tid][index]), + ) + + for index, region_tids in excitation_inds.items(): + tid_left, tid_right = tuple(region_tids) + left = excitation_inds[index][tid_left] + right = excitation_inds[index][tid_right] + ml = bp.messages[index, tid_left] + mr = bp.messages[index, tid_right] + + if _uses_symmray(bp.tn): + projector = _symmray_d2_operator( + bp.tn, + index, + _symmray_rank_one_d2_projector( + bp.tn, index, ml, mr, layout="series" + ), + complement=True, + layout="series", + ) + else: + vacuum = ar.do( + "einsum", + "i,j->ij", + ml.reshape(-1), + mr.reshape(-1), + ) + projector = ar.do("eye", ar.do("shape", vacuum)[0]) - vacuum + projector = ar.do( + "reshape", + projector, + ar.do("shape", ml) + ar.do("shape", mr), + ) + + local |= qtn.Tensor( + projector, + inds=(*left, *right), + ) + + return local + + +def _get_d2_edge_partial_trace_excited( bp, - terms, + tids, + *, + excited_edges=(), + partial_trace_map=(), + exclude=(), + gate=None, + gate_inds=(), +): + """Build a D2 local RDM network with explicit P/Q edge choices. + + Every internal virtual bond of ``tids`` receives ``P`` except the bonds + named by ``excited_edges``, which receive ``Q = I - P``. Bonds in + ``exclude`` are traced directly. This is deliberately separate from + :func:`_get_d2_partial_trace_excited`, whose non-excluded bonds are all + ``Q`` for Quimb's local-region convention. + """ + import quimb.tensor as qtn + + stn = bp.tn._select_tids(tids) + excited_edges = set(excited_edges) + exclude = set(exclude) + kixmaps = {tid: {} for tid in stn.tensor_map} + bixmaps = {tid: {} for tid in stn.tensor_map} + projector_inds = {} + boundary_inds = [] + + for index, region_tids in stn.ind_map.items(): + region_tids = tuple(region_tids) + if index in bp.output_inds: + if index in partial_trace_map: + (tid,) = region_tids + bixmaps[tid][index] = partial_trace_map[index] + elif index in exclude: + bix = qtn.rand_uuid() + for tid in region_tids: + bixmaps[tid][index] = bix + elif index in stn._inner_inds: + for tid in region_tids: + kix = qtn.rand_uuid() + bix = qtn.rand_uuid() + kixmaps[tid][index] = kix + bixmaps[tid][index] = bix + projector_inds.setdefault(index, {})[tid] = (bix, kix) + else: + (tid,) = region_tids + kix = qtn.rand_uuid() + bix = qtn.rand_uuid() + kixmaps[tid][index] = kix + bixmaps[tid][index] = bix + boundary_inds.append((index, tid)) + + local = qtn.TensorNetwork() + if gate is None: + ket_tn = stn + else: + ket_tn = qtn.tensor_network_gate_inds( + stn, + gate, + gate_inds, + contract=False, + tags=[], + info=None, + inplace=False, + ) + + for tid in stn.tensor_map: + local |= ket_tn.tensor_map[tid].reindex(kixmaps[tid]) + bra_reindex = { + bp.index_dual_map.get(index, index): new_index + for index, new_index in bixmaps[tid].items() + } + local |= bp.tensor_dual_map[tid].reindex(bra_reindex) + for tid, tensor in ket_tn.tensor_map.items(): + if tid not in stn.tensor_map: + local |= tensor + + for index, tid in boundary_inds: + local |= qtn.Tensor( + bp.messages[index, tid], + inds=(bixmaps[tid][index], kixmaps[tid][index]), + ) + + for index, region_tids in projector_inds.items(): + tid_left, tid_right = tuple(region_tids) + left = projector_inds[index][tid_left] + right = projector_inds[index][tid_right] + ml = bp.messages[index, tid_left] + mr = bp.messages[index, tid_right] + if _uses_symmray(bp.tn): + p0 = _symmray_rank_one_d2_projector( + bp.tn, index, ml, mr, layout="series" + ) + projector = _symmray_d2_operator( + bp.tn, + index, + p0, + complement=index in excited_edges, + layout="series", + ) + else: + p0 = ar.do("einsum", "i,j->ij", ml.reshape(-1), mr.reshape(-1)) + projector = ( + ar.do("eye", ar.do("shape", p0)[0]) - p0 + if index in excited_edges + else p0 + ) + projector = ar.do( + "reshape", + projector, + ar.do("shape", ml) + ar.do("shape", mr), + ) + local |= qtn.Tensor(projector, inds=(*left, *right)) + + return local + + +def _rho_trace(rho): + """Trace a reduced density matrix, including omitted Symmray blocks.""" + if _is_symmray_array(rho): + return np.trace(_symmray_to_dense(rho)) + return ar.do("trace", rho) + + +def _term_sites(tn, where): + """Normalize a Quimb local-term key to an ordered site tuple.""" + has_site = getattr(tn, "has_site", None) + if callable(has_site) and has_site(where): + return (where,) + if isinstance(where, (str, bytes)): + return (where,) + try: + sites = tuple(where) + except TypeError: + return (where,) + if not sites: + raise ValueError("a local expectation term must have at least one site") + return sites + + +def _partial_trace_loop_series( + bp, + where, + gloops, *, + normalized, + grow_from, + strict_size, multi_excitation_correct, - tol_correction, - maxiter_correction, - strip_exponent, optimize, contract_opts, - normalize, + info, ): - if normalize: - if bp.__class__.__name__ == "D2BP": - _align_symmray_d2bp_messages(bp) - bp.normalize_message_pairs() - bp.normalize_tensors() + """Contract the native D2 local loop-series density matrices.""" + if bp.__class__.__name__ != "D2BP": + raise ValueError( + "partial_trace_loop_series_expand currently requires norm='2norm'" + ) + if normalized == "prod": + normalized = True + if normalized not in (True, False, "local", "separate"): + raise ValueError( + "normalized must be one of True, False, 'prod', 'local', or " + "'separate'" + ) - weights = {} - for term in terms: - weights[term.edges] = _get_edge_excited(bp, term).contract( - optimize=optimize, **contract_opts + _align_symmray_d2bp_messages(bp) + bp.normalize_message_pairs() + bp.normalize_tensors() + + tags = [bp.tn.site_tag(coo) for coo in where] + tids = frozenset(bp.tn._get_tids_from_tags(tags, "any")) + if not tids: + raise ValueError("where must contain at least one site in the network") + + kix = [bp.tn.site_ind(coo) for coo in where] + import quimb.tensor as qtn + + bix = [qtn.rand_uuid() for _ in where] + partial_trace_map = dict(zip(kix, bix)) + output_inds = (*kix, *bix) + + regions = tuple( + frozenset(region) + for region in bp.tn.get_local_gloops( + tids=tids, + gloops=gloops, + grow_from=grow_from, + strict_size=strict_size, ) - estimate, correction, suppression = _process_weights( - weights, - mantissa=bp.sign, - exponent=bp.exponent, - multi_excitation_correct=multi_excitation_correct, - tol_correction=tol_correction, - maxiter_correction=maxiter_correction, - strip_exponent=strip_exponent, ) - return estimate, weights, correction, suppression + base_region = frozenset(tids) + if base_region not in regions: + regions = (base_region, *regions) + # Preserve the first occurrence while making the term keys hashable. + regions = tuple(dict.fromkeys(regions)) + + inner_bonds = bp.tn._select_tids(tids).inner_inds() + term_cache = {} if info is None else info.setdefault("rho_terms", {}) + rho_terms = {} + for region in regions: + cache_key = (region, tuple(where)) + try: + rho_e = term_cache[cache_key] + except KeyError: + excited = _get_d2_partial_trace_excited( + bp, + region, + partial_trace_map=partial_trace_map, + exclude=inner_bonds, + ) + rho_e = excited.contract( + output_inds=output_inds, + optimize=optimize, + **contract_opts, + ).to_dense(kix, bix) + term_cache[cache_key] = rho_e + + if normalized == "local" and region != base_region: + rho_e = rho_e / (1 + _rho_trace(rho_e)) + rho_terms[region] = rho_e + + weights = { + region: _rho_trace(rho_e) for region, rho_e in rho_terms.items() + } + if multi_excitation_correct: + correction_weights = { + region: weight + for region, weight in weights.items() + if region != base_region + } + if correction_weights: + from quimb.tensor.belief_propagation.bp_common import ( + process_loop_series_expansion_weights, + ) + + suppression = process_loop_series_expansion_weights( + correction_weights, + return_all=True, + ) + else: + suppression = {} + else: + suppression = {} + suppression[base_region] = 1.0 + for region in regions: + suppression.setdefault(region, 1.0) + + rho = functools.reduce( + operator.add, + (rho_terms[region] * suppression[region] for region in regions), + ) + if normalized: + rho = rho / _rho_trace(rho) + elif (bp.sign, bp.exponent) != (1.0, 0.0): + rho = rho * bp.sign * 10**bp.exponent + if info is not None: + info["rho_weights"] = weights + info["rho_suppression_factors"] = suppression + info["rho_regions"] = regions + return rho -def _build_bp( - tn, + +def _partial_trace_loop_cluster( + bp, + where, + gloops, *, - norm, - messages, - gauges, - run_bp, - bp_runner, - relay_opts, - max_iterations, - tol, - tol_abs, - tol_rolling_diff, - diis, - damping, - update, + combine, + normalized, + autocomplete, + grow_from, + strict_size, optimize, - bp_opts, - progbar, - validate_graph=True, + contract_opts, + info, ): - key, bp_cls = _cluster_bp_class(norm) - if key == "2norm" and _uses_symmray(tn): - tn = _restore_fermionic_dummy_modes(tn) - if key == "1norm" and _uses_symmray(tn): - tn = _dense_bp_tn(tn) - messages = _dense_message_tree(messages) - gauges = _dense_message_tree(gauges) - if bp_runner not in {"plain", "relay"}: - raise ValueError("bp_runner must be either 'plain' or 'relay'") - if gauges is not None and messages is not None: - raise ValueError("pass either messages or gauges, not both") + """Contract native D2BP generalized-loop cluster density matrices.""" + if bp.__class__.__name__ != "D2BP": + raise ValueError( + "partial_trace_loop_cluster_expand currently requires norm='2norm'" + ) + if combine not in {"sum", "prod"}: + raise ValueError("combine must be 'sum' or 'prod'") + if normalized == "prod": + normalized = True + if normalized is True: + normalized = "local" + if normalized not in (False, "local", "separate"): + raise ValueError( + "normalized must be one of True, False, 'prod', 'local', or " + "'separate'" + ) - if key == "1norm" and validate_graph: - from .gauges import _validate_d1_graph + _align_symmray_d2bp_messages(bp) - _validate_d1_graph(tn) - else: - from .gauges import _validate_d2_graph + tags = [bp.tn.site_tag(coo) for coo in where] + tids = frozenset(bp.tn._get_tids_from_tags(tags, "any")) + if not tids: + raise ValueError("where must contain at least one site in the network") - _validate_d2_graph(tn) + kix = [bp.tn.site_ind(coo) for coo in where] + import quimb.tensor as qtn - info = {} - if gauges is not None: - from .gauges import ( - d1bp_from_simple_update_gauges, - d2bp_from_simple_update_gauges, + bix = [qtn.rand_uuid() for _ in where] + partial_trace_map = dict(zip(kix, bix)) + output_inds = (*kix, *bix) + regions = tuple( + bp.tn.get_local_gloops( + tids=tids, + gloops=gloops, + grow_from=grow_from, + strict_size=strict_size, ) + ) - gauge_builder = ( - d1bp_from_simple_update_gauges - if key == "1norm" - else d2bp_from_simple_update_gauges + from quimb.tensor.belief_propagation import gen_region_counts + + term_cache = {} if info is None else info.setdefault("cluster_rho_terms", {}) + rhos = [] + counts = [] + for region, count in gen_region_counts(regions, autocomplete=autocomplete): + region = frozenset(region) + cache_key = (region, tuple(where)) + try: + rho_r = term_cache[cache_key] + except KeyError: + cluster = bp.get_cluster_norm( + region, + partial_trace_map=partial_trace_map, + ) + rho_r = cluster.contract( + output_inds=output_inds, + optimize=optimize, + **contract_opts, + ).to_dense(kix, bix) + term_cache[cache_key] = rho_r + + if normalized == "local": + rho_r = rho_r / _rho_trace(rho_r) + rhos.append(rho_r) + counts.append(count) + + if not rhos: + raise ValueError("no generalized-loop cluster regions were generated") + if combine == "sum": + rho = functools.reduce( + operator.add, + (count * rho_r for rho_r, count in zip(rhos, counts)), ) - gauge_opts = dict(bp_opts) - if key == "2norm": - gauge_opts.setdefault("optimize", optimize) + else: + rho = functools.reduce( + operator.mul, + (rho_r**count for rho_r, count in zip(rhos, counts)), + ) + + if normalized == "separate" or (normalized and combine == "prod"): + rho = rho / _rho_trace(rho) + elif not normalized and (bp.sign, bp.exponent) != (1.0, 0.0): + rho = rho * bp.sign * 10**bp.exponent + + if info is not None: + info["cluster_rho_regions"] = tuple( + (region, count) + for region, count in gen_region_counts( + regions, autocomplete=autocomplete + ) + ) + return rho + + +def _get_d2_cluster_norm( + bp, + tids, + *, + partial_trace_map=(), + gate=None, + gate_inds=(), +): + """Build a D2BP message-closed cluster, optionally with a graded gate.""" + import quimb.tensor as qtn + + ket_base = bp.tn._select_tids(tids, virtual=False) + ket = ket_base + if gate is not None: + ket = qtn.tensor_network_gate_inds( + ket, + gate, + gate_inds, + contract=False, + tags=[], + info=None, + inplace=False, + ) + bra = qtn.TensorNetwork(bp.tensor_dual_map[tid] for tid in tids) + if partial_trace_map: + bra.reindex_(partial_trace_map) + cluster = bra | ket + for index in ket_base.outer_inds(): + if index in partial_trace_map or index in bp.output_inds: + continue + (tid,) = ket_base.ind_map[index] + cluster |= qtn.Tensor( + bp.messages[index, tid], + inds=(bp.index_dual_map[index], index), + ) + return cluster + + +def _local_expectation_loop_series( + bp, + where, + gate, + gloops, + *, + normalized, + grow_from, + strict_size, + multi_excitation_correct, + optimize, + contract_opts, + info, +): + """Contract one gate through the graded D2BP loop-series network.""" + if normalized == "prod": + normalized = True + if normalized not in (True, False, "local", "separate"): + raise ValueError( + "normalized must be one of True, False, 'prod', 'local', or " + "'separate'" + ) + + _align_symmray_d2bp_messages(bp) + bp.normalize_message_pairs() + bp.normalize_tensors() + + tags = [bp.tn.site_tag(coo) for coo in where] + tids = frozenset(bp.tn._get_tids_from_tags(tags, "any")) + kix = [bp.tn.site_ind(coo) for coo in where] + regions = tuple( + frozenset(region) + for region in bp.tn.get_local_gloops( + tids=tids, + gloops=gloops, + grow_from=grow_from, + strict_size=strict_size, + ) + ) + base_region = frozenset(tids) + if base_region not in regions: + regions = (base_region, *regions) + regions = tuple(dict.fromkeys(regions)) + + inner_bonds = bp.tn._select_tids(tids).inner_inds() + term_cache = ( + {} if info is None else info.setdefault("series_norm_terms", {}) + ) + norm_terms = {} + gate_terms = {} + for region in regions: + cache_key = (region, tuple(where)) + try: + norm_e = term_cache[cache_key] + except KeyError: + norm_tn = _get_d2_partial_trace_excited( + bp, + region, + partial_trace_map={}, + exclude=inner_bonds, + ) + norm_e = norm_tn.contract( + optimize=optimize, + **contract_opts, + ) + term_cache[cache_key] = norm_e + + gated = _get_d2_partial_trace_excited( + bp, + region, + partial_trace_map={}, + exclude=inner_bonds, + gate=gate, + gate_inds=kix, + ) + gate_e = gated.contract(optimize=optimize, **contract_opts) + if normalized == "local" and region != base_region: + scale = 1 + norm_e + norm_e = norm_e / scale + gate_e = gate_e / scale + norm_terms[region] = norm_e + gate_terms[region] = gate_e + + weights = { + region: norm_e for region, norm_e in norm_terms.items() + } + if multi_excitation_correct: + correction_weights = { + region: weight + for region, weight in weights.items() + if region != base_region + } + if correction_weights: + from quimb.tensor.belief_propagation.bp_common import ( + process_loop_series_expansion_weights, + ) + + suppression = process_loop_series_expansion_weights( + correction_weights, + return_all=True, + ) + else: + suppression = {} + else: + suppression = {} + suppression[base_region] = 1.0 + for region in regions: + suppression.setdefault(region, 1.0) + + norm = functools.reduce( + operator.add, + (norm_terms[region] * suppression[region] for region in regions), + ) + value = functools.reduce( + operator.add, + (gate_terms[region] * suppression[region] for region in regions), + ) + if normalized: + value = value / norm + elif (bp.sign, bp.exponent) != (1.0, 0.0): + value = value * bp.sign * 10**bp.exponent + + return value, norm + + +def _edge_series_terms_for_support(bp, tids, gloops, *, cache): + """Parse canonical edge terms and validate the local observable support.""" + terms = _parse_gloops(bp.tn, gloops, cache=cache) + inner_bonds = frozenset(bp.tn._select_tids(tids).inner_inds()) + crossing = [ + edge + for term in terms + for edge in term.edges + if edge in inner_bonds + ] + if crossing: + raise ValueError( + "explicit edge loop-series terms cannot currently excite a " + "virtual bond internal to the observable support; choose a " + "support without that bond or use the local-region API" + ) + return terms, inner_bonds + + +def _edge_series_suppression( + weights, + *, + multi_excitation_correct, + tol_correction, + maxiter_correction, +): + if not multi_excitation_correct or not weights: + return {edges: 1.0 for edges in weights} + from quimb.tensor.belief_propagation.bp_common import ( + process_loop_series_expansion_weights, + ) + + return process_loop_series_expansion_weights( + weights, + multi_excitation_correct=True, + tol_correction=tol_correction, + maxiter_correction=maxiter_correction, + return_all=True, + ) + + +def _partial_trace_edge_loop_series( + bp, + where, + gloops, + *, + normalized, + multi_excitation_correct, + tol_correction, + maxiter_correction, + optimize, + contract_opts, + cache, + info, +): + """Explicit edge-subset P/Q expansion for a local density matrix.""" + if normalized == "prod": + normalized = True + if normalized not in (True, False, "local", "separate"): + raise ValueError( + "normalized must be one of True, False, 'prod', 'local', or " + "'separate'" + ) + _align_symmray_d2bp_messages(bp) + bp.normalize_message_pairs() + bp.normalize_tensors() + + tags = [bp.tn.site_tag(coo) for coo in where] + tids = frozenset(bp.tn._get_tids_from_tags(tags, "any")) + if not tids: + raise ValueError("where must contain at least one site in the network") + terms, inner_bonds = _edge_series_terms_for_support( + bp, tids, gloops, cache=cache + ) + + import quimb.tensor as qtn + + kix = [bp.tn.site_ind(coo) for coo in where] + bix = [qtn.rand_uuid() for _ in where] + partial_trace_map = dict(zip(kix, bix)) + output_inds = (*kix, *bix) + term_cache = {} if info is None else info.setdefault("edge_rho_terms", {}) + rho_terms = {} + + for term in terms: + region = frozenset((*tids, *term.tids)) + cache_key = (term.edges, region, tuple(where)) + try: + rho_e = term_cache[cache_key] + except KeyError: + rho_e = _get_d2_edge_partial_trace_excited( + bp, + region, + excited_edges=term.edges, + partial_trace_map=partial_trace_map, + exclude=inner_bonds, + ).contract( + output_inds=output_inds, + optimize=optimize, + **contract_opts, + ).to_dense(kix, bix) + term_cache[cache_key] = rho_e + rho_terms[term.edges] = rho_e + + base = _get_d2_edge_partial_trace_excited( + bp, + tids, + partial_trace_map=partial_trace_map, + exclude=inner_bonds, + ).contract( + output_inds=output_inds, + optimize=optimize, + **contract_opts, + ).to_dense(kix, bix) + + weights = {edges: _rho_trace(rho) for edges, rho in rho_terms.items()} + suppression = _edge_series_suppression( + weights, + multi_excitation_correct=multi_excitation_correct, + tol_correction=tol_correction, + maxiter_correction=maxiter_correction, + ) + if normalized == "local": + rho_terms = { + edges: rho / (1 + weights[edges]) + for edges, rho in rho_terms.items() + } + rho = base + for edges, rho_e in rho_terms.items(): + rho = rho + rho_e * suppression[edges] + if normalized: + rho = rho / _rho_trace(rho) + elif (bp.sign, bp.exponent) != (1.0, 0.0): + rho = rho * bp.sign * 10**bp.exponent + if info is not None: + info["edge_rho_weights"] = weights + info["edge_rho_suppression_factors"] = suppression + info["edge_rho_terms"] = terms + return rho + + +def _local_expectation_edge_loop_series( + bp, + where, + gate, + gloops, + *, + normalized, + multi_excitation_correct, + tol_correction, + maxiter_correction, + optimize, + contract_opts, + cache, + info, +): + """Direct graded scalar counterpart of the explicit edge RDM series.""" + if normalized == "prod": + normalized = True + if normalized not in (True, False, "local", "separate"): + raise ValueError( + "normalized must be one of True, False, 'prod', 'local', or " + "'separate'" + ) + _align_symmray_d2bp_messages(bp) + bp.normalize_message_pairs() + bp.normalize_tensors() + + tags = [bp.tn.site_tag(coo) for coo in where] + tids = frozenset(bp.tn._get_tids_from_tags(tags, "any")) + terms, inner_bonds = _edge_series_terms_for_support( + bp, tids, gloops, cache=cache + ) + if _uses_symmray(bp.tn) and len(where) > 1 and terms: + raise NotImplementedError( + "fermionic explicit-edge loop corrections for multi-site gates " + "are not supported yet; use gloops=0, a one-site term, or a " + "separate exact/path observable route" + ) + kix = [bp.tn.site_ind(coo) for coo in where] + norm_cache = {} if info is None else info.setdefault("edge_series_norm_terms", {}) + norm_terms = {} + gate_terms = {} + for term in terms: + region = frozenset((*tids, *term.tids)) + cache_key = (term.edges, region, tuple(where)) + try: + norm_e = norm_cache[cache_key] + except KeyError: + norm_e = _get_d2_edge_partial_trace_excited( + bp, + region, + excited_edges=term.edges, + exclude=inner_bonds, + ).contract(optimize=optimize, **contract_opts) + norm_cache[cache_key] = norm_e + gate_e = _get_d2_edge_partial_trace_excited( + bp, + region, + excited_edges=term.edges, + exclude=inner_bonds, + gate=gate, + gate_inds=kix, + ).contract(optimize=optimize, **contract_opts) + norm_terms[term.edges] = norm_e + gate_terms[term.edges] = gate_e + + base_norm = _get_d2_edge_partial_trace_excited( + bp, tids, exclude=inner_bonds + ).contract(optimize=optimize, **contract_opts) + base_value = _get_d2_edge_partial_trace_excited( + bp, + tids, + exclude=inner_bonds, + gate=gate, + gate_inds=kix, + ).contract(optimize=optimize, **contract_opts) + suppression = _edge_series_suppression( + norm_terms, + multi_excitation_correct=multi_excitation_correct, + tol_correction=tol_correction, + maxiter_correction=maxiter_correction, + ) + if normalized == "local": + for edges in norm_terms: + scale = 1 + norm_terms[edges] + norm_terms[edges] /= scale + gate_terms[edges] /= scale + norm = base_norm + sum( + norm_terms[edges] * suppression[edges] for edges in norm_terms + ) + value = base_value + sum( + gate_terms[edges] * suppression[edges] for edges in gate_terms + ) + if normalized: + value = value / norm + elif (bp.sign, bp.exponent) != (1.0, 0.0): + value = value * bp.sign * 10**bp.exponent + if info is not None: + info["edge_series_weights"] = dict(norm_terms) + info["edge_series_suppression_factors"] = suppression + info["edge_series_terms"] = terms + return value, norm + + +def _local_expectation_loop_cluster( + bp, + where, + gate, + gloops, + *, + combine, + normalized, + autocomplete, + grow_from, + strict_size, + optimize, + contract_opts, + info, +): + """Contract one gate through the graded D2BP loop-cluster network.""" + if combine != "sum": + raise ValueError( + "graded loop-cluster expectations currently require combine='sum'; " + "the product construction is an elementwise rho operation" + ) + if normalized == "prod": + normalized = True + if normalized is True: + normalized = "local" + if normalized not in (False, "local", "separate"): + raise ValueError( + "normalized must be one of True, False, 'prod', 'local', or " + "'separate'" + ) + + _align_symmray_d2bp_messages(bp) + tags = [bp.tn.site_tag(coo) for coo in where] + tids = frozenset(bp.tn._get_tids_from_tags(tags, "any")) + kix = [bp.tn.site_ind(coo) for coo in where] + regions = tuple( + bp.tn.get_local_gloops( + tids=tids, + gloops=gloops, + grow_from=grow_from, + strict_size=strict_size, + ) + ) + from quimb.tensor.belief_propagation import gen_region_counts + + term_cache = ( + {} if info is None else info.setdefault("cluster_norm_terms", {}) + ) + norm_terms = [] + gate_terms = [] + counts = [] + for region, count in gen_region_counts(regions, autocomplete=autocomplete): + region = frozenset(region) + cache_key = (region, tuple(where)) + try: + norm_e = term_cache[cache_key] + except KeyError: + norm_e = _get_d2_cluster_norm(bp, region).contract( + optimize=optimize, + **contract_opts, + ) + term_cache[cache_key] = norm_e + gate_e = _get_d2_cluster_norm( + bp, + region, + gate=gate, + gate_inds=kix, + ).contract(optimize=optimize, **contract_opts) + if normalized == "local": + gate_e = gate_e / norm_e + norm_e = 1.0 + norm_terms.append(norm_e) + gate_terms.append(gate_e) + counts.append(count) + + norm = sum(count * value for count, value in zip(counts, norm_terms)) + value = sum(count * value for count, value in zip(counts, gate_terms)) + if normalized == "separate": + value = value / norm + elif not normalized and (bp.sign, bp.exponent) != (1.0, 0.0): + value = value * bp.sign * 10**bp.exponent + return value, norm + + +def _process_weights( + weights, + *, + mantissa, + exponent, + multi_excitation_correct, + tol_correction, + maxiter_correction, + strip_exponent, +): + """Use Quimb's loop-series resummation with edge-degree keys.""" + from quimb.tensor.belief_propagation.bp_common import ( + process_loop_series_expansion_weights, + ) + + suppression = process_loop_series_expansion_weights( + weights, + multi_excitation_correct=multi_excitation_correct, + tol_correction=tol_correction, + maxiter_correction=maxiter_correction, + return_all=True, + ) + correction = -sum( + weight * suppression[edges] for edges, weight in weights.items() + ) + estimate = process_loop_series_expansion_weights( + weights, + mantissa=mantissa, + exponent=exponent, + multi_excitation_correct=multi_excitation_correct, + tol_correction=tol_correction, + maxiter_correction=maxiter_correction, + strip_exponent=strip_exponent, + ) + return estimate, correction, suppression + + +def _contract_loop_series( + bp, + terms, + *, + multi_excitation_correct, + tol_correction, + maxiter_correction, + strip_exponent, + optimize, + contract_opts, + normalize, +): + if normalize: + if bp.__class__.__name__ == "D2BP": + _align_symmray_d2bp_messages(bp) + bp.normalize_message_pairs() + bp.normalize_tensors() + + weights = {} + for term in terms: + weights[term.edges] = _get_edge_excited(bp, term).contract( + optimize=optimize, **contract_opts + ) + estimate, correction, suppression = _process_weights( + weights, + mantissa=bp.sign, + exponent=bp.exponent, + multi_excitation_correct=multi_excitation_correct, + tol_correction=tol_correction, + maxiter_correction=maxiter_correction, + strip_exponent=strip_exponent, + ) + return estimate, weights, correction, suppression + + +def _build_bp( + tn, + *, + norm, + messages, + gauges, + run_bp, + bp_runner, + relay_opts, + max_iterations, + tol, + tol_abs, + tol_rolling_diff, + diis, + damping, + update, + optimize, + bp_opts, + progbar, + validate_graph=True, +): + key, bp_cls = _cluster_bp_class(norm) + if key == "2norm" and _uses_symmray(tn): + tn = _restore_fermionic_dummy_modes(tn) + if key == "1norm" and _uses_symmray(tn): + tn = _dense_bp_tn(tn) + messages = _dense_message_tree(messages) + gauges = _dense_message_tree(gauges) + if bp_runner not in {"plain", "relay"}: + raise ValueError("bp_runner must be either 'plain' or 'relay'") + if gauges is not None and messages is not None: + raise ValueError("pass either messages or gauges, not both") + + if key == "1norm" and validate_graph: + from .gauges import _validate_d1_graph + + _validate_d1_graph(tn) + else: + from .gauges import _validate_d2_graph + + _validate_d2_graph(tn) + + info = {} + if gauges is not None: + from .gauges import ( + d1bp_from_simple_update_gauges, + d2bp_from_simple_update_gauges, + ) + + gauge_builder = ( + d1bp_from_simple_update_gauges + if key == "1norm" + else d2bp_from_simple_update_gauges + ) + gauge_opts = dict(bp_opts) + if key == "2norm": + gauge_opts.setdefault("optimize", optimize) bp = gauge_builder( tn, gauges, @@ -658,94 +1668,625 @@ def _build_bp( update=update, **gauge_opts, ) - if run_bp and bp_runner == "relay": - from .relay import relay_bp + if run_bp and bp_runner == "relay": + from .relay import relay_bp + + relay_kwargs = {} if relay_opts is None else dict(relay_opts) + init_messages = { + key_: ar.do("copy", value) + for key_, value in bp.messages.items() + } + bp_result = relay_bp( + bp.tn, + method="d1bp" if key == "1norm" else "d2bp", + init_messages=init_messages, + max_iterations=max_iterations, + tol=tol, + tol_abs=tol_abs, + tol_rolling_diff=tol_rolling_diff, + damping=damping, + update=update, + **relay_kwargs, + **_filter_gauge_init_only_bp_opts(bp_opts), + ) + bp = bp_result.bp + info = { + "converged": bp_result.converged, + "iterations": bp_result.iterations, + "max_mdiff": bp_result.max_mdiff, + } + elif run_bp: + info = _run_plain_bp( + bp, + max_iterations=max_iterations, + tol=tol, + tol_abs=tol_abs, + tol_rolling_diff=tol_rolling_diff, + diis=diis, + progbar=progbar, + ) + return bp, info + + if run_bp and bp_runner == "relay": + from .relay import relay_bp + + relay_kwargs = {} if relay_opts is None else dict(relay_opts) + relay_bp_opts = dict(bp_opts) + if key == "2norm": + relay_bp_opts["optimize"] = optimize + bp_result = relay_bp( + tn, + method="d1bp" if key == "1norm" else "d2bp", + init_messages=messages, + max_iterations=max_iterations, + tol=tol, + tol_abs=tol_abs, + tol_rolling_diff=tol_rolling_diff, + damping=damping, + update=update, + **relay_kwargs, + **relay_bp_opts, + ) + return bp_result.bp, { + "converged": bp_result.converged, + "iterations": bp_result.iterations, + "max_mdiff": bp_result.max_mdiff, + } + + ctor = {"messages": messages, "damping": damping, "update": update} + if key == "2norm": + ctor["optimize"] = optimize + ctor.update(bp_opts) + bp = bp_cls(tn, **ctor) + if run_bp: + info = _run_plain_bp( + bp, + max_iterations=max_iterations, + tol=tol, + tol_abs=tol_abs, + tol_rolling_diff=tol_rolling_diff, + diis=diis, + progbar=progbar, + ) + return bp, info + + +def loop_series_expand( + tn, + gloops=None, + *, + norm: str = "2norm", + messages=None, + gauges=None, + run_bp: bool = True, + bp_runner: str = "plain", + relay_opts: dict[str, Any] | None = None, + max_iterations: int = 1000, + tol: float = 5e-6, + tol_abs: float | None = None, + tol_rolling_diff: float | None = 0.0, + diis: bool | dict[str, Any] = False, + damping: float = 0.0, + update: str = "sequential", + require_fixed_point: bool = True, + cache: LoopSeriesCache | None = None, + multi_excitation_correct: bool = True, + tol_correction: float = 1e-12, + maxiter_correction: int = 100, + optimize: str = "auto-hq", + strip_exponent: bool = False, + progbar: bool = False, + contract_opts: dict[str, Any] | None = None, + **bp_opts, +) -> LoopSeriesResult: + """Estimate a D1/D2 BP contraction with an edge-resolved loop series. + + ``gloops`` may be an integer maximum *excited-bond degree*, an explicit + iterable of :class:`LoopSeriesTerm` objects, or the legacy Quimb iterable + of tensor-id regions. Integer cutoffs enumerate every connected edge set + for which every incident tensor has at least two excited bonds. Distinct + embeddings are retained separately. Disconnected products are supplied + by the multi-excitation resummation, rather than being collapsed into one + region term. + + ``norm="1norm"`` uses D1BP on a closed scalar tensor network. The default + ``norm="2norm"`` uses D2BP on a PEPS-like network with dangling physical + indices. ``gauges`` can initialize either BP family from simple-update + bond gauges. The loop-series formal cancellation assumes a fixed point; + set ``require_fixed_point=False`` only for an explicitly exploratory + boundary approximation. + """ + if run_bp and ( + not isinstance(max_iterations, (int, np.integer)) or max_iterations < 1 + ): + raise ValueError("max_iterations must be a positive integer when run_bp=True") + if not isinstance(maxiter_correction, (int, np.integer)) or maxiter_correction < 1: + raise ValueError("maxiter_correction must be a positive integer") + if tol_correction < 0: + raise ValueError("tol_correction must be non-negative") + + contract_opts = {} if contract_opts is None else dict(contract_opts) + bp, info = _build_bp( + tn, + norm=norm, + messages=messages, + gauges=gauges, + run_bp=run_bp, + bp_runner=bp_runner, + relay_opts=relay_opts, + max_iterations=max_iterations, + tol=tol, + tol_abs=tol_abs, + tol_rolling_diff=tol_rolling_diff, + diis=diis, + damping=damping, + update=update, + optimize=optimize, + bp_opts=bp_opts, + progbar=progbar, + ) + if require_fixed_point and run_bp and not info.get("converged", False): + raise RuntimeError( + "loop_series_expand requires converged BP messages; pass " + "require_fixed_point=False for an exploratory boundary estimate" + ) + + cache = cache or LoopSeriesCache() + terms = _parse_gloops(bp.tn, gloops, cache=cache) + estimate, weights, correction, suppression = _contract_loop_series( + bp, + terms, + multi_excitation_correct=multi_excitation_correct, + tol_correction=tol_correction, + maxiter_correction=maxiter_correction, + strip_exponent=strip_exponent, + optimize=optimize, + contract_opts=contract_opts, + normalize=True, + ) + return LoopSeriesResult( + estimate=estimate, + gloops=gloops, + norm=str(norm).lower(), + terms=terms, + loop_weights=weights, + free_energy_correction=correction, + suppression_factors=suppression, + multi_excitation_correct=multi_excitation_correct, + bp_converged=info.get("converged"), + bp_iterations=info.get("iterations"), + bp_max_mdiff=info.get("max_mdiff"), + bp=bp, + _cache=cache, + _contract_defaults={ + "tol_correction": tol_correction, + "maxiter_correction": maxiter_correction, + }, + ) + + +def partial_trace_loop_series_expand( + tn, + where, + gloops=None, + *, + messages=None, + gauges=None, + run_bp: bool = True, + bp_runner: str = "plain", + relay_opts: dict[str, Any] | None = None, + max_iterations: int = 1000, + tol: float = 5e-6, + tol_abs: float | None = None, + tol_rolling_diff: float | None = 0.0, + diis: bool | dict[str, Any] = False, + damping: float = 0.0, + update: str = "sequential", + require_fixed_point: bool = True, + normalized: bool | str = True, + grow_from: str = "alldangle", + strict_size: bool = False, + multi_excitation_correct: bool = True, + optimize: str = "auto-hq", + info: dict[str, Any] | None = None, + contract_opts: dict[str, Any] | None = None, + **bp_opts, +): + """Compute a reduced density matrix with the D2BP loop series. + + This is the local-observable counterpart of :func:`loop_series_expand`. + The selected physical sites remain open while BP ``P`` and loop-series + ``Q = I - P`` projectors are inserted on virtual bonds. The returned + matrix is ordered as ``where`` on the ket side followed by ``where`` on + the bra side, fused into a two-dimensional array. + + Parameters + ---------- + tn : TensorNetwork + A PEPS-like tensor network for the native 2-norm BP calculation. + where : sequence + The physical sites whose reduced density matrix is requested. + gloops : int or iterable, optional + Local generalized-loop cutoff or explicit tensor regions. Unlike the + global :func:`loop_series_expand` edge cutoff, an integer here follows + Quimb's local-region loop-series convention. + normalized : bool or {"local", "separate"}, optional + Whether to normalize the final density matrix. ``"local"`` also + normalizes each non-base local contribution before combining it. + grow_from : {"alldangle", "all", "any"}, optional + How local loop regions are generated around ``where``. + multi_excitation_correct : bool, optional + Apply the existing loop-series multi-excitation resummation to the + traces of the local density-matrix contributions. + info : dict, optional + Reusable cache and diagnostics. Reuse only for the same network, + messages, and ``where``. + + Returns + ------- + array_like + The reduced density matrix, with ket and bra site groups fused. + + Notes + ----- + The implementation is D2BP-only because a PEPS wavefunction and native + fermionic Symmray arrays require the two-norm construction. The virtual + projectors remain native Symmray arrays when ``tn`` is fermionic or + block-sparse. The returned physical density matrix also remains native + when possible; its small local trace is materialized densely so omitted + Symmray charge blocks are included in normalization. + """ + if contract_opts is None: + contract_opts = {} + else: + contract_opts = dict(contract_opts) + + where = tuple(where) + if not where: + raise ValueError("where must contain at least one site") + if grow_from not in {"alldangle", "all", "any"}: + raise ValueError( + "grow_from must be one of 'alldangle', 'all', or 'any'" + ) + + bp, bp_info = _build_bp( + tn, + norm="2norm", + messages=messages, + gauges=gauges, + run_bp=run_bp, + bp_runner=bp_runner, + relay_opts=relay_opts, + max_iterations=max_iterations, + tol=tol, + tol_abs=tol_abs, + tol_rolling_diff=tol_rolling_diff, + diis=diis, + damping=damping, + update=update, + optimize=optimize, + bp_opts=bp_opts, + progbar=False, + ) + if require_fixed_point and run_bp and not bp_info.get("converged", False): + raise RuntimeError( + "partial_trace_loop_series_expand requires converged BP messages; " + "pass require_fixed_point=False for an exploratory estimate" + ) + + return _partial_trace_loop_series( + bp, + where, + gloops, + normalized=normalized, + grow_from=grow_from, + strict_size=strict_size, + multi_excitation_correct=multi_excitation_correct, + optimize=optimize, + contract_opts=contract_opts, + info=info, + ) + + +def partial_trace_edge_loop_series_expand( + tn, + where, + gloops=None, + *, + messages=None, + gauges=None, + run_bp: bool = True, + bp_runner: str = "plain", + relay_opts: dict[str, Any] | None = None, + max_iterations: int = 1000, + tol: float = 5e-6, + tol_abs: float | None = None, + tol_rolling_diff: float | None = 0.0, + diis: bool | dict[str, Any] = False, + damping: float = 0.0, + update: str = "sequential", + require_fixed_point: bool = True, + normalized: bool | str = True, + multi_excitation_correct: bool = True, + tol_correction: float = 1e-12, + maxiter_correction: int = 1000, + cache: LoopSeriesCache | None = None, + optimize: str = "auto-hq", + info: dict[str, Any] | None = None, + contract_opts: dict[str, Any] | None = None, + **bp_opts, +): + """Compute a local RDM from canonical, explicit ``P + Q`` edge terms. + + This is the edge-resolved counterpart of + :func:`partial_trace_loop_series_expand`. Here an integer ``gloops`` is + an excited-*edge* degree cutoff, exactly as for + :func:`loop_series_expand`; an iterable can contain + :class:`LoopSeriesTerm` objects or explicit virtual-edge sets. It does + not use Quimb's local-region ``get_local_gloops`` convention. + + The ``where`` support is retained as a directly traced physical region. + Consequently, an explicit term may not put ``Q`` on a virtual bond whose + two endpoint tensors both belong to ``where``. This first edge API covers + the standard one-site and separated-support observables, and makes that + limitation explicit rather than silently changing the term. + + For fermionic networks the returned RDM is useful for charge-block and + trace diagnostics. Evaluate a fermionic operator with + :func:`compute_local_expectation_edge_loop_series`, which inserts the + gate before constructing the graded bra network. + """ + contract_opts = {} if contract_opts is None else dict(contract_opts) + where = tuple(where) + if not where: + raise ValueError("where must contain at least one site") + bp, bp_info = _build_bp( + tn, + norm="2norm", + messages=messages, + gauges=gauges, + run_bp=run_bp, + bp_runner=bp_runner, + relay_opts=relay_opts, + max_iterations=max_iterations, + tol=tol, + tol_abs=tol_abs, + tol_rolling_diff=tol_rolling_diff, + diis=diis, + damping=damping, + update=update, + optimize=optimize, + bp_opts=bp_opts, + progbar=False, + ) + if require_fixed_point and run_bp and not bp_info.get("converged", False): + raise RuntimeError( + "partial_trace_edge_loop_series_expand requires converged BP " + "messages; pass require_fixed_point=False for an exploratory " + "estimate" + ) + return _partial_trace_edge_loop_series( + bp, + where, + gloops, + normalized=normalized, + multi_excitation_correct=multi_excitation_correct, + tol_correction=tol_correction, + maxiter_correction=maxiter_correction, + optimize=optimize, + contract_opts=contract_opts, + cache=cache or LoopSeriesCache(), + info=info, + ) + + +def partial_trace_loop_cluster_expand( + tn, + where, + gloops=None, + *, + messages=None, + gauges=None, + run_bp: bool = True, + bp_runner: str = "plain", + relay_opts: dict[str, Any] | None = None, + max_iterations: int = 1000, + tol: float = 5e-6, + tol_abs: float | None = None, + tol_rolling_diff: float | None = 0.0, + diis: bool | dict[str, Any] = False, + damping: float = 0.0, + update: str = "sequential", + require_fixed_point: bool = True, + combine: str = "sum", + normalized: bool | str = True, + autocomplete: bool = True, + grow_from: str = "alldangle", + strict_size: bool = False, + optimize: str = "auto-hq", + info: dict[str, Any] | None = None, + contract_opts: dict[str, Any] | None = None, + **bp_opts, +): + """Compute a local RDM with D2BP generalized-loop clusters. + + This is the reduced-density-matrix counterpart of + :func:`pepsy.bp.loop_cluster_expand`. It contracts BP-message-closed + generalized-loop regions, combines them with their inclusion--exclusion + counts, and leaves the selected physical ket and bra legs open. + + ``combine="sum"`` is the physical default. ``combine="prod"`` follows + Quimb's elementwise product convention and is primarily useful for + compatibility experiments. Native fermionic Symmray arrays remain native + through every virtual contraction; local traces include omitted charge + sectors before normalization. + """ + if contract_opts is None: + contract_opts = {} + else: + contract_opts = dict(contract_opts) + where = tuple(where) + if not where: + raise ValueError("where must contain at least one site") + if grow_from not in {"alldangle", "all", "any"}: + raise ValueError( + "grow_from must be one of 'alldangle', 'all', or 'any'" + ) + if run_bp and ( + not isinstance(max_iterations, (int, np.integer)) or max_iterations < 1 + ): + raise ValueError("max_iterations must be a positive integer when run_bp=True") + + bp, bp_info = _build_bp( + tn, + norm="2norm", + messages=messages, + gauges=gauges, + run_bp=run_bp, + bp_runner=bp_runner, + relay_opts=relay_opts, + max_iterations=max_iterations, + tol=tol, + tol_abs=tol_abs, + tol_rolling_diff=tol_rolling_diff, + diis=diis, + damping=damping, + update=update, + optimize=optimize, + bp_opts=bp_opts, + progbar=False, + ) + if require_fixed_point and run_bp and not bp_info.get("converged", False): + raise RuntimeError( + "partial_trace_loop_cluster_expand requires converged BP messages; " + "pass require_fixed_point=False for an exploratory estimate" + ) + + return _partial_trace_loop_cluster( + bp, + where, + gloops, + combine=combine, + normalized=normalized, + autocomplete=autocomplete, + grow_from=grow_from, + strict_size=strict_size, + optimize=optimize, + contract_opts=contract_opts, + info=info, + ) - relay_kwargs = {} if relay_opts is None else dict(relay_opts) - init_messages = { - key_: ar.do("copy", value) - for key_, value in bp.messages.items() - } - bp_result = relay_bp( - bp.tn, - method="d1bp" if key == "1norm" else "d2bp", - init_messages=init_messages, - max_iterations=max_iterations, - tol=tol, - tol_abs=tol_abs, - tol_rolling_diff=tol_rolling_diff, - damping=damping, - update=update, - **relay_kwargs, - **_filter_gauge_init_only_bp_opts(bp_opts), - ) - bp = bp_result.bp - info = { - "converged": bp_result.converged, - "iterations": bp_result.iterations, - "max_mdiff": bp_result.max_mdiff, - } - elif run_bp: - info = _run_plain_bp( - bp, - max_iterations=max_iterations, - tol=tol, - tol_abs=tol_abs, - tol_rolling_diff=tol_rolling_diff, - diis=diis, - progbar=progbar, - ) - return bp, info - if run_bp and bp_runner == "relay": - from .relay import relay_bp +def compute_local_expectation_loop_cluster( + tn, + terms, + gloops=None, + *, + messages=None, + gauges=None, + run_bp: bool = True, + bp_runner: str = "plain", + relay_opts: dict[str, Any] | None = None, + max_iterations: int = 1000, + tol: float = 5e-6, + tol_abs: float | None = None, + tol_rolling_diff: float | None = 0.0, + diis: bool | dict[str, Any] = False, + damping: float = 0.0, + update: str = "sequential", + require_fixed_point: bool = True, + combine: str = "sum", + normalized: bool | str = True, + autocomplete: bool = True, + grow_from: str = "alldangle", + strict_size: bool = False, + optimize: str = "auto-hq", + info: dict[str, Any] | None = None, + return_all: bool = False, + contract_opts: dict[str, Any] | None = None, + **bp_opts, +): + """Compute local expectations from D2BP loop-cluster RDMs. - relay_kwargs = {} if relay_opts is None else dict(relay_opts) - relay_bp_opts = dict(bp_opts) - if key == "2norm": - relay_bp_opts["optimize"] = optimize - bp_result = relay_bp( - tn, - method="d1bp" if key == "1norm" else "d2bp", - init_messages=messages, - max_iterations=max_iterations, - tol=tol, - tol_abs=tol_abs, - tol_rolling_diff=tol_rolling_diff, - damping=damping, - update=update, - **relay_kwargs, - **relay_bp_opts, + The call accepts the usual ``{site_or_sites: operator}`` term mapping and + shares one D2BP solve between all supports. It is the scalar companion to + :func:`partial_trace_loop_cluster_expand`. + """ + if not hasattr(terms, "items"): + raise TypeError("terms must be a mapping from sites to operators") + if not terms: + raise ValueError("terms must contain at least one operator") + if normalized == "prod": + normalized = True + if grow_from not in {"alldangle", "all", "any"}: + raise ValueError( + "grow_from must be one of 'alldangle', 'all', or 'any'" ) - return bp_result.bp, { - "converged": bp_result.converged, - "iterations": bp_result.iterations, - "max_mdiff": bp_result.max_mdiff, - } + if run_bp and ( + not isinstance(max_iterations, (int, np.integer)) or max_iterations < 1 + ): + raise ValueError("max_iterations must be a positive integer when run_bp=True") - ctor = {"messages": messages, "damping": damping, "update": update} - if key == "2norm": - ctor["optimize"] = optimize - ctor.update(bp_opts) - bp = bp_cls(tn, **ctor) - if run_bp: - info = _run_plain_bp( + contract_opts = {} if contract_opts is None else dict(contract_opts) + bp, bp_info = _build_bp( + tn, + norm="2norm", + messages=messages, + gauges=gauges, + run_bp=run_bp, + bp_runner=bp_runner, + relay_opts=relay_opts, + max_iterations=max_iterations, + tol=tol, + tol_abs=tol_abs, + tol_rolling_diff=tol_rolling_diff, + diis=diis, + damping=damping, + update=update, + optimize=optimize, + bp_opts=bp_opts, + progbar=False, + ) + if require_fixed_point and run_bp and not bp_info.get("converged", False): + raise RuntimeError( + "compute_local_expectation_loop_cluster requires converged BP " + "messages; pass require_fixed_point=False for an exploratory " + "estimate" + ) + + term_info = ( + {} if info is None else info.setdefault("cluster_normalization_by_term", {}) + ) + expecs = {} + for where, gate in terms.items(): + sites = _term_sites(bp.tn, where) + value, normalization = _local_expectation_loop_cluster( bp, - max_iterations=max_iterations, - tol=tol, - tol_abs=tol_abs, - tol_rolling_diff=tol_rolling_diff, - diis=diis, - progbar=progbar, + sites, + gate, + gloops, + combine=combine, + normalized=normalized, + autocomplete=autocomplete, + grow_from=grow_from, + strict_size=strict_size, + optimize=optimize, + contract_opts=contract_opts, + info=info, ) - return bp, info + term_info[where] = normalization + expecs[where] = value + if return_all: + return expecs + return functools.reduce(operator.add, expecs.values()) -def loop_series_expand( + +def compute_local_expectation_loop_series( tn, + terms, gloops=None, *, - norm: str = "2norm", messages=None, gauges=None, run_bp: bool = True, @@ -759,46 +2300,53 @@ def loop_series_expand( damping: float = 0.0, update: str = "sequential", require_fixed_point: bool = True, - cache: LoopSeriesCache | None = None, + normalized: bool | str = True, + grow_from: str = "alldangle", + strict_size: bool = False, multi_excitation_correct: bool = True, - tol_correction: float = 1e-12, - maxiter_correction: int = 100, optimize: str = "auto-hq", - strip_exponent: bool = False, - progbar: bool = False, + info: dict[str, Any] | None = None, + return_all: bool = False, contract_opts: dict[str, Any] | None = None, **bp_opts, -) -> LoopSeriesResult: - """Estimate a D1/D2 BP contraction with an edge-resolved loop series. - - ``gloops`` may be an integer maximum *excited-bond degree*, an explicit - iterable of :class:`LoopSeriesTerm` objects, or the legacy Quimb iterable - of tensor-id regions. Integer cutoffs enumerate every connected edge set - for which every incident tensor has at least two excited bonds. Distinct - embeddings are retained separately. Disconnected products are supplied - by the multi-excitation resummation, rather than being collapsed into one - region term. - - ``norm="1norm"`` uses D1BP on a closed scalar tensor network. The default - ``norm="2norm"`` uses D2BP on a PEPS-like network with dangling physical - indices. ``gauges`` can initialize either BP family from simple-update - bond gauges. The loop-series formal cancellation assumes a fixed point; - set ``require_fixed_point=False`` only for an explicitly exploratory - boundary approximation. +): + """Compute local operator expectations from D2BP loop-series RDMs. + + This is the scalar companion to + :func:`partial_trace_loop_series_expand`. ``terms`` has the familiar + ``{site_or_sites: operator}`` form used by Quimb's local-expectation + methods. A single D2BP solve is shared by all terms, while each support + gets its own local loop-series reduced density matrix. + + ``normalized="prod"`` is accepted as a compatibility spelling for a + normalized local RDM. Unlike Quimb's generalized-loop *cluster* + expectation API, this is the D2 ``P + Q`` loop-series construction, so + there is no separate inclusion--exclusion ``combine`` mode. """ + if not hasattr(terms, "items"): + raise TypeError("terms must be a mapping from sites to operators") + if not terms: + raise ValueError("terms must contain at least one operator") + if normalized == "prod": + normalized = True + if normalized not in (True, False, "local", "separate"): + raise ValueError( + "normalized must be one of True, False, 'prod', 'local', or " + "'separate'" + ) + if grow_from not in {"alldangle", "all", "any"}: + raise ValueError( + "grow_from must be one of 'alldangle', 'all', or 'any'" + ) if run_bp and ( not isinstance(max_iterations, (int, np.integer)) or max_iterations < 1 ): raise ValueError("max_iterations must be a positive integer when run_bp=True") - if not isinstance(maxiter_correction, (int, np.integer)) or maxiter_correction < 1: - raise ValueError("maxiter_correction must be a positive integer") - if tol_correction < 0: - raise ValueError("tol_correction must be non-negative") contract_opts = {} if contract_opts is None else dict(contract_opts) - bp, info = _build_bp( + bp, bp_info = _build_bp( tn, - norm=norm, + norm="2norm", messages=messages, gauges=gauges, run_bp=run_bp, @@ -813,43 +2361,141 @@ def loop_series_expand( update=update, optimize=optimize, bp_opts=bp_opts, - progbar=progbar, + progbar=False, ) - if require_fixed_point and run_bp and not info.get("converged", False): + if require_fixed_point and run_bp and not bp_info.get("converged", False): raise RuntimeError( - "loop_series_expand requires converged BP messages; pass " - "require_fixed_point=False for an exploratory boundary estimate" + "compute_local_expectation_loop_series requires converged BP " + "messages; pass require_fixed_point=False for an exploratory " + "estimate" ) - cache = cache or LoopSeriesCache() - terms = _parse_gloops(bp.tn, gloops, cache=cache) - estimate, weights, correction, suppression = _contract_loop_series( - bp, - terms, - multi_excitation_correct=multi_excitation_correct, - tol_correction=tol_correction, - maxiter_correction=maxiter_correction, - strip_exponent=strip_exponent, + term_info = ( + {} if info is None else info.setdefault("normalization_by_term", {}) + ) + expecs = {} + for where, gate in terms.items(): + sites = _term_sites(bp.tn, where) + value, normalization = _local_expectation_loop_series( + bp, + sites, + gate, + gloops, + normalized=normalized, + grow_from=grow_from, + strict_size=strict_size, + multi_excitation_correct=multi_excitation_correct, + optimize=optimize, + contract_opts=contract_opts, + info=info, + ) + term_info[where] = normalization + expecs[where] = value + + if return_all: + return expecs + return functools.reduce(operator.add, expecs.values()) + + +def compute_local_expectation_edge_loop_series( + tn, + terms, + gloops=None, + *, + messages=None, + gauges=None, + run_bp: bool = True, + bp_runner: str = "plain", + relay_opts: dict[str, Any] | None = None, + max_iterations: int = 1000, + tol: float = 5e-6, + tol_abs: float | None = None, + tol_rolling_diff: float | None = 0.0, + diis: bool | dict[str, Any] = False, + damping: float = 0.0, + update: str = "sequential", + require_fixed_point: bool = True, + normalized: bool | str = True, + multi_excitation_correct: bool = True, + tol_correction: float = 1e-12, + maxiter_correction: int = 1000, + cache: LoopSeriesCache | None = None, + optimize: str = "auto-hq", + info: dict[str, Any] | None = None, + return_all: bool = False, + contract_opts: dict[str, Any] | None = None, + **bp_opts, +): + """Compute graded local expectations from explicit edge loop terms. + + The operator mapping has the usual ``{site_or_sites: gate}`` form. In + contrast to :func:`compute_local_expectation_loop_series`, ``gloops`` is + parsed by :func:`loop_series_expand`: integers count Q edges and explicit + :class:`LoopSeriesTerm` objects preserve their exact virtual-edge set. + The gate is inserted directly into the ket before the graded bra layer is + built, so this is the fermion-safe scalar path. + """ + if not hasattr(terms, "items"): + raise TypeError("terms must be a mapping from sites to operators") + if not terms: + raise ValueError("terms must contain at least one operator") + if normalized == "prod": + normalized = True + if normalized not in (True, False, "local", "separate"): + raise ValueError( + "normalized must be one of True, False, 'prod', 'local', or " + "'separate'" + ) + contract_opts = {} if contract_opts is None else dict(contract_opts) + bp, bp_info = _build_bp( + tn, + norm="2norm", + messages=messages, + gauges=gauges, + run_bp=run_bp, + bp_runner=bp_runner, + relay_opts=relay_opts, + max_iterations=max_iterations, + tol=tol, + tol_abs=tol_abs, + tol_rolling_diff=tol_rolling_diff, + diis=diis, + damping=damping, + update=update, optimize=optimize, - contract_opts=contract_opts, - normalize=True, + bp_opts=bp_opts, + progbar=False, ) - return LoopSeriesResult( - estimate=estimate, - gloops=gloops, - norm=str(norm).lower(), - terms=terms, - loop_weights=weights, - free_energy_correction=correction, - suppression_factors=suppression, - multi_excitation_correct=multi_excitation_correct, - bp_converged=info.get("converged"), - bp_iterations=info.get("iterations"), - bp_max_mdiff=info.get("max_mdiff"), - bp=bp, - _cache=cache, - _contract_defaults={ - "tol_correction": tol_correction, - "maxiter_correction": maxiter_correction, - }, + if require_fixed_point and run_bp and not bp_info.get("converged", False): + raise RuntimeError( + "compute_local_expectation_edge_loop_series requires converged " + "BP messages; pass require_fixed_point=False for an exploratory " + "estimate" + ) + + cache = cache or LoopSeriesCache() + term_info = ( + {} if info is None else info.setdefault("edge_normalization_by_term", {}) ) + expecs = {} + for where, gate in terms.items(): + sites = _term_sites(bp.tn, where) + value, normalization = _local_expectation_edge_loop_series( + bp, + sites, + gate, + gloops, + normalized=normalized, + multi_excitation_correct=multi_excitation_correct, + tol_correction=tol_correction, + maxiter_correction=maxiter_correction, + optimize=optimize, + contract_opts=contract_opts, + cache=cache, + info=info, + ) + term_info[where] = normalization + expecs[where] = value + if return_all: + return expecs + return functools.reduce(operator.add, expecs.values()) From 840ee101a16d79fdbd226dd2f7671fd41c78eb8b Mon Sep 17 00:00:00 2001 From: rezaquant Date: Thu, 30 Jul 2026 12:05:56 -0600 Subject: [PATCH 30/70] Add compressed fermionic BP path observables --- .github/skills/belief-propagation/SKILL.md | 60 ++ docs/api/bp.md | 118 +++ src/pepsy/bp/__init__.py | 8 + src/pepsy/bp/observables.py | 485 +++++++++ tests/test_bp_symmray.py | 1109 ++++++++++++++++++++ 5 files changed, 1780 insertions(+) create mode 100644 src/pepsy/bp/observables.py diff --git a/.github/skills/belief-propagation/SKILL.md b/.github/skills/belief-propagation/SKILL.md index f07b8d8..081797d 100644 --- a/.github/skills/belief-propagation/SKILL.md +++ b/.github/skills/belief-propagation/SKILL.md @@ -84,6 +84,15 @@ over `quimb.tensor.belief_propagation`; the annotated paper trail lives in positive-weight passing on a closed pairwise network. Call `result.projectors(rank=r)` and pass the returned projectors to PNE on the returned gauge-transformed network. +- Long-range PEPS observable helpers: `compute_boundary_expectation(tn, terms, + max_bond=chi, ...)` batches one- and two-site terms, including separated + support, through Quimb's boundary environment. For a connected local + approximation, `compute_path_cluster_expectation(tn, terms, + max_distance=..., gauges=su_gauges, ...)` joins a two-site support by a graph + path and uses simple-update bond vectors to close the cluster boundary. + `compute_bp_path_expectation(...)` is the safe fermionic convenience route: + native D2BP -> Pepsy BP-to-SU conversion -> path-cluster expectation. + Terms use Quimb's mapping form, e.g. `{((x0, y0), (x1, y1)): operator}`. - SU/simple-gauge bridge helpers: `simple_update_messages_from_gauges`, `d1bp_from_simple_update_gauges`, `run_d1bp_from_simple_update_gauges`, `simple_update_bp_residual`, `d2bp_from_simple_update_gauges`, @@ -249,6 +258,57 @@ Loop corrections split by how much they need a **converged BP fixed point**: - `weight_pass` is intentionally restricted to closed pairwise networks; for a D2 calculation, obtain the environment on the appropriate closed double-layer network before supplying projectors to D2 PNE. + +## Long-range fermionic PEPS observables + +Use the native two-site Symmray operator built by `Fermion` or +`Fermion.operator_term`; never construct separated odd operators with a plain +Kronecker product or add a Jordan--Wigner string to the native path. Symmray +supplies the graded signs when the complete ordered native operator is +contracted with the native PEPS. + +For a finite path cluster, `gauges` must be SU/simple-update *bond vectors*. +Do **not** pass D2BP's directed positive-semidefinite matrix messages to +Quimb's `compute_local_expectation_cluster`. Use `compute_bp_path_expectation` +or `simple_update_core_and_gauges_from_d2bp` to convert them first. The +returned `core` and external gauge vectors represent the same state only after +`core.copy().gauge_simple_insert(gauges)`; the cluster routine uses the vectors +only to close its cut boundary. + +For native Symmray PEPS, path-cluster compression (`max_bond=chi`) uses Pepsy's +graded adapter around Quimb's public compressed-contraction API. It keeps the +observable RDM physical legs unfused, aligns zero-weight virtual charge sectors +on a private cluster copy, and uses Symmray's fermionic `squeeze` before +Quimb's QR/SVD steps. Thus standard Quimb/Cotengra path optimizers can be +supplied through `optimize`. A finite cluster or finite-chi boundary discrepancy +is an environment approximation, not by itself a fermionic-sign failure; +enlarge the region/chi against an exact small reference. + +### Required sign regression + +Any change to native fermionic D2BP, SU gauging, BP-to-SU conversion, or +long-range measurement must preserve these tests in +`tests/test_bp_symmray.py`: + +- `test_fermionic_long_range_hopping_sign_survives_su_and_bp_gauges` prepares + a controlled spinless U1 2x2 Fock PEPS whose diagonal hopping correlator + crosses an occupied mode in row-major Jordan--Wigner order. +- `test_spinful_long_range_hopping_sign_survives_su_and_bp_gauges` repeats the + parity-sensitive up-fermion correlator for spinful U1 and U1U1 PEPS. +- `test_spinful_eta_pair_measurement_survives_su_and_bp_gauges` prepares a + long-range eta-pair observable with the public routed 2D gate path and + verifies its imaginary-time expectation for spinful U1 and U1U1 PEPS. + +Each has an independent dense JW oracle; the hopping cases also compare it to +the deliberately no-string bosonic control, which has the opposite sign. The +native exact contraction, a distinct contraction path, direct SU +reconstruction, D2BP-to-SU reconstruction, and the public BP path helper must +all match the JW value. The spinless regression also checks the native +compressed path-cluster route at a sufficiently large `max_bond`. Also assert +that native tensors/messages/vectors retain their Symmray types. This is the +sign oracle; matching only density observables or only two native contraction +routes is insufficient. + - Native Symmray fermionic BP coverage is in `tests/test_bp_symmray.py` and exercises `U1`, `U1U1`, and `Z2` SU↔D2BP round trips, D2 relay/loop-series/ PNE corrections, and closed-scalar 1-norm compatibility. Preserve the array diff --git a/docs/api/bp.md b/docs/api/bp.md index 16b02f0..95b9a2a 100644 --- a/docs/api/bp.md +++ b/docs/api/bp.md @@ -6,3 +6,121 @@ the top-level package. > API details are maintained as handwritten Markdown in this page. + +## Long-range PEPS expectations + +Use `compute_boundary_expectation` for batched one- and two-site operators, +including separated supports. It delegates to Quimb's PEPS boundary +environment and preserves native fermionic Symmray operators. + +```python +from pepsy.bp import compute_boundary_expectation + +value = compute_boundary_expectation( + peps.tn, + terms, # {(site_a, site_b): native two-site operator} + max_bond=chi, + normalized=True, +) +``` + +For a controlled finite-region estimate, use +`compute_path_cluster_expectation`. For two-site support, Quimb connects the +sites by a graph path, expands it by `max_distance`, and optionally fills +lattice corners. Supply only compatible simple-update/SU bond vectors through +`gauges`; D2BP matrix messages are not SU gauges. + +```python +from pepsy.bp import compute_path_cluster_expectation + +value = compute_path_cluster_expectation( + peps.tn, + terms, + max_distance=1, + fillin=True, + gauges=su_gauges, + max_bond=chi, + optimize="auto-hq", +) +``` + +`compute_bp_path_expectation` runs Pepsy's native fermionic D2BP, converts its +messages through the tested BP-to-SU bridge, and then evaluates the connected +path cluster. Native Symmray path clusters also accept `max_bond=chi`: Pepsy +keeps the local RDM physical legs unfused, pads zero-weight charge sectors on +the private cluster copy, and lets Quimb perform its usual QR/SVD compressed +contraction. Pass `optimize="auto-hq"`, another Quimb optimizer string, or a +standard Cotengra path optimizer through either path-cluster helper. + +## Local reduced density matrices + +For a D2BP loop-series estimate of a local reduced density matrix, keep the +requested physical sites open with `partial_trace_loop_series_expand`: + +```python +from pepsy.bp import partial_trace_loop_series_expand + +rho = partial_trace_loop_series_expand( + peps.tn, + where=((0, 0), (1, 1)), + gloops=2, + normalized=True, +) +``` + +The matrix is ordered as the selected sites on the ket side followed by the +same sites on the bra side. The D2BP virtual messages and `P`/`Q` projectors +remain native for fermionic Symmray PEPS; use `rho.to_dense()` only when a +dense physical matrix is needed for inspection or an external observable. + +For a mapping of Hamiltonian terms, use the scalar companion and take `.real` +when the Hamiltonian is Hermitian: + +```python +from pepsy.bp import compute_local_expectation_loop_series + +energy = compute_local_expectation_loop_series( + peps.tn, + mag_terms, + gloops=2, + normalized="prod", # compatibility spelling for a normalized local rho +).real +``` + +Native fermionic operators are preferred for Symmray PEPS: their charge +sectors and graded swaps are handled by inserting the gate into the ket before +forming the BP double layer. Do not reconstruct a fermionic expectation as a +dense `trace(rho @ operator)`: the returned rho is useful for diagnostics, but +that contraction misses the physical graded ordering. + +### Explicit edge-subset loop series + +The APIs above follow Quimb's *local-region* convention: their integer +`gloops` is a region cutoff. For a brute-force-compatible expansion over the +canonical virtual Q-edge sets used by `loop_series_expand`, use the explicit +edge path instead: + +```python +from pepsy.bp import compute_local_expectation_edge_loop_series + +energy = compute_local_expectation_edge_loop_series( + peps.tn, + mag_terms, + gloops=4, # maximum number of explicitly excited Q edges + normalized="prod", +).real +``` + +`partial_trace_edge_loop_series_expand` supplies the matching diagnostic RDM. +The scalar function is the fermion-safe choice for its supported cases. +Explicit terms can be passed as `LoopSeriesTerm` objects (or virtual-edge sets); at present they cannot put +Q on a bond wholly internal to the selected observable support. A nonzero-Q +fermionic scalar correction is currently restricted to one-site gates; the +multi-site graded-Q contraction is rejected explicitly while its block routing +is completed. + +`partial_trace_loop_cluster_expand` and +`compute_local_expectation_loop_cluster` provide the parallel D2BP +generalized-loop-cluster route. Its default `combine="sum"` uses the usual +inclusion--exclusion region counts; reserve `combine="prod"` for compatibility +experiments with Quimb's elementwise product convention. diff --git a/src/pepsy/bp/__init__.py b/src/pepsy/bp/__init__.py index 98b0185..71403fe 100644 --- a/src/pepsy/bp/__init__.py +++ b/src/pepsy/bp/__init__.py @@ -93,6 +93,11 @@ relay_bp, two_norm_bp, ) +from .observables import ( + compute_boundary_expectation, + compute_bp_path_expectation, + compute_path_cluster_expectation, +) from .reduced_update import ( ExactReducedUpdateProblem, LoopClusterReducedUpdateProblem, @@ -183,4 +188,7 @@ "solve_reduced_als", "su_cluster_reduced_update_problem", "two_norm_bp", + "compute_boundary_expectation", + "compute_path_cluster_expectation", + "compute_bp_path_expectation", ] diff --git a/src/pepsy/bp/observables.py b/src/pepsy/bp/observables.py new file mode 100644 index 0000000..dc3849d --- /dev/null +++ b/src/pepsy/bp/observables.py @@ -0,0 +1,485 @@ +"""Long-range PEPS expectation helpers with BP-aware environment choices. + +The boundary route delegates to Quimb's 2D PEPS environment contraction. The +path-cluster route is useful when a two-site operator has widely separated +support: Quimb explicitly connects the sites, expands a graph-distance +neighbourhood, and contracts the resulting finite cluster. Native Symmray +clusters use a small adapter around Quimb's compressed-contraction API so +QR/SVD truncations retain their charge sectors and graded metadata. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +import autoray as ar + +from ._symmray import ( + is_symmray_array as _is_symmray_array, + uses_symmray as _uses_symmray, +) + +__all__ = [ + "compute_boundary_expectation", + "compute_path_cluster_expectation", + "compute_bp_path_expectation", +] + + +def _validate_terms(terms): + if not isinstance(terms, Mapping): + raise TypeError("terms must be a mapping from site support to operators") + if not terms: + raise ValueError("terms must contain at least one operator") + return terms + + +def _term_sites(tn, where): + """Normalize a Quimb local-term key to an ordered site tuple.""" + has_site = getattr(tn, "has_site", None) + if callable(has_site) and has_site(where): + return (where,) + if isinstance(where, (str, bytes)): + return (where,) + try: + sites = tuple(where) + except TypeError: + return (where,) + if not sites: + raise ValueError("a local expectation term must have at least one site") + return sites + + +def _squeeze_native_singleton_bonds(tn): + """Remove dimension-one bonds with Symmray's graded squeeze operation. + + Quimb's compressed-contraction implementation correctly detects such bonds + but removes them with a generic reshape. For a fermionic Symmray array, + ``squeeze`` additionally records an odd singlet as a dummy mode, which is + necessary to retain its graded ordering. Doing this before the public + ``contract_compressed`` call avoids the dense-only reshape path without + changing Quimb's contraction or SVD algorithms. + """ + for index in tuple(tn.inner_inds()): + if tn.ind_size(index) != 1: + continue + for tid in tuple(tn.ind_map[index]): + tensor = tn.tensor_map[tid] + axis = tensor.inds.index(index) + data = ar.do("squeeze", tensor.data, axis=axis) + inds = tuple(ind for ind in tensor.inds if ind != index) + left_inds = ( + None + if tensor.left_inds is None + else tuple(ind for ind in tensor.left_inds if ind != index) + ) + tensor.modify(data=data, inds=inds, left_inds=left_inds) + + +def _align_native_cluster_bonds(tn): + """Pad cluster bond sectors so both Symmray endpoints share one layout. + + A BP-to-SU conversion can legitimately leave a zero-weight charge sector + absent from one endpoint tensor while the opposite endpoint retains it. + Exact Symmray contractions handle that sparse support, but Quimb's + compressed-contraction bookkeeping uses one dense bond size per index. + Pad only the private cluster copy with zero blocks before compression. + """ + for index in tuple(tn.inner_inds()): + tids = tuple(tn.ind_map[index]) + if len(tids) != 2: + continue + left_tid, right_tid = tids + left_tensor = tn.tensor_map[left_tid] + right_tensor = tn.tensor_map[right_tid] + if not ( + _is_symmray_array(left_tensor.data) + and _is_symmray_array(right_tensor.data) + ): + continue + + left_axis = left_tensor.inds.index(index) + right_axis = right_tensor.inds.index(index) + left_index = left_tensor.data.indices[left_axis] + right_index = right_tensor.data.indices[right_axis] + left_map = dict(left_index.chargemap) + right_map = dict(right_index.chargemap) + for charge in left_map.keys() & right_map.keys(): + if left_map[charge] != right_map[charge]: + raise ValueError( + "incompatible Symmray virtual charge dimensions on " + f"bond {index!r} for compressed path contraction" + ) + shared_map = {**left_map, **right_map} + if left_map == right_map: + continue + + left_indices = list(left_tensor.data.indices) + left_indices[left_axis] = left_index.copy_with( + chargemap=shared_map, + dual=left_index.dual, + ) + left_data = left_tensor.data.copy_with(indices=tuple(left_indices)) + left_data.fill_missing_blocks() + left_tensor.modify(data=left_data) + + right_indices = list(right_tensor.data.indices) + right_indices[right_axis] = right_index.copy_with( + chargemap=shared_map, + dual=right_index.dual, + ) + right_data = right_tensor.data.copy_with(indices=tuple(right_indices)) + right_data.fill_missing_blocks() + right_tensor.modify(data=right_data) + + +def _native_compressed_path_expectation( + tn, + terms, + *, + max_distance, + mode, + fillin, + grow_from, + gauges, + smudge, + power, + max_bond, + normalized, + optimize, + return_all, + contract_opts, +): + """Contract native Symmray path clusters with Quimb's compression engine. + + Quimb's generic ``local_expectation`` route fuses the resulting RDM and + contracts it as a dense matrix. Keep the RDM physical legs unfused instead: + that retains their fermionic index metadata, lets the native operator + supply the graded ordering, and still delegates every contraction and + truncation to Quimb/Cotengra/Symmray. + """ + if normalized not in (True, False): + raise ValueError( + "native Symmray compressed path clusters require normalized to " + "be True or False" + ) + + options = dict(contract_opts) + flatten = options.pop("flatten", True) + reduce = options.pop("reduce", False) + symmetrized = options.pop("symmetrized", "auto") + rehearse = options.pop("rehearse", False) + method = options.pop("method", "contract_compressed") + if reduce: + raise NotImplementedError( + "reduce=True is not yet supported for native Symmray compressed " + "path clusters" + ) + if rehearse: + raise NotImplementedError( + "rehearse is not yet supported for native Symmray compressed " + "path clusters" + ) + if method != "contract_compressed": + raise ValueError( + "native Symmray compressed path clusters use " + "method='contract_compressed'" + ) + if symmetrized == "auto": + symmetrized = not flatten + + user_post_contract = options.pop("callback_post_contract", None) + user_post_compress = options.pop("callback_post_compress", None) + + def _post_contract(work_tn, tid): + # A contraction can drop a zero-weight charge sector at one endpoint. + # Re-pad it before a following compression inspects dense bond sizes. + _align_native_cluster_bonds(work_tn) + # A compression can also create a new singleton bond after the initial + # cleanup. Remove it before Quimb's next generic squeeze pass. + _squeeze_native_singleton_bonds(work_tn) + if user_post_contract is not None: + user_post_contract(work_tn, tid) + + def _post_compress(work_tn, tids): + _align_native_cluster_bonds(work_tn) + _squeeze_native_singleton_bonds(work_tn) + if user_post_compress is not None: + user_post_compress(work_tn, tids) + + options["callback_post_contract"] = _post_contract + options["callback_post_compress"] = _post_compress + + expecs = {} + for where, gate in terms.items(): + if not _is_symmray_array(gate): + raise TypeError( + "native Symmray compressed path clusters require native " + "Symmray local operators" + ) + + sites = _term_sites(tn, where) + cluster = tn.get_cluster( + sites, + gauges=gauges, + max_distance=max_distance, + mode=mode, + fillin=fillin, + grow_from=grow_from, + smudge=smudge, + power=power, + ).copy() + _align_native_cluster_bonds(cluster) + ket_inds = tuple(map(cluster.site_ind, sites)) + bra_ind_id = "_pepsy_bra{}" + bra_inds = tuple(map(bra_ind_id.format, sites)) + rdm = cluster.make_reduced_density_matrix( + sites, + bra_ind_id=bra_ind_id, + ) + + if flatten: + for site in cluster.gen_site_coos(): + if site not in sites or flatten == "all": + tag = rdm.site_tag(site) + if tag in rdm.tag_map: + rdm ^= tag + + rdm.fuse_multibonds_() + _squeeze_native_singleton_bonds(rdm) + rho_tensor = rdm.contract_compressed( + optimize, + max_bond=max_bond, + output_inds=ket_inds + bra_inds, + **options, + ) + rho = rho_tensor.data + if normalized: + norm = ar.do("trace", rho_tensor.to_dense(ket_inds, bra_inds)) + rho = rho / norm + if symmetrized: + rho = (rho + ar.do("dag", rho)) / 2 + + nsites = len(sites) + if ar.do("ndim", gate) != 2 * nsites: + gate = ar.do("reshape", gate, ar.do("shape", rho)) + expecs[where] = ar.do( + "tensordot", + rho, + gate, + axes=( + tuple(range(2 * nsites)), + tuple(range(nsites, 2 * nsites)) + tuple(range(nsites)), + ), + ) + + if return_all: + return expecs + return sum(expecs.values()) + + +def compute_boundary_expectation( + tn, + terms, + *, + max_bond=None, + cutoff=1.0e-10, + canonize=True, + mode="mps", + layer_tags=("KET", "BRA"), + normalized=True, + autogroup=True, + contract_optimize="auto-hq", + return_all=False, + plaquette_envs=None, + plaquette_map=None, + **plaquette_env_options, +): + """Compute batched PEPS expectations using Quimb's boundary environment. + + ``terms`` accepts the same one- and two-site support mapping as Quimb's + ``TensorNetwork2DVector.compute_local_expectation``. In particular, a + key such as ``((x0, y0), (x1, y1))`` remains a connected long-range + operator support; it is not replaced by a product of endpoint estimates. + + Parameters are forwarded to Quimb's PEPS boundary implementation. The + returned value is the sum of locally normalized terms unless + ``return_all=True``. + """ + _validate_terms(terms) + if not hasattr(tn, "compute_local_expectation"): + raise TypeError( + "tn must provide Quimb's compute_local_expectation method" + ) + + return tn.compute_local_expectation( + terms, + max_bond=max_bond, + cutoff=cutoff, + canonize=canonize, + mode=mode, + layer_tags=layer_tags, + normalized=normalized, + autogroup=autogroup, + contract_optimize=contract_optimize, + return_all=return_all, + plaquette_envs=plaquette_envs, + plaquette_map=plaquette_map, + **plaquette_env_options, + ) + + +def compute_path_cluster_expectation( + tn, + terms, + *, + max_distance=0, + mode="graphdistance", + fillin=True, + grow_from="all", + gauges=None, + smudge=1.0e-12, + power=1.0, + max_bond=None, + normalized=True, + optimize="auto-hq", + return_all=False, + **contract_opts, +): + """Compute expectations on connected, distance-expanded PEPS clusters. + + For a two-site term, Quimb first adds a graph path between the sites, then + expands that path by ``max_distance``. ``fillin=True`` adds lattice corner + tensors. If ``gauges`` is supplied, it must contain simple-update/SU-style + bond vectors used to close the cluster boundary; D2BP matrix messages must + not be passed here directly. + + Native Symmray clusters use a Pepsy adapter around Quimb's public + compressed-contraction API when ``max_bond`` is supplied. The adapter + preserves the unfused fermionic RDM legs and removes singleton virtual + bonds with Symmray's graded squeeze operation before Quimb performs its + QR/SVD truncations. Thus ``optimize`` accepts Quimb/Cotengra contraction + paths for both exact and compressed native clusters. + """ + _validate_terms(terms) + if not hasattr(tn, "compute_local_expectation_cluster"): + raise TypeError( + "tn must provide Quimb's compute_local_expectation_cluster method" + ) + if max_distance < 0: + raise ValueError("max_distance must be nonnegative") + if _uses_symmray(tn) and max_bond is not None: + return _native_compressed_path_expectation( + tn, + terms, + max_distance=max_distance, + mode=mode, + fillin=fillin, + grow_from=grow_from, + gauges=gauges, + smudge=smudge, + power=power, + max_bond=max_bond, + normalized=normalized, + optimize=optimize, + return_all=return_all, + contract_opts=contract_opts, + ) + + return tn.compute_local_expectation_cluster( + terms, + max_distance=max_distance, + mode=mode, + fillin=fillin, + grow_from=grow_from, + gauges=gauges, + smudge=smudge, + power=power, + max_bond=max_bond, + normalized=normalized, + optimize=optimize, + return_all=return_all, + **contract_opts, + ) + + +def compute_bp_path_expectation( + tn, + terms, + *, + max_distance=0, + mode="graphdistance", + fillin=True, + max_bond=None, + normalized=True, + optimize="auto-hq", + return_all=False, + bp_options: dict[str, Any] | None = None, + conversion_options: dict[str, Any] | None = None, + require_converged=True, + **contract_opts, +): + """Compute path-cluster expectations using a D2BP-derived SU closure. + + This is the safe BP route for fermionic Symmray PEPS. D2BP is run on the + physical network, then Pepsy's tested BP-to-SU bridge converts its native + positive-semidefinite matrix messages into SU-style bond vectors. The + converted core and vectors are passed to Quimb's connected path-cluster + expectation routine; D2BP matrices are never passed as SU gauges. + + ``bp_options`` is forwarded to :func:`pepsy.bp.gauge_all` under its + ``bp_options`` argument, for example ``{"run_opts": {"diis": False}}``. + ``conversion_options`` is forwarded to the BP-to-SU conversion and + defaults to a small regularization of singular message eigenvalues. + + Native Symmray path clusters support ``max_bond`` through the same graded + compressed path as :func:`compute_path_cluster_expectation`. Set + ``require_converged=False`` only for diagnostic experiments with an + unconverged BP fixed point. + """ + _validate_terms(terms) + from .gauges import gauge_all + + bp_options = {} if bp_options is None else dict(bp_options) + conversion_options = ( + {"smudge": 1.0e-12} + if conversion_options is None + else dict(conversion_options) + ) + bridge = gauge_all( + tn, + start="bp", + target="su", + norm="2norm", + bp_options=bp_options, + conversion_options=conversion_options, + ) + + if require_converged and ( + bridge.bp_result is None or not bridge.bp_result.converged + ): + diagnostic = ( + None + if bridge.bp_result is None + else bridge.bp_result.max_mdiff + ) + raise RuntimeError( + "D2BP did not converge before the path-cluster expectation was " + f"requested (max_mdiff={diagnostic!r}); pass " + "require_converged=False for a diagnostic estimate" + ) + + return compute_path_cluster_expectation( + bridge.core, + terms, + max_distance=max_distance, + mode=mode, + fillin=fillin, + gauges=bridge.gauges, + max_bond=max_bond, + normalized=normalized, + optimize=optimize, + return_all=return_all, + **contract_opts, + ) diff --git a/tests/test_bp_symmray.py b/tests/test_bp_symmray.py index a421a66..6cd5ff4 100644 --- a/tests/test_bp_symmray.py +++ b/tests/test_bp_symmray.py @@ -7,23 +7,850 @@ sr = pytest.importorskip("symmray") qtn = pytest.importorskip("quimb.tensor") +ctg = pytest.importorskip("cotengra") from pepsy.bp import ( # noqa: E402 + compute_boundary_expectation, + compute_bp_path_expectation, + compute_local_expectation_edge_loop_series, + compute_local_expectation_loop_cluster, + compute_local_expectation_loop_series, + compute_path_cluster_expectation, gauge_all, loop_cluster_expand, loop_series_expand, one_norm_bp, + partial_trace_edge_loop_series_expand, + partial_trace_loop_series_expand, + partial_trace_loop_cluster_expand, partitioned_expand, relay_bp, two_norm_bp, weight_pass, ) from pepsy.tensors import ( # noqa: E402 + Fermion, SymPEPS, + ps_to_peps, site_charge_alternating, ) +def _long_range_density_term(where, *, symmetry="U1"): + """Return a neutral native density-density observable at ``where``.""" + fermion = Fermion(spinful=True, symmetry=symmetry) + return {where: fermion.density_operator()} + + +def _jw_annihilator(num_modes, mode): + """Return a spinless Jordan--Wigner annihilator in site-major order.""" + identity = np.eye(2, dtype="complex128") + parity = np.diag([1.0, -1.0]).astype("complex128") + annihilate = np.array([[0.0, 1.0], [0.0, 0.0]], dtype="complex128") + factors = [ + parity if site < mode else annihilate if site == mode else identity + for site in range(num_modes) + ] + out = factors[0] + for factor in factors[1:]: + out = np.kron(out, factor) + return out + + +def _jw_hopping_operator(num_modes, left, right): + """Return ``c_left^dag c_right + h.c.`` with its JW parity string.""" + left_annihilate = _jw_annihilator(num_modes, left) + right_annihilate = _jw_annihilator(num_modes, right) + return ( + left_annihilate.conj().T @ right_annihilate + + right_annihilate.conj().T @ left_annihilate + ) + + +def _bosonic_hopping_operator(num_modes, left, right): + """Deliberately omit the parity string for the negative-control oracle.""" + identity = np.eye(2, dtype="complex128") + annihilate = np.array([[0.0, 1.0], [0.0, 0.0]], dtype="complex128") + + def local_annihilator(mode): + factors = [annihilate if site == mode else identity for site in range(num_modes)] + out = factors[0] + for factor in factors[1:]: + out = np.kron(out, factor) + return out + + left_annihilate = local_annihilator(left) + right_annihilate = local_annihilator(right) + return ( + left_annihilate.conj().T @ right_annihilate + + right_annihilate.conj().T @ left_annihilate + ) + + +def _unitary_from_hermitian(hamiltonian, dt): + """Exponentiate a small dense Hermitian Hamiltonian without SciPy.""" + values, vectors = np.linalg.eigh(hamiltonian) + return (vectors * np.exp(-1j * dt * values)) @ vectors.conj().T + + +def _imaginary_time_from_hermitian(hamiltonian, dt): + """Apply ``exp(-dt * hamiltonian)`` for a small dense reference state.""" + values, vectors = np.linalg.eigh(hamiltonian) + return (vectors * np.exp(-dt * values)) @ vectors.conj().T + + +def _jw_eta_pair_operator(left, right): + """Return ``Delta_left^dag Delta_right + h.c.`` in JW mode order.""" + num_modes = 8 + + def pair_create(site): + up = _jw_annihilator(num_modes, 2 * site).conj().T + down = _jw_annihilator(num_modes, 2 * site + 1).conj().T + return up @ down + + def pair_annihilate(site): + up = _jw_annihilator(num_modes, 2 * site) + down = _jw_annihilator(num_modes, 2 * site + 1) + return down @ up + + return ( + pair_create(left) @ pair_annihilate(right) + + pair_create(right) @ pair_annihilate(left) + ) + + +def _spinless_sign_sensitive_peps(): + """Prepare a 2x2 U1 PEPS whose diagonal hop needs a JW minus sign.""" + fermion = Fermion(spinful=False, symmetry="U1") + occupations = { + (0, 0): 1, + (0, 1): 1, + (1, 0): 0, + (1, 1): 0, + } + peps = ps_to_peps( + 2, + 2, + fermion=fermion, + occupations=occupations, + dtype="complex128", + ) + state = SymPEPS( + peps=peps, + symmetry="U1", + edges=tuple(qtn.edges_2d_square(2, 2)), + fermionic=True, + phys_sectors=fermion.physical_sectors, + site_charge=occupations, + site_ind_id="k{},{}", + ) + dt = 0.37 + gate = fermion.hopping_gate(dt, t=1.0) + state.apply_gates( + ( + (gate, ((0, 0), (1, 0))), + (gate, ((1, 0), (1, 1))), + ), + method="direct", + contract="split", + max_bond=4, + cutoff=0.0, + ) + + initial = np.zeros(16, dtype="complex128") + initial[np.ravel_multi_index((1, 1, 0, 0), (2,) * 4)] = 1.0 + hop_02 = _jw_hopping_operator(4, 0, 2) + hop_23 = _jw_hopping_operator(4, 2, 3) + dense_state = _unitary_from_hermitian(-hop_23, dt) @ ( + _unitary_from_hermitian(-hop_02, dt) @ initial + ) + return state, fermion, dense_state + + +def test_fermionic_long_range_hopping_sign_survives_su_and_bp_gauges(): + """A JW sign-sensitive long-range hop is invariant under fermionic gauges.""" + state, fermion, dense_state = _spinless_sign_sensitive_peps() + where = ((0, 0), (1, 1)) + terms = {where: fermion.hopping_operator()} + + jw_operator = _jw_hopping_operator(4, 0, 3) + bosonic_operator = _bosonic_hopping_operator(4, 0, 3) + jw_value = dense_state.conj() @ jw_operator @ dense_state + bosonic_value = dense_state.conj() @ bosonic_operator @ dense_state + assert abs(jw_value) > 1.0e-6 + assert jw_value == pytest.approx(-bosonic_value, abs=1e-12) + + exact_auto = state.tn.compute_local_expectation_exact( + terms, + normalized=True, + optimize="auto-hq", + ) + exact_greedy = state.tn.compute_local_expectation_exact( + terms, + normalized=True, + optimize="greedy", + ) + assert exact_auto == pytest.approx(jw_value, rel=1e-10, abs=1e-10) + assert exact_greedy == pytest.approx(exact_auto, rel=1e-10, abs=1e-10) + + su = gauge_all( + state.tn, + start="su", + target="su", + norm="2norm", + su_options={"max_iterations": 8, "tol": 0.0}, + ) + assert all( + type(gauge).__module__.startswith("symmray") + for gauge in su.gauges.values() + ) + su_reconstructed = su.core.copy() + su_reconstructed.gauge_simple_insert(su.gauges) + su_exact = su_reconstructed.compute_local_expectation_exact( + terms, + normalized=True, + optimize="auto-hq", + ) + assert su_exact == pytest.approx(exact_auto, rel=1e-10, abs=1e-10) + + bp_options = { + "run_opts": { + "max_iterations": 150, + "tol": 1e-10, + "diis": False, + } + } + bridge = gauge_all( + state.tn, + start="bp", + target="su", + norm="2norm", + bp_options=bp_options, + conversion_options={"smudge": 1e-12}, + ) + assert bridge.bp_result.converged + assert all( + type(message).__name__ == "U1FermionicArray" + for message in bridge.messages.values() + ) + bp_reconstructed = bridge.core.copy() + bp_reconstructed.gauge_simple_insert(bridge.gauges) + bp_exact = bp_reconstructed.compute_local_expectation_exact( + terms, + normalized=True, + optimize="auto-hq", + ) + assert bp_exact == pytest.approx(exact_auto, rel=1e-10, abs=1e-10) + + cluster = compute_path_cluster_expectation( + bridge.core, + terms, + gauges=bridge.gauges, + max_distance=1, + fillin=True, + max_bond=None, + normalized=True, + optimize="auto-hq", + ) + assert cluster == pytest.approx(exact_auto, rel=1e-10, abs=1e-10) + + with pytest.warns(UserWarning, match="not a compressed one"): + compressed_cluster = compute_path_cluster_expectation( + bridge.core, + terms, + gauges=bridge.gauges, + max_distance=1, + fillin=True, + max_bond=2, + normalized=True, + optimize="auto-hq", + ) + assert compressed_cluster == pytest.approx( + jw_value, + rel=1e-10, + abs=1e-10, + ) + + bp_helper = compute_bp_path_expectation( + state.tn, + terms, + max_distance=1, + fillin=True, + max_bond=None, + normalized=True, + optimize="auto-hq", + bp_options=bp_options, + conversion_options={"smudge": 1e-12}, + ) + assert bp_helper == pytest.approx(exact_auto, rel=1e-10, abs=1e-10) + + +def _spinful_sign_sensitive_peps(symmetry): + """Prepare a controlled spinful state with an occupied JW string.""" + fermion = Fermion(spinful=True, symmetry=symmetry) + doublon = 2 if symmetry == "U1" else (1, 1) + empty = 0 if symmetry == "U1" else (0, 0) + occupations = { + (0, 0): doublon, + (0, 1): doublon, + (1, 0): empty, + (1, 1): empty, + } + peps = ps_to_peps( + 2, + 2, + fermion=fermion, + occupations=occupations, + dtype="complex128", + ) + state = SymPEPS( + peps=peps, + symmetry=symmetry, + edges=tuple(qtn.edges_2d_square(2, 2)), + fermionic=True, + phys_sectors=fermion.physical_sectors, + site_charge=occupations, + site_ind_id="k{},{}", + ) + dt = 0.29 + gate = fermion.hopping_gate(dt, t=(1.0, 0.0)) + state.apply_gates( + ( + (gate, ((0, 0), (1, 0))), + (gate, ((1, 0), (1, 1))), + ), + method="direct", + contract="split", + max_bond=4, + cutoff=0.0, + ) + + initial = np.zeros(256, dtype="complex128") + initial[np.ravel_multi_index((1, 1, 1, 1, 0, 0, 0, 0), (2,) * 8)] = 1.0 + hop_04 = _jw_hopping_operator(8, 0, 4) + hop_46 = _jw_hopping_operator(8, 4, 6) + dense_state = _unitary_from_hermitian(-hop_46, dt) @ ( + _unitary_from_hermitian(-hop_04, dt) @ initial + ) + return state, fermion, dense_state + + +@pytest.mark.parametrize("symmetry", ("U1", "U1U1")) +def test_spinful_long_range_hopping_sign_survives_su_and_bp_gauges(symmetry): + """Spinful U1 gauges preserve a JW-sensitive up-fermion correlator.""" + state, fermion, dense_state = _spinful_sign_sensitive_peps(symmetry) + where = ((0, 0), (1, 1)) + terms = {where: fermion.hopping_operator(spin="up")} + + jw_operator = _jw_hopping_operator(8, 0, 6) + bosonic_operator = _bosonic_hopping_operator(8, 0, 6) + jw_value = dense_state.conj() @ jw_operator @ dense_state + bosonic_value = dense_state.conj() @ bosonic_operator @ dense_state + assert abs(jw_value) > 1.0e-6 + assert jw_value == pytest.approx(-bosonic_value, abs=1e-12) + + exact_auto = state.tn.compute_local_expectation_exact( + terms, + normalized=True, + optimize="auto-hq", + ) + exact_greedy = state.tn.compute_local_expectation_exact( + terms, + normalized=True, + optimize="greedy", + ) + assert exact_auto == pytest.approx(jw_value, rel=1e-10, abs=1e-10) + assert exact_greedy == pytest.approx(exact_auto, rel=1e-10, abs=1e-10) + + su = gauge_all( + state.tn, + start="su", + target="su", + norm="2norm", + su_options={"max_iterations": 8, "tol": 0.0}, + ) + assert all( + type(gauge).__module__.startswith("symmray") + for gauge in su.gauges.values() + ) + su_reconstructed = su.core.copy() + su_reconstructed.gauge_simple_insert(su.gauges) + su_exact = su_reconstructed.compute_local_expectation_exact( + terms, + normalized=True, + optimize="auto-hq", + ) + assert su_exact == pytest.approx(exact_auto, rel=1e-10, abs=1e-10) + + bp_options = { + "run_opts": { + "max_iterations": 150, + "tol": 1e-10, + "diis": False, + } + } + bridge = gauge_all( + state.tn, + start="bp", + target="su", + norm="2norm", + bp_options=bp_options, + conversion_options={"smudge": 1e-12}, + ) + assert bridge.bp_result.converged + expected_message_type = ( + "U1FermionicArray" if symmetry == "U1" else "U1U1FermionicArray" + ) + assert all( + type(message).__name__ == expected_message_type + for message in bridge.messages.values() + ) + bp_reconstructed = bridge.core.copy() + bp_reconstructed.gauge_simple_insert(bridge.gauges) + bp_exact = bp_reconstructed.compute_local_expectation_exact( + terms, + normalized=True, + optimize="auto-hq", + ) + assert bp_exact == pytest.approx(exact_auto, rel=1e-10, abs=1e-10) + + cluster = compute_path_cluster_expectation( + bridge.core, + terms, + gauges=bridge.gauges, + max_distance=1, + fillin=True, + max_bond=None, + normalized=True, + optimize="auto-hq", + ) + bp_helper = compute_bp_path_expectation( + state.tn, + terms, + max_distance=1, + fillin=True, + max_bond=None, + normalized=True, + optimize="auto-hq", + bp_options=bp_options, + conversion_options={"smudge": 1e-12}, + ) + assert cluster == pytest.approx(exact_auto, rel=1e-10, abs=1e-10) + assert bp_helper == pytest.approx(exact_auto, rel=1e-10, abs=1e-10) + + +@pytest.mark.parametrize("symmetry", ("U1", "U1U1")) +def test_spinful_eta_pair_measurement_survives_su_and_bp_gauges(symmetry): + """Native eta-pair measurements agree with JW through both gauge routes.""" + fermion = Fermion(spinful=True, symmetry=symmetry) + doublon = 2 if symmetry == "U1" else (1, 1) + empty = 0 if symmetry == "U1" else (0, 0) + occupations = { + (0, 0): doublon, + (0, 1): doublon, + (1, 0): empty, + (1, 1): empty, + } + peps = ps_to_peps( + 2, + 2, + fermion=fermion, + occupations=occupations, + dtype="complex128", + ) + state = SymPEPS( + peps=peps, + symmetry=symmetry, + edges=tuple(qtn.edges_2d_square(2, 2)), + fermionic=True, + phys_sectors=fermion.physical_sectors, + site_charge=occupations, + site_ind_id="k{},{}", + ) + where = ((0, 0), (1, 1)) + operator = fermion.eta_pair_operator() + terms = {where: operator} + dt = 0.17 + state.apply_gates( + ((fermion.operator_gate(operator, dt, imaginary=True), where),), + method="gate", + max_bond=16, + cutoff=0.0, + ) + + initial = np.zeros(256, dtype="complex128") + initial[np.ravel_multi_index((1, 1, 1, 1, 0, 0, 0, 0), (2,) * 8)] = 1.0 + jw_operator = _jw_eta_pair_operator(0, 3) + dense_state = _imaginary_time_from_hermitian(jw_operator, dt) @ initial + jw_value = ( + dense_state.conj() @ jw_operator @ dense_state + ) / (dense_state.conj() @ dense_state) + assert abs(jw_value) > 1.0e-6 + + exact_auto = state.tn.compute_local_expectation_exact( + terms, + normalized=True, + optimize="auto-hq", + ) + exact_greedy = state.tn.compute_local_expectation_exact( + terms, + normalized=True, + optimize="greedy", + ) + assert exact_auto == pytest.approx(jw_value, rel=1e-10, abs=1e-10) + assert exact_greedy == pytest.approx(exact_auto, rel=1e-10, abs=1e-10) + + su = gauge_all( + state.tn, + start="su", + target="su", + norm="2norm", + su_options={"max_iterations": 8, "tol": 0.0}, + ) + su_reconstructed = su.core.copy() + su_reconstructed.gauge_simple_insert(su.gauges) + su_exact = su_reconstructed.compute_local_expectation_exact( + terms, + normalized=True, + optimize="auto-hq", + ) + assert su_exact == pytest.approx(exact_auto, rel=1e-10, abs=1e-10) + + bp_options = { + "run_opts": { + "max_iterations": 150, + "tol": 1e-10, + "diis": False, + } + } + bridge = gauge_all( + state.tn, + start="bp", + target="su", + norm="2norm", + bp_options=bp_options, + conversion_options={"smudge": 1e-12}, + ) + assert bridge.bp_result.converged + bp_reconstructed = bridge.core.copy() + bp_reconstructed.gauge_simple_insert(bridge.gauges) + bp_exact = bp_reconstructed.compute_local_expectation_exact( + terms, + normalized=True, + optimize="auto-hq", + ) + assert bp_exact == pytest.approx(exact_auto, rel=1e-10, abs=1e-10) + + cluster = compute_path_cluster_expectation( + bridge.core, + terms, + gauges=bridge.gauges, + max_distance=1, + fillin=True, + max_bond=None, + normalized=True, + optimize="auto-hq", + ) + bp_helper = compute_bp_path_expectation( + state.tn, + terms, + max_distance=1, + fillin=True, + max_bond=None, + normalized=True, + optimize="auto-hq", + bp_options=bp_options, + conversion_options={"smudge": 1e-12}, + ) + assert cluster == pytest.approx(exact_auto, rel=1e-10, abs=1e-10) + assert bp_helper == pytest.approx(exact_auto, rel=1e-10, abs=1e-10) + + +def test_fermionic_long_range_boundary_expectation_matches_exact(): + """The boundary route preserves a distant native density operator.""" + state = SymPEPS.random( + 2, + 2, + symmetry="U1", + bond_dim=3, + phys_dim=4, + fermionic=True, + seed=1500, + dtype="complex128", + ) + terms = _long_range_density_term(((0, 1), (1, 0))) + + exact = state.tn.compute_local_expectation_exact( + terms, + normalized=True, + optimize="auto-hq", + ) + measured = compute_boundary_expectation( + state.tn, + terms, + max_bond=8, + mode="mps", + normalized=True, + contract_optimize="auto-hq", + ) + + assert measured == pytest.approx(exact, rel=1e-10, abs=1e-10) + + +def test_fermionic_path_cluster_accepts_native_su_gauges(): + """A distant fermionic density cluster accepts native SU bond vectors.""" + state = SymPEPS.random( + 3, + 3, + symmetry="U1", + bond_dim=2, + phys_dim=4, + fermionic=True, + seed=1501, + dtype="complex128", + ) + where = ((0, 1), (2, 1)) + terms = _long_range_density_term(where) + exact = state.tn.compute_local_expectation_exact( + terms, + normalized=True, + optimize="auto-hq", + ) + norm_before = complex(state.tn.norm()) + + su = gauge_all( + state.tn, + start="su", + target="su", + norm="2norm", + su_options={"max_iterations": 8, "tol": 0.0}, + ) + assert su.gauges + assert all( + type(gauge).__module__.startswith("symmray") + for gauge in su.gauges.values() + ) + + ungauged = compute_path_cluster_expectation( + su.core, + terms, + max_distance=0, + fillin=False, + max_bond=None, + normalized=True, + optimize="auto-hq", + ) + measured = compute_path_cluster_expectation( + su.core, + terms, + max_distance=0, + fillin=False, + gauges=su.gauges, + max_bond=None, + normalized=True, + optimize="auto-hq", + ) + assert np.isfinite(measured) + assert measured != pytest.approx(ungauged, rel=1e-6, abs=1e-10) + + full_cluster = compute_path_cluster_expectation( + su.core, + terms, + max_distance=1, + fillin=True, + gauges=su.gauges, + max_bond=None, + normalized=True, + optimize="auto-hq", + ) + + assert full_cluster == pytest.approx(exact, rel=1e-10, abs=1e-10) + assert complex(state.tn.norm()) == pytest.approx(norm_before) + + +@pytest.mark.parametrize("symmetry", ("U1", "U1U1")) +def test_fermionic_path_cluster_compression_preserves_native_symmetry(symmetry): + """Native path clusters support Symmray-aware QR/SVD compression.""" + site_charge = None + if symmetry == "U1U1": + site_charge = site_charge_alternating((1, 0), (0, 1)) + state = SymPEPS.random( + 3, + 3, + symmetry=symmetry, + bond_dim=2, + phys_dim=4, + fermionic=True, + site_charge=site_charge, + seed=1620, + dtype="complex128", + ) + where = ((0, 1), (2, 1)) + terms = _long_range_density_term(where, symmetry=symmetry) + su = gauge_all( + state.tn, + start="su", + target="su", + norm="2norm", + su_options={"max_iterations": 8, "tol": 0.0}, + ) + + exact_cluster = compute_path_cluster_expectation( + su.core, + terms, + gauges=su.gauges, + max_distance=1, + fillin=True, + max_bond=None, + optimize="auto-hq", + ) + with pytest.warns(UserWarning, match="not a compressed one"): + high_chi = compute_path_cluster_expectation( + su.core, + terms, + gauges=su.gauges, + max_distance=1, + fillin=True, + max_bond=8, + optimize="auto-hq", + ) + with pytest.warns(UserWarning, match="not a compressed one"): + low_chi = compute_path_cluster_expectation( + su.core, + terms, + gauges=su.gauges, + max_distance=1, + fillin=True, + max_bond=2, + optimize="auto-hq", + ) + + assert high_chi == pytest.approx(exact_cluster, rel=1e-10, abs=1e-10) + assert np.isfinite(low_chi) + assert low_chi != pytest.approx(high_chi, rel=1e-6, abs=1e-10) + assert all( + type(tensor.data).__name__ + == ("U1FermionicArray" if symmetry == "U1" else "U1U1FermionicArray") + for tensor in su.core.tensors + ) + + cotengra_optimizer = ctg.HyperOptimizer( + max_repeats=2, + parallel=False, + progbar=False, + ) + with pytest.warns(UserWarning, match="not a compressed one"): + via_cotengra = compute_path_cluster_expectation( + su.core, + terms, + gauges=su.gauges, + max_distance=1, + fillin=True, + max_bond=2, + optimize=cotengra_optimizer, + ) + assert np.isfinite(via_cotengra) + + +def test_fermionic_bp_path_expectation_uses_native_bp_to_su_bridge(): + """BP-derived SU gauges close a distant fermionic density cluster.""" + state = SymPEPS.random( + 3, + 3, + symmetry="U1", + bond_dim=2, + phys_dim=4, + fermionic=True, + seed=1502, + dtype="complex128", + ) + where = ((0, 1), (2, 1)) + terms = _long_range_density_term(where) + norm_before = complex(state.tn.norm()) + bp_options = { + "run_opts": { + "max_iterations": 150, + "tol": 1e-10, + "diis": False, + } + } + conversion_options = {"smudge": 1e-12} + + bridge = gauge_all( + state.tn, + start="bp", + target="su", + norm="2norm", + bp_options=bp_options, + conversion_options=conversion_options, + ) + assert bridge.bp_result.converged + assert all( + type(message).__name__ == "U1FermionicArray" + for message in bridge.messages.values() + ) + via_bridge = compute_path_cluster_expectation( + bridge.core, + terms, + max_distance=0, + fillin=False, + gauges=bridge.gauges, + max_bond=None, + normalized=True, + optimize="auto-hq", + ) + + measured = compute_bp_path_expectation( + state.tn, + terms, + max_distance=0, + fillin=False, + max_bond=None, + normalized=True, + optimize="auto-hq", + bp_options=bp_options, + conversion_options=conversion_options, + ) + + assert np.isfinite(via_bridge) + assert measured == pytest.approx(via_bridge, rel=1e-10, abs=1e-10) + + with pytest.warns(UserWarning, match="not a compressed one"): + compressed_via_bridge = compute_path_cluster_expectation( + bridge.core, + terms, + max_distance=1, + fillin=True, + gauges=bridge.gauges, + max_bond=2, + normalized=True, + optimize="auto-hq", + ) + with pytest.warns(UserWarning, match="not a compressed one"): + compressed_measured = compute_bp_path_expectation( + state.tn, + terms, + max_distance=1, + fillin=True, + max_bond=2, + normalized=True, + optimize="auto-hq", + bp_options=bp_options, + conversion_options=conversion_options, + ) + assert np.isfinite(compressed_via_bridge) + assert compressed_measured == pytest.approx( + compressed_via_bridge, + rel=1e-10, + abs=1e-10, + ) + assert complex(state.tn.norm()) == pytest.approx(norm_before) + assert all( + type(tensor.data).__name__ == "U1FermionicArray" + for tensor in state.tn.tensors + ) + + @pytest.mark.parametrize("bond_dim", (2, 3)) def test_fermionic_u1_two_norm_bp_is_exact_on_a_tree(bond_dim): """D2BP preserves the native graded contraction on a PEPS tree.""" @@ -314,6 +1141,288 @@ def test_fermionic_u1_d2bp_corrections_and_relay_use_native_messages(): ) +def test_partial_trace_loop_series_matches_quimb_on_dense_peps(): + """The Pepsy local rho wrapper matches Quimb's dense D2BP reference.""" + state = qtn.PEPS.rand( + 2, + 2, + bond_dim=2, + seed=1901, + dtype="complex128", + ) + where = ((0, 0), (1, 1)) + + rho = partial_trace_loop_series_expand( + state, + where, + gloops=2, + max_iterations=200, + tol=1e-10, + diis=False, + ) + reference_bp = two_norm_bp( + state, + max_iterations=200, + tol=1e-10, + diis=False, + ) + reference = reference_bp.bp.partial_trace_loop_series_expansion( + where=where, + gloops=2, + normalized=True, + ) + + np.testing.assert_allclose(rho, reference, rtol=1e-10, atol=1e-12) + np.testing.assert_allclose(np.trace(rho), 1.0, rtol=1e-10, atol=1e-12) + + gate = np.diag([1.0, -1.0, -1.0, 1.0]) + value = compute_local_expectation_loop_series( + state, + {where: gate}, + gloops=2, + normalized="prod", + max_iterations=200, + tol=1e-10, + diis=False, + ) + np.testing.assert_allclose(value, np.trace(reference @ gate)) + + cluster_rho = partial_trace_loop_cluster_expand( + state, + where, + gloops=2, + max_iterations=200, + tol=1e-10, + diis=False, + ) + cluster_reference = reference_bp.bp.partial_trace_gloop_expand( + where, + gloops=2, + combine="sum", + normalized=True, + ) + np.testing.assert_allclose( + cluster_rho, + cluster_reference, + rtol=1e-10, + atol=1e-12, + ) + cluster_value = compute_local_expectation_loop_cluster( + state, + {where: gate}, + gloops=2, + max_iterations=200, + tol=1e-10, + diis=False, + ) + np.testing.assert_allclose(cluster_value, np.trace(cluster_reference @ gate)) + + +def test_fermionic_partial_trace_loop_series_keeps_native_rho(): + """Local D2 loop-series rho keeps fermionic Symmray block structure.""" + state = SymPEPS.random( + 2, + 2, + symmetry="U1", + bond_dim=3, + phys_dim=2, + fermionic=True, + seed=1903, + dtype="complex128", + ) + where = ((0, 0), (1, 1)) + + rho = partial_trace_loop_series_expand( + state.tn, + where, + gloops=2, + max_iterations=200, + tol=1e-10, + diis=False, + ) + + assert type(rho).__name__ == "U1FermionicArray" + assert rho.ndim == 2 + np.testing.assert_allclose( + np.trace(rho.to_dense()), + 1.0, + rtol=1e-10, + atol=1e-12, + ) + + cluster_rho = partial_trace_loop_cluster_expand( + state.tn, + where, + gloops=2, + max_iterations=200, + tol=1e-10, + diis=False, + ) + assert type(cluster_rho).__name__ == "U1FermionicArray" + np.testing.assert_allclose( + np.trace(cluster_rho.to_dense()), + 1.0, + rtol=1e-10, + atol=1e-12, + ) + + +def test_fermionic_local_expectation_loop_series_aligns_charge_support(): + """Graded gate insertion is exact on a fermionic D2BP tree.""" + state = SymPEPS.random( + 1, + 4, + symmetry="U1", + bond_dim=3, + phys_dim=4, + fermionic=True, + seed=1904, + dtype="complex128", + ) + where = ((0, 1), (0, 2)) + gate = Fermion(spinful=True, symmetry="U1").eta_pair_operator() + exact = state.tn.compute_local_expectation_exact( + {where: gate}, + normalized=True, + optimize="auto-hq", + ) + series_value = compute_local_expectation_loop_series( + state.tn, + {where: gate}, + gloops=0, + max_iterations=200, + tol=1e-10, + diis=False, + ) + cluster_value = compute_local_expectation_loop_cluster( + state.tn, + {where: gate}, + gloops=0, + max_iterations=200, + tol=1e-10, + diis=False, + ) + + np.testing.assert_allclose(series_value, exact, rtol=1e-10, atol=1e-12) + np.testing.assert_allclose(cluster_value, exact, rtol=1e-10, atol=1e-12) + + reverse_where = where[::-1] + reverse_gate = gate.transpose((1, 0, 3, 2)) + reverse_exact = state.tn.compute_local_expectation_exact( + {reverse_where: reverse_gate}, + normalized=True, + optimize="auto-hq", + ) + reverse_series = compute_local_expectation_loop_series( + state.tn, + {reverse_where: reverse_gate}, + gloops=0, + max_iterations=200, + tol=1e-10, + diis=False, + ) + np.testing.assert_allclose( + reverse_series, + reverse_exact, + rtol=1e-10, + atol=1e-12, + ) + + +def test_explicit_edge_loop_series_uses_fermion_safe_gate_path(): + """The explicit-edge scalar API is exact on a fermionic tree.""" + state = SymPEPS.random( + 1, + 4, + symmetry="U1", + bond_dim=3, + phys_dim=4, + fermionic=True, + seed=1905, + dtype="complex128", + ) + where = ((0, 2),) + gate = Fermion(spinful=True, symmetry="U1").chemical_potential_operator() + exact = state.tn.compute_local_expectation_exact( + {where: gate}, normalized=True, optimize="auto-hq" + ) + value = compute_local_expectation_edge_loop_series( + state.tn, + {where: gate}, + gloops=0, + max_iterations=200, + tol=1e-10, + diis=False, + ) + rho = partial_trace_edge_loop_series_expand( + state.tn, + where, + gloops=0, + max_iterations=200, + tol=1e-10, + diis=False, + ) + + np.testing.assert_allclose(value, exact, rtol=1e-10, atol=1e-12) + assert type(rho).__name__ == "U1FermionicArray" + np.testing.assert_allclose(np.trace(rho.to_dense()), 1.0) + + +def test_explicit_edge_loop_series_preserves_dense_edge_degree_terms(): + """Edge-degree terms are distinct from the local-region cutoff API.""" + state = qtn.PEPS.rand(2, 2, bond_dim=2, seed=1906, dtype="complex128") + where = ((0, 0),) + gate = np.diag([1.0, -1.0]) + info = {} + rho = partial_trace_edge_loop_series_expand( + state, + where, + gloops=4, + max_iterations=200, + tol=1e-10, + diis=False, + info=info, + ) + value = compute_local_expectation_edge_loop_series( + state, + {where: gate}, + gloops=4, + max_iterations=200, + tol=1e-10, + diis=False, + ) + + assert info["edge_rho_terms"] + assert all(len(term.edges) <= 4 for term in info["edge_rho_terms"]) + np.testing.assert_allclose(value, np.trace(rho @ gate)) + + +def test_explicit_edge_loop_series_rejects_multisite_fermionic_q_terms(): + """Unsupported graded multi-site Q terms fail before contraction.""" + state = SymPEPS.random( + 2, + 2, + symmetry="U1", + bond_dim=2, + phys_dim=4, + fermionic=True, + seed=1907, + dtype="complex128", + ) + fermion = Fermion(spinful=True, symmetry="U1") + where = ((0, 0), (1, 1)) + gate = fermion.eta_pair_operator() + with pytest.raises(NotImplementedError, match="multi-site gates"): + compute_local_expectation_edge_loop_series( + state.tn, + {where: gate}, + gloops=4, + max_iterations=200, + tol=1e-10, + diis=False, + ) + + def _closed_u1_scalar_network(): """Build a small native Symmray scalar network for the D1 compatibility path.""" maps = ({0: 0, 1: 1},) * 2 From 50b46b889dcb10b1118f68841377bccbb91a84a9 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Thu, 30 Jul 2026 12:08:11 -0600 Subject: [PATCH 31/70] Update fermion and VMC compatibility --- .../skills/pepsy-fermion-operators/SKILL.md | 13 +- .../references/design.md | 53 +++++---- docs/api/tensors/symmetric.md | 2 +- docs/api/vmc.md | 20 ++-- src/pepsy/optimizers/mps/optimizer.py | 8 ++ src/pepsy/tensors/symmetric.py | 4 +- src/pepsy/vmc/netket.py | 85 ++++++++----- tests/test_netket_flat_z2.py | 112 ++++++++++++++++++ 8 files changed, 226 insertions(+), 71 deletions(-) diff --git a/.github/skills/pepsy-fermion-operators/SKILL.md b/.github/skills/pepsy-fermion-operators/SKILL.md index c48a15d..056f5fd 100644 --- a/.github/skills/pepsy-fermion-operators/SKILL.md +++ b/.github/skills/pepsy-fermion-operators/SKILL.md @@ -50,13 +50,20 @@ Keep one model-facing object responsible for: - matching `SymHamiltonian` construction for spinful Hubbard or spinless t-V; - occupation/charge helpers and model metadata. -Prefer an explicit constructor such as: +Prefer an explicit local-space constructor, with physical couplings supplied +at the term, Hamiltonian, or gate-stream call site: ```python -Fermion(spinful=False, symmetry="U1", t=1.0, V=0.5) -Fermion(spinful=True, symmetry="U1U1", t=1.0, U=8.0) +spinless = Fermion(spinful=False, symmetry="U1") +spinful = Fermion(spinful=True, symmetry="U1U1") + +spinless.hamiltonian(edges, t=1.0, V=0.5) +spinful.strang_gate_stream(edges, dt=0.01, t=1.0, U=8.0) ``` +`Fermion` must not store `t`, `U`, `V`, or `mu`. Reject a Hubbard `U` on a +spinless object instead of silently dropping it. + Do not hide the spinless/spinful choice behind a model-name string when a boolean or explicit local-space descriptor can make it clear. Model-name strings may remain accepted for compatibility and Hamiltonian dispatch. diff --git a/.github/skills/pepsy-fermion-operators/references/design.md b/.github/skills/pepsy-fermion-operators/references/design.md index 31d39b7..516f5ad 100644 --- a/.github/skills/pepsy-fermion-operators/references/design.md +++ b/.github/skills/pepsy-fermion-operators/references/design.md @@ -6,24 +6,20 @@ Symmray implementation. ## Current Pepsy state -Pepsy currently exposes `SpinfulFermion` from `pepsy` and `pepsy.tensors`. The -implementation lives in `src/pepsy/tensors/symmetric.py`, is re-exported by -`src/pepsy/tensors/symm_fermions.py`, and currently provides: +Pepsy exposes the canonical `Fermion` helper from `pepsy` and `pepsy.tensors`. +The implementation lives in `src/pepsy/tensors/symmetric.py`, is re-exported +by `src/pepsy/tensors/symm_fermions.py`, and provides: -- spinful `U1` and `U1U1` local spaces; +- spinless `U1`/`Z2` and spinful `U1`/`Z2`/`U1U1`/`Z2Z2` local spaces; - dense local operators and aliases; - native Symmray one-site observables; -- onsite interaction and two-site hopping gates; -- deterministic second-order edge-colored gate streams; -- `SymHamiltonian` construction for `fermi_hubbard` and - `fermi_hubbard_u1u1`; -- `SpinfulFermionHubbard = SpinfulFermion` compatibility alias. - -Pepsy already supports spinless model metadata and Hamiltonian paths through -`fermi_hubbard_spinless`, but the model-facing helper does not yet unify that -space with the spinful helper. The intended evolution is a canonical -`Fermion(spinful=...)` helper, retaining the existing spinful names as -compatibility aliases. +- onsite, hopping, density, field, and parity-preserving pairing gates; +- deterministic first- and second-order edge-colored gate streams; +- `SymHamiltonian` construction for spinful Hubbard and spinless t-V models. + +`SpinfulFermion` and `SpinfulFermionHubbard` are compatibility constructors: +they deliberately fix `spinful=True`, while `SymmFermions.spinful(...)` and +`SymmFermions.spinless(...)` provide corresponding factories. ## Symmray local conventions @@ -68,15 +64,15 @@ Use one object with an explicit local-space switch. The exact name can change only with user direction; `Fermion` is the recommended canonical name. ```python -spinless = Fermion(spinful=False, symmetry="U1", t=1.0, V=0.5, mu=0.0) -spinful = Fermion(spinful=True, symmetry="U1U1", t=1.0, U=8.0, mu=0.0) +spinless = Fermion(spinful=False, symmetry="U1") +spinful = Fermion(spinful=True, symmetry="U1U1") spinless.dense_operator("number") spinful.observable("number_up") -spinful.hopping_gate(dt=0.01) -spinful.interaction_gate(dt=0.01, site=3) -spinful.strang_gate_stream(edges, dt=0.01, sites=range(L)) -spinful.hamiltonian(edges) +spinful.hopping_gate(dt=0.01, t=1.0) +spinful.interaction_gate(dt=0.01, site=3, U=8.0) +spinful.strang_gate_stream(edges, dt=0.01, sites=range(L), t=1.0, U=8.0) +spinful.hamiltonian(edges, t=1.0, U=8.0, mu=0.0) ``` Recommended metadata and methods: @@ -96,7 +92,10 @@ Recommended metadata and methods: | `strang_gate_stream(...)` | create deterministic canonical bundled entries | The spinless parameter should use `V` for nearest-neighbor density interaction; -do not silently interpret a spinless `U` as a doublon interaction. Spinless +do not silently interpret a spinless `U` as a doublon interaction: the public +methods reject it. `Fermion` stores only local-space and backend metadata; +pass `t`, `U`, `V`, and `mu` to a term, Hamiltonian, or gate-stream call. +Spinless pairing or superconducting terms should be added only with explicit symmetry and charge semantics, since they break particle-number conservation. @@ -170,14 +169,16 @@ once. This is correct for native local-term measurements, but the one-site interaction is not a separate dictionary entry. For examples that prioritize visible bookkeeping, build hopping-only edge -terms with `U=0` and `mu=0`, then add one-site native observables: +terms with `U=0` and `mu=0`, then add one-site native observables. Keep the +couplings as ordinary local variables rather than Fermion attributes: ```python -hop = fermion.hamiltonian(edges, U=0.0, mu=0.0) +U, mu = 8.0, 0.0 +hop = fermion.hamiltonian(edges, t=1.0, U=0.0, mu=0.0) energy_terms = dict(hop.terms) onsite = ( - fermion.U * fermion.observable("double") - - fermion.mu * fermion.observable("number") + U * fermion.observable("double") + - mu * fermion.observable("number") ) energy_terms.update({(site,): onsite for site in range(L)}) ``` diff --git a/docs/api/tensors/symmetric.md b/docs/api/tensors/symmetric.md index d6fc201..f55fb34 100644 --- a/docs/api/tensors/symmetric.md +++ b/docs/api/tensors/symmetric.md @@ -288,7 +288,7 @@ two-state mode layout. This returns qMERA ``LocalTerm`` objects rather than four-state site tensors, matching ``QMeraGeometry(site_modes=("up", "down"))``: ```python -from pepsy.optimizers.mera import QMeraGeometry +from pepsy.optimizers.qmera import QMeraGeometry geometry = QMeraGeometry(shape=3, site_modes=("up", "down")) qmera_terms = spinful.local_terms(geometry, layout="qmera") diff --git a/docs/api/vmc.md b/docs/api/vmc.md index 472ff6d..b0dd0e0 100644 --- a/docs/api/vmc.md +++ b/docs/api/vmc.md @@ -75,8 +75,9 @@ For fermions, the supported JIT configuration intentionally combines a flat `SpinOrbitalFermions` sampling sector. The `Z2` label describes the JAX-friendly tensor storage; it does not relax the sampler's separate particle number conservation. The fermion builders emit `SymmetryFallbackWarning` to -make this distinction visible. A block-sparse `U1U1` PEPS still cannot enter -NetKet's jitted `MCState` until Symmray supplies a flat `U1U1` backend. +make this distinction visible. Block-sparse `U1`, `U1U1`, and `Z2Z2` PEPS still +cannot enter NetKet's jitted `MCState` until Symmray supplies matching flat +fermionic backends. ```python import pepsy.vmc as pvmc @@ -1256,12 +1257,13 @@ native sparse U1U1 and exact contraction; approximate boundary contractions continue to use the established sparse path. The same option is forwarded by `TorchFermionVMC` when its contraction is exact. -For `U1U1`, `fermionic_peps_rand("U1U1", ...)` builds the block-sparse ansatz -and `make_fermionic_peps_batched_amplitude_function(..., jit=False)` is -validated with `contraction="exact"`, `"hotrg"`, `"ctmrg"`, and -`"boundary"`/`"mps"`. Full NetKet `MCState` VMC still requires a jitted Flax -model, so `build_fermi_hubbard_vmc(...)` raises clearly for block-sparse -`U1U1` PEPS until Symmray provides a flat U1U1 fermionic backend. +For `U1`, `U1U1`, and `Z2Z2`, `fermionic_peps_rand(...)` builds the +block-sparse ansatz and `make_fermionic_peps_batched_amplitude_function(..., +jit=False)` is validated with `contraction="exact"`, `"hotrg"`, `"ctmrg"`, +and `"boundary"`/`"mps"`. Full NetKet `MCState` VMC still requires a jitted +Flax/JAX model, so `build_fermi_hubbard_vmc(...)` raises clearly for these +block-sparse PEPS until Symmray provides a flat backend for the requested +symmetry. For an actual fixed-sector sparse-block VMC loop, use `build_sparse_fermi_hubbard_vmc(...)`: it builds the NetKet `SpinOrbitalFermions` Hilbert space and Fermi-Hubbard operator metadata, then @@ -1273,7 +1275,7 @@ sectors it can enumerate the NetKet Hilbert space, while larger runs can pass ```python peps = pvmc.fermionic_peps_rand( - "U1U1", + "U1", # or "U1U1" for separately fixed spin sectors Lx=2, Ly=2, bond_dim=3, diff --git a/src/pepsy/optimizers/mps/optimizer.py b/src/pepsy/optimizers/mps/optimizer.py index 64aa4d1..ac65fa7 100644 --- a/src/pepsy/optimizers/mps/optimizer.py +++ b/src/pepsy/optimizers/mps/optimizer.py @@ -2788,6 +2788,14 @@ def _to_state_backend(self, array): source_signature = _array_backend_signature(array) if source_signature == target_signature: return array + if self._is_symmray_array(array) and self._is_symmray_array(like): + # Symmray arrays deliberately do not implement Autoray's generic + # ``array(..., like=symmray_array)`` constructor. Their outer + # object has no scalar dtype either, so the generic dtype fast + # path below cannot establish compatibility. Native Symmray gates + # already carry their own block backend and must pass through as + # graded arrays rather than being rebuilt as dense payloads. + return array if target_signature[0] == "symmray" and source_signature[0] != "symmray": raise TypeError( "Cannot convert a dense gate/operator payload into a native " diff --git a/src/pepsy/tensors/symmetric.py b/src/pepsy/tensors/symmetric.py index bab2e0e..fbcf74a 100644 --- a/src/pepsy/tensors/symmetric.py +++ b/src/pepsy/tensors/symmetric.py @@ -9890,7 +9890,7 @@ def local_terms(self, edges, *, layout="site", **params): if layout in {"site", "sites", "native"}: return self.hamiltonian(edges, **params).terms if layout in {"qmera", "qmera_modes", "modes"}: - from ..optimizers.mera import ( # pylint: disable=import-outside-toplevel + from ..optimizers.qmera import ( # pylint: disable=import-outside-toplevel qmera_symmray_fermi_hubbard_terms, ) @@ -9900,7 +9900,7 @@ def local_terms(self, edges, *, layout="site", **params): **params, ) if layout in {"majorana", "qmera_majorana"}: - from ..optimizers.mera import ( # pylint: disable=import-outside-toplevel + from ..optimizers.qmera import ( # pylint: disable=import-outside-toplevel qmera_symmray_majorana_terms, ) diff --git a/src/pepsy/vmc/netket.py b/src/pepsy/vmc/netket.py index baf31fb..c3efe5a 100644 --- a/src/pepsy/vmc/netket.py +++ b/src/pepsy/vmc/netket.py @@ -145,6 +145,7 @@ class PackedPEPS: site_inds: tuple[Any, ...] = () uses_flat_symmray: bool | None = None phys_charges: tuple[Any, ...] = () + symmray_symmetry: str | None = None @property def n_sites(self): @@ -1758,7 +1759,8 @@ class NetKetSparseFermiHubbardVMC: NetKet supplies the Hilbert space, graph, and Hamiltonian metadata. The actual samples, amplitudes, and local energies use the non-jitted torch - PEPS path so block-sparse ``U1U1`` Symmray tensors can be evaluated today. + PEPS path so block-sparse ``U1``, ``U1U1``, and ``Z2Z2`` Symmray tensors + can be evaluated today. """ hilbert: Any @@ -2265,6 +2267,18 @@ def _is_symmray_array(value): return type(value).__module__.split(".", 1)[0] == "symmray" +def _symmray_symmetry_name(tn): + """Return the first Symmray tensor's symmetry name, if available.""" + for tensor in tn: + data = getattr(tensor, "data", None) + if not _is_symmray_array(data): + continue + symmetry = getattr(data, "symmetry", None) + if symmetry is not None: + return str(symmetry) + return None + + def _uses_flat_symmray_arrays(tn): """Return whether Symmray arrays in ``tn`` use flat JAX-friendly storage.""" symmray_seen = False @@ -2334,8 +2348,8 @@ def prepare_fermionic_peps_for_netket(peps, *, device=None): blocks, and then moved to JAX. ``device=None`` follows JAX's default device; ``PEPSY_FH_JAX_DEVICE`` can be used for a notebook-level override. - Non-Z2 Symmray tensors are left intact so the existing clear U1U1 error is - raised by the jitted NetKet model. Dense, already-flat, and non-Symmray + Non-Z2 Symmray tensors are left block-sparse so the jitted NetKet model can + raise a clear capability error. Dense, already-flat, and non-Symmray tensors only go through the JAX backend conversion. """ work = peps.copy() @@ -2348,7 +2362,7 @@ def prepare_fermionic_peps_for_netket(peps, *, device=None): type_name = type(data).__name__ if ( _is_symmray_array(data) - and "Z2" in type_name + and str(getattr(data, "symmetry", "")) == "Z2" and "Flat" not in type_name ): if sr is None: @@ -2457,10 +2471,12 @@ def _require_jittable_fermionic_ansatz(ansatz): """Raise a clear error when NetKet's jitted VMC path cannot use ``ansatz``.""" if ansatz.uses_flat_symmray is not False: return - if _spinful_phys_lookup(getattr(ansatz, "phys_charges", ())) is not None: - symmetry = "U1U1" - else: - symmetry = "non-flat Symmray" + symmetry = getattr(ansatz, "symmray_symmetry", None) + if symmetry is None: + if _spinful_phys_lookup(getattr(ansatz, "phys_charges", ())) is not None: + symmetry = "U1U1 or Z2Z2" + else: + symmetry = "non-flat Symmray" raise NotImplementedError( "NetKet MCState JIT-compiles the PEPS log-amplitude model, but this " f"{symmetry} fermionic PEPS uses block-sparse Symmray arrays rather " @@ -2468,7 +2484,7 @@ def _require_jittable_fermionic_ansatz(ansatz): "make_fermionic_peps_batched_amplitude_function(..., jit=False) with " "contraction='exact', 'hotrg', 'ctmrg', or 'boundary' for validation, " "or use a flat Z2 fermionic PEPS for full NetKet VMC until Symmray " - "provides a flat U1U1 fermionic backend." + "provides flat backends for the requested symmetry." ) @@ -3113,6 +3129,7 @@ def _pack_peps_ansatz( n_params=n_params, uses_flat_symmray=_uses_flat_symmray_arrays(tn), phys_charges=_peps_phys_charges(tn), + symmray_symmetry=_symmray_symmetry_name(tn), ) @@ -4173,18 +4190,19 @@ def fermionic_peps_rand( """Build a random fermionic PEPS for VMC, symmetry-aware. ``"Z2"`` uses the flat (``jax.jit``/``vmap``-friendly) Symmray backend and a - ``phys_dim=4`` parity-resolved physical index. ``"U1U1"`` uses the - block-sparse backend (Symmray currently has no flat ``U1U1`` fermionic - array) with a per-spin ``(n_up, n_down)`` physical charge map and a default - site-charge summing to ``n_fermions_per_spin`` (half filling if omitted). + ``phys_dim=4`` parity-resolved physical index. ``"U1"``, ``"U1U1"``, and + ``"Z2Z2"`` use block-sparse backends. ``U1`` receives the total local + occupation charge, while the latter two receive a per-spin + ``(n_up, n_down)`` physical charge map. The default site-charge map sums + to ``n_fermions_per_spin`` (half filling if omitted). Note ---- - A ``U1U1`` ansatz cannot yet be driven through the NetKet Monte-Carlo state - (which JIT-compiles the model) because the flat backend is missing upstream. - The block-sparse ``U1U1`` PEPS still evaluates correctly through the - non-jitted amplitude functions (``jit=False``) for validation and exact/dense - sums. + A ``U1``, ``U1U1``, or ``Z2Z2`` ansatz cannot yet be driven through the NetKet + Monte-Carlo state (which JIT-compiles the model) because the corresponding + flat backend is missing upstream. The block-sparse PEPS still evaluates + correctly through the non-jitted amplitude functions (``jit=False``) for + validation and exact/dense sums. """ sr = _require_symmray() sym = str(symmetry).upper().replace("-", "").replace("_", "") @@ -4205,7 +4223,7 @@ def fermionic_peps_rand( **kwargs, ) - if sym == "U1U1": + if sym in {"U1", "U1U1", "Z2Z2"}: from pepsy.tensors import ( default_physical_sectors, site_charge_from_occupations, @@ -4215,26 +4233,32 @@ def fermionic_peps_rand( if n_fermions_per_spin is None: n_fermions_per_spin = (n_sites // 2, n_sites // 2) n_up, n_down = (int(x) for x in n_fermions_per_spin) - site_charge = site_charge_from_occupations( - _default_u1u1_flux_occupations(Lx, Ly, n_up, n_down) + occupations = _default_u1u1_flux_occupations( + Lx, Ly, n_up, n_down ) + if sym == "U1": + occupations = { + site: sum(charge) + for site, charge in occupations.items() + } + site_charge = site_charge_from_occupations(occupations) use_flat = False if flat == "auto" else bool(flat) if use_flat: warnings.warn( - "Symmray has no flat U1U1 fermionic backend; falling back to " + f"Symmray has no flat {sym} fermionic backend; falling back to " "block-sparse (flat=False). NetKet MC sampling JIT-compiles the " "model and needs a flat backend, so use the non-jit amplitude " - "functions for U1U1 until flat U1U1 lands upstream.", + f"functions for {sym} until a flat {sym} backend lands upstream.", RuntimeWarning, stacklevel=2, ) use_flat = False return sr.networks.PEPS_fermionic_rand( - "U1U1", + sym, Lx, Ly, bond_dim, - phys_dim=default_physical_sectors(model="fermi_hubbard_u1u1"), + phys_dim=default_physical_sectors(sym, 4), site_charge=site_charge, flat=use_flat, seed=seed, @@ -4244,7 +4268,7 @@ def fermionic_peps_rand( raise ValueError( f"Unsupported symmetry {symmetry!r} for fermionic_peps_rand; " - "use 'Z2' or 'U1U1'." + "use 'Z2', 'U1', 'Z2Z2', or 'U1U1'." ) @@ -5635,10 +5659,11 @@ def build_sparse_fermi_hubbard_vmc( """Create a sparse-block PEPS VMC setup for the Fermi-Hubbard model. This path is intended for Symmray block-sparse fermionic PEPS, including - ``U1U1`` tensors that cannot yet be used by NetKet's jitted ``MCState``. - NetKet still defines the Hilbert sector, graph, and Hamiltonian metadata, - while Pepsy's torch kernels do Metropolis sweeps and local-energy - evaluation with exact, HOTRG, CTMRG, or boundary contractions. + ``U1``, ``U1U1``, and ``Z2Z2`` tensors that cannot yet be used by NetKet's + jitted ``MCState``. It is a Pepsy VMC loop with NetKet Hilbert/graph/Hamiltonian + metadata rather than an ``nk.driver.VMC`` instance: Pepsy's torch kernels + do Metropolis sweeps and local-energy evaluation with exact, HOTRG, CTMRG, + or boundary contractions. """ nk = _require_netket() n_sites = int(Lx) * int(Ly) diff --git a/tests/test_netket_flat_z2.py b/tests/test_netket_flat_z2.py index 3fb71b7..9d8c7bb 100644 --- a/tests/test_netket_flat_z2.py +++ b/tests/test_netket_flat_z2.py @@ -86,6 +86,118 @@ def collect(entries): assert collect(actual) == pytest.approx(collect(expected)) +@pytest.mark.smoke +def test_z2z2_peps_stays_sparse_and_supports_nonjit_amplitudes(): + """Z2xZ2 follows Fermion's native symmetry, but not NetKet's JIT path.""" + pytest.importorskip("jax") + nk = pytest.importorskip("netket") + + import pepsy as py + from pepsy.vmc.netket import ( + fermionic_peps_rand, + make_fermionic_peps_batched_amplitude_function, + netket_spin_orbital_columns, + pack_fermionic_peps_ansatz, + prepare_fermionic_peps_for_netket, + ) + + fermion = py.Fermion(spinful=True, symmetry="Z2Z2") + assert fermion.physical_sectors == { + (0, 0): 1, + (0, 1): 1, + (1, 0): 1, + (1, 1): 1, + } + + with pytest.warns(RuntimeWarning, match="no flat Z2Z2"): + peps = fermionic_peps_rand( + "Z2Z2", + 2, + 2, + 2, + n_fermions_per_spin=(2, 2), + seed=17, + dtype="float32", + flat=True, + ) + + assert type(peps[(0, 0)].data).__name__ == "Z2Z2FermionicArray" + prepared = prepare_fermionic_peps_for_netket(peps) + assert type(prepared[(0, 0)].data).__name__ == "Z2Z2FermionicArray" + + ansatz = pack_fermionic_peps_ansatz(prepared, lattice_shape=(2, 2)) + assert ansatz.uses_flat_symmray is False + assert ansatz.symmray_symmetry == "Z2Z2" + + hilbert = nk.hilbert.SpinOrbitalFermions( + 4, + s=1 / 2, + n_fermions_per_spin=(2, 2), + ) + columns = netket_spin_orbital_columns(hilbert) + with pytest.raises(NotImplementedError, match="Z2Z2"): + make_fermionic_peps_batched_amplitude_function( + ansatz, + columns, + contraction="exact", + jit=True, + ) + + rows = np.zeros((1, 8), dtype=np.int8) + rows[0, list(columns.down[:2])] = 1 + rows[0, list(columns.up[:2])] = 1 + amplitude = make_fermionic_peps_batched_amplitude_function( + ansatz, + columns, + contraction="exact", + output="log", + jit=False, + )(rows) + assert np.asarray(amplitude).shape == (1,) + + +@pytest.mark.smoke +@pytest.mark.parametrize("symmetry", ["U1", "U1U1"]) +def test_sparse_u1_vmc_builder_uses_nonjit_path(symmetry): + """Total-U1 and U1xU1 use NetKet metadata with Pepsy's sparse VMC path.""" + pytest.importorskip("netket") + + import pepsy.vmc as pvmc + from pepsy.vmc import SamplingConfig + + with pytest.warns(RuntimeWarning, match=f"no flat {symmetry}"): + peps = pvmc.fermionic_peps_rand( + symmetry, + 2, + 2, + 2, + n_fermions_per_spin=(2, 2), + seed=23, + dtype="float32", + flat=True, + ) + + setup = pvmc.build_sparse_fermi_hubbard_vmc( + peps, + Lx=2, + Ly=2, + n_fermions_per_spin=(2, 2), + n_samples=2, + seed=23, + sampler_seed=29, + device="cpu", + ) + assert setup.ansatz.symmray_symmetry == symmetry + assert type(setup.model).__name__ == "TorchPEPSAmplitude" + assert np.isfinite(setup.energy_estimate().detach().cpu().numpy()) + + samples = setup.sample( + SamplingConfig(n_samples_per_chain=2, n_chains=2, burn_in=0, thin=1) + ) + assert samples.configs.shape == (2, 2, 4) + assert np.isfinite(setup.energy_estimate().detach().cpu().numpy()) + + @pytest.mark.smoke def test_warmup_summary_reports_stage_times_and_amplitude_rows(capsys): """Warmup uses NetKet's JIT forward and gradient routes at chunk size.""" From 79fa9602a9dfa0be14c0b80c512f283e9344edb6 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Thu, 30 Jul 2026 15:26:33 -0700 Subject: [PATCH 32/70] improve layout visualization diagnostics --- docs/api/optimizers/mps.md | 27 + docs/api/optimizers/tree.md | 95 ++ src/pepsy/optimizers/_layout_visualization.py | 129 +++ src/pepsy/optimizers/mps/layout.py | 262 +++++ src/pepsy/optimizers/mps/optimizer.py | 22 + src/pepsy/optimizers/tree/layout.py | 937 ++++++++++++++++++ src/pepsy/optimizers/tree/optimizer.py | 78 ++ tests/test_optimize_mps.py | 53 + tests/test_optimize_tree.py | 219 ++++ 9 files changed, 1822 insertions(+) create mode 100644 src/pepsy/optimizers/_layout_visualization.py diff --git a/docs/api/optimizers/mps.md b/docs/api/optimizers/mps.md index a9269a5..2f47f95 100644 --- a/docs/api/optimizers/mps.md +++ b/docs/api/optimizers/mps.md @@ -160,5 +160,32 @@ temporarily permutes the working MPS and restores the returned MPS to the original site order. Layout-aware replay prints a concise report by default; pass `layout_report=False` to silence it. +The layout can be inspected graphically without changing the optimizer. The +finder returns a Matplotlib `(fig, ax)` pair. The original lattice and gate +connectivity remain a light grey background, while the colored arrow chain +shows the selected MPS permutation directly. The default plot is axis-free and +does not number the background lattice; use `show_site_labels=True` and +`show_axes=True` when those annotations are useful: + +```python +finder = opt.layout_finder() +plan = finder.run(order="quality") +fig, ax = finder.plot( + plan, + site_coords={q: (q % 4, q // 4) for q in range(opt.p.nsites)}, +) +``` + +`opt.plot_layout(plan, site_coords=...)` is the equivalent convenience wrapper. +Coordinates are optional; tuple-valued site labels are interpreted as `(x, y)` +automatically, and ordinary labels fall back to a 1D line. Install the +optional `viz` profile to enable plotting. A stream-order colorbar is not shown +by default; pass `colorbar=True` only when the MPS-position scale is useful. +The default plot contains visible `0` through `last` order labels but no title, +chain sentence, or other text. The styling follows Quimb's axis-free schematic +drawings while retaining Pepsy's ordinary `(fig, ax)` return value. +Pass `show_order_labels=False` to hide the position labels, or use +`show_chain_label=True` and `show_title=True` for additional annotations. + > API details are maintained as handwritten Markdown in this page. diff --git a/docs/api/optimizers/tree.md b/docs/api/optimizers/tree.md index baaad03..9e20133 100644 --- a/docs/api/optimizers/tree.md +++ b/docs/api/optimizers/tree.md @@ -655,6 +655,101 @@ backend. The dominant lever for accuracy at fixed `chi` is the tree structure, so the finder and optimizer expose diagnostics to choose it: +The same diagnostics are available as a Cotengra-style tent plot. +`TreeLayoutFinder.plot(plan)` is the default tent view (also available as +`plot_tent(plan)`): it keeps the raw graph at the bottom and lifts +the selected hierarchy above its descendant sites: the raw lattice and gate +connectivity are gray, while circular nodes use stable scale colors. Hierarchy +edges use one uniform solid color by default; pass `edge_color=None` to restore +scale/order-colored edges. Midpoint arrows are enabled by default and indicate +the parent-to-child direction; pass `show_edge_arrows=False` to hide them. +Nearest-neighbor gate edges are not duplicated over the lattice. Supplying +`site_coords={qubit: (x, y)}` places the physical sites on an existing lattice. +It returns `(fig, ax)` and does not mutate the plan or live TTN: + +```python +finder = py.TreeLayoutFinder(gates, n=n, objective="congestion") +plan = finder.run() +fig, ax = finder.plot( + plan, + site_coords=logical_lattice_coords, + color_by="scale", + edge_color=None, + edge_cmap="turbo", + node_cmap="YlOrRd", + order=False, + show_edge_arrows=True, +) + +# For a live optimizer, the same plot is available without changing its state. +fig, ax = opt.plot_layout(site_coords=logical_lattice_coords) +``` + +The default plot is therefore the hierarchy that `TreeLayoutFinder` selected: +one hierarchy edge per parent-child connection, drawn over the physical lattice +and gate connectivity. For a background-free binary check, hide the physical +background with: + +```python +fig, ax = finder.plot( + plan, + lattice=False, + show_gate_connectivity=False, +) +``` + +This leaves only the selected hierarchy. The public tent plot intentionally +does not draw gate-by-gate route overlays. + +Use `finder.plot_rubberband(...)` for the same hierarchy in physical-lattice +rubberband form. The optional `viz` +profile provides Matplotlib. The plot uses the same visual +idea as Cotengra's circuit/rubberband views: the source interaction structure +remains visible underneath, and band color can encode either tree scale or +post-order. The default styling is axis-free, following Quimb's schematic +drawings, and the +background lattice is not numbered. Pass `show_axes=True` or +`show_site_labels=True` when those annotations are wanted. A stream-order +colorbar is hidden by default; pass `colorbar=True` when that diagnostic is +wanted. + +For a scale-invariant tree view, use `color_by="scale"`: + +```python +fig, ax = finder.plot( + plan, + site_coords=logical_lattice_coords, + color_by="scale", + edge_color=None, + edge_cmap="turbo", + node_cmap="YlOrRd", +) +``` + +Here leaves are scale zero and nodes use stable colors for their hierarchical +scale. To also color hierarchy edges by scale, pass `edge_color=None`. +Midpoint arrows show the direction from each parent to its children. The +mapping is independent of the number or order of gates; +`colorbar=True` then labels tree scale rather than gate-stream order. The plot +has no title by default; pass `show_title=True` if a title is wanted. + +For a physical-lattice view closer to Quimb's rubberband drawing, use: + +```python +fig, ax = finder.plot_rubberband( + plan, + site_coords=logical_lattice_coords, + color_by="gate", +) + +# The live optimizer exposes the same non-mutating view. +fig, ax = opt.plot_rubberband(site_coords=logical_lattice_coords) +``` + +This keeps the lattice sites and gate connectivity grey and wraps each +non-root tree cluster in a rounded, translucent band. Use `color_by="scale"` +for one stable band color per tree scale. + - `TreeLayoutFinder.report(plan=None)` summarises the physical-node geodesic lengths over the interaction graph (`score`, `max_path`, `mean_path`, `weighted_mean_path`) and compares against a balanced index tree diff --git a/src/pepsy/optimizers/_layout_visualization.py b/src/pepsy/optimizers/_layout_visualization.py new file mode 100644 index 0000000..00e5fa5 --- /dev/null +++ b/src/pepsy/optimizers/_layout_visualization.py @@ -0,0 +1,129 @@ +"""Small optional-matplotlib helpers for gate-stream layout plots.""" + +from __future__ import annotations + +from collections.abc import Mapping + + +def matplotlib_modules(): + """Import and return the Matplotlib pieces used by layout plots.""" + try: + import matplotlib.pyplot as plt # pylint: disable=import-outside-toplevel + from matplotlib import colormaps # pylint: disable=import-outside-toplevel + from matplotlib.cm import ScalarMappable # pylint: disable=import-outside-toplevel + from matplotlib.colors import Normalize # pylint: disable=import-outside-toplevel + from matplotlib.patches import FancyArrowPatch # pylint: disable=import-outside-toplevel + except ImportError as exc: # pragma: no cover - optional dependency + raise ImportError( + "Layout plotting requires matplotlib. " + "Install it with: pip install pepsy[viz]." + ) from exc + return plt, colormaps, ScalarMappable, Normalize, FancyArrowPatch + + +def resolve_site_coords(sites, site_coords=None): + """Resolve plotting coordinates for an ordered collection of sites.""" + sites = tuple(sites) + if site_coords is None: + # A tuple-valued site label is a useful zero-configuration lattice + # convention, e.g. ``(x, y)``. Otherwise the safe fallback is a 1D + # logical-site line. + if sites and all( + isinstance(site, tuple) + and len(site) == 2 + and all(isinstance(value, (int, float)) for value in site) + for site in sites + ): + return {site: (float(site[0]), float(site[1])) for site in sites} + return {site: (float(position), 0.0) for position, site in enumerate(sites)} + + if isinstance(site_coords, Mapping): + missing = [site for site in sites if site not in site_coords] + if missing: + raise ValueError( + "site_coords is missing plotting coordinates for site(s): " + f"{missing!r}." + ) + coords = {site: tuple(site_coords[site]) for site in sites} + else: + try: + values = tuple(site_coords) + except TypeError as exc: + raise TypeError( + "site_coords must be a mapping or a sequence of (x, y) pairs." + ) from exc + if len(values) != len(sites): + raise ValueError( + "a coordinate sequence must have one (x, y) pair per site." + ) + coords = dict(zip(sites, values)) + + for site, point in coords.items(): + if len(point) != 2: + raise ValueError( + f"plotting coordinate for site {site!r} must have length two." + ) + try: + coords[site] = (float(point[0]), float(point[1])) + except (TypeError, ValueError) as exc: + raise ValueError( + f"plotting coordinate for site {site!r} must be numeric." + ) from exc + return coords + + +def coordinate_lattice_edges(coords): + """Return unit horizontal/vertical edges present in site coordinates.""" + sites = tuple(coords) + edges = [] + for index, left in enumerate(sites): + x0, y0 = coords[left] + for right in sites[index + 1:]: + x1, y1 = coords[right] + manhattan = abs(x0 - x1) + abs(y0 - y1) + if abs(manhattan - 1.0) < 1.0e-9: + edges.append((left, right)) + return tuple(edges) + + +def coordinate_lattice_edge_keys(coords): + """Return unordered keys for the unit lattice edges in ``coords``.""" + return frozenset( + frozenset(edge) for edge in coordinate_lattice_edges(coords) + ) + + +def add_order_colorbar(fig, ax, colormaps, ScalarMappable, Normalize, cmap, + n_events, *, label="gate stream order"): + """Add a sequential event-order colorbar when events are present.""" + if n_events < 1: + return + normalizer = Normalize(vmin=0, vmax=max(1, n_events - 1)) + fig.colorbar( + ScalarMappable(norm=normalizer, cmap=colormaps.get_cmap(cmap)), + ax=ax, + pad=0.02, + fraction=0.046, + label=label, + ) + + +def event_color(colormaps, cmap, index, n_events): + """Return a stable sequential color for one gate-stream event.""" + normalizer = max(1, n_events - 1) + return colormaps.get_cmap(cmap)(float(index) / normalizer) + + +def scale_color(colormaps, cmap, scale, n_scales): + """Return a color for a tree scale independent of gate-stream length.""" + normalizer = max(1, n_scales - 1) + return colormaps.get_cmap(cmap)(float(scale) / normalizer) + + +def finish_schematic_axes(ax, *, title=None, margins=0.12): + """Apply the axis-free styling used by quimb's schematic drawings.""" + ax.set_axis_off() + ax.set_aspect("equal", adjustable="datalim") + ax.margins(margins) + if title is not None: + ax.set_title(title) diff --git a/src/pepsy/optimizers/mps/layout.py b/src/pepsy/optimizers/mps/layout.py index 4a85b0e..60d25c9 100644 --- a/src/pepsy/optimizers/mps/layout.py +++ b/src/pepsy/optimizers/mps/layout.py @@ -10,6 +10,14 @@ import numpy as np from ...operators.gates import _normalize_gate_entries +from .._layout_visualization import ( + coordinate_lattice_edge_keys, + coordinate_lattice_edges, + event_color, + finish_schematic_axes, + matplotlib_modules, + resolve_site_coords, +) __all__ = ["MpsGateStreamLayoutFinder"] @@ -1419,3 +1427,257 @@ def map_where(self, where, plan): def mapped_where_sequence(self, plan): """Return mapped locations for the stored stream.""" return tuple(self.map_where(where, plan) for where in self.where) + + def plot( + self, + plan=None, + *, + site_coords=None, + ax=None, + figsize=(10, 7), + cmap="turbo", + lattice=True, + show_mps_order=True, + show_chain_arrows=True, + show_order_labels=True, + show_gate_connectivity=True, + show_site_labels=False, + show_event_labels=False, + colorbar=False, + show_axes=False, + show_title=False, + show_chain_label=False, + node_size=52, + event_linewidth=1.8, + event_alpha=0.62, + ): + """Plot the logical interaction graph with the proposed MPS layout. + + The faint graph is the original lattice, with solid grey edges for + gate connectivity. The selected MPS chain is the only colored route: + its arrows run through the logical lattice in exact MPS order, and the + optional node labels show both the logical site and its MPS position. + The default presentation is axis-free, following quimb's schematic + drawing style and contains no text; set ``show_title`` or one of the + label options to add annotations, or ``show_axes=True`` to retain + Matplotlib axes. + ``site_coords`` can be a mapping from logical labels to ``(x, y)`` or + a sequence aligned with :attr:`sites`. Tuple-valued ``(x, y)`` labels + are recognized automatically; otherwise sites are drawn on a line. + + Returns + ------- + (matplotlib.figure.Figure, matplotlib.axes.Axes) + The figure and axes, ready for further customization or saving. + """ + plt, colormaps, ScalarMappable, Normalize, FancyArrowPatch = ( + matplotlib_modules() + ) + if plan is None: + plan = self.run() + if not isinstance(plan, Mapping) or "site_order" not in plan: + raise TypeError("plan must be a layout mapping returned by run().") + + created_ax = ax is None + if created_ax: + _, ax = plt.subplots(figsize=figsize) + if not show_axes: + ax.figure.subplots_adjust(left=0, right=1, bottom=0, top=1) + fig = ax.figure + coords = resolve_site_coords(self.sites, site_coords) + site_order = tuple(plan["site_order"]) + position = {site: index for index, site in enumerate(site_order)} + + # Draw the physical lattice first. This is deliberately separate from + # the gate graph so a long-range gate cannot be mistaken for an MPS + # bond or a lattice edge. + if lattice: + for left, right in coordinate_lattice_edges(coords): + x0, y0 = coords[left] + x1, y1 = coords[right] + ax.plot( + (x0, x1), + (y0, y1), + color="#d5d9de", + linewidth=1.0, + alpha=0.78, + zorder=1, + ) + + if show_gate_connectivity: + lattice_pairs = ( + coordinate_lattice_edge_keys(coords) + if lattice + else set() + ) + seen_pairs = {} + for support in self.supports: + unique = tuple(dict.fromkeys(support)) + for left, right in zip(unique, unique[1:]): + key = frozenset((left, right)) + if key in lattice_pairs: + continue + seen_pairs[key] = seen_pairs.get(key, 0) + 1 + for pair, multiplicity in seen_pairs.items(): + left, right = tuple(pair) + x0, y0 = coords[left] + x1, y1 = coords[right] + ax.plot( + (x0, x1), + (y0, y1), + color="#7e8995", + linewidth=( + 0.45 + 0.18 * min(multiplicity, 4) + + 0.1 * event_linewidth + ), + linestyle="-", + alpha=event_alpha, + zorder=2, + ) + + # The colored arrows are the MPS chain itself, not stream events. + # This is the key visual distinction: every site has exactly one + # incoming/outgoing chain edge, while the grey graph above may + # contain arbitrary gate connectivity. + if show_mps_order and site_order: + for chain_index, (left, right) in enumerate( + zip(site_order, site_order[1:]) + ): + x0, y0 = coords[left] + x1, y1 = coords[right] + sign = -1.0 if chain_index % 2 else 1.0 + radius = sign * (0.045 + 0.012 * (chain_index % 3)) + ax.add_patch( + FancyArrowPatch( + (x0, y0), + (x1, y1), + arrowstyle="-|>" if show_chain_arrows else "-", + mutation_scale=10, + connectionstyle=f"arc3,rad={radius}", + linewidth=2.65, + color=event_color( + colormaps, cmap, chain_index, len(site_order) + ), + alpha=0.88, + zorder=4, + ) + ) + + # Color logical sites by their position in the proposed MPS chain. + if self.sites: + site_values = [position[site] for site in self.sites] + scatter = ax.scatter( + [coords[site][0] for site in self.sites], + [coords[site][1] for site in self.sites], + c=site_values, + cmap=colormaps.get_cmap(cmap), + vmin=0, + vmax=max(1, len(site_order) - 1), + s=node_size, + edgecolors="#41464c", + linewidths=0.65, + zorder=5, + ) + else: + scatter = None + + if show_event_labels and self.supports: + for event_index, support in enumerate(self.supports): + support = tuple(dict.fromkeys(support)) + if len(support) < 2: + continue + left, right = support[:2] + x = (coords[left][0] + coords[right][0]) / 2.0 + y = (coords[left][1] + coords[right][1]) / 2.0 + ax.text( + x, + y, + str(event_index), + color="#59636e", + fontsize=7, + ha="center", + va="center", + zorder=8, + ) + + if show_site_labels: + for site in self.sites: + x, y = coords[site] + ax.annotate( + str(site), + (x, y), + xytext=(0, 7), + textcoords="offset points", + ha="center", + fontsize=8, + color="#41464c", + zorder=9, + ) + if show_order_labels and show_mps_order: + for site in self.sites: + x, y = coords[site] + ax.annotate( + str(position[site]), + (x, y), + xytext=(0, -15), + textcoords="offset points", + ha="center", + fontsize=8, + fontweight="bold", + color="#1f2937", + bbox={ + "boxstyle": "round,pad=0.18", + "facecolor": "white", + "edgecolor": "#9ca3af", + "linewidth": 0.55, + "alpha": 0.92, + }, + zorder=10, + ) + + if colorbar and scatter is not None: + fig.colorbar( + ScalarMappable( + norm=Normalize(vmin=0, vmax=max(1, len(site_order) - 1)), + cmap=colormaps.get_cmap(cmap), + ), + ax=ax, + pad=0.02, + fraction=0.046, + label="MPS position", + ) + + title = ( + "MPS layout finder" + + (f" — {plan['selected_order']}" if plan.get("selected_order") else "") + ) + if show_axes: + if show_title: + ax.set_title(title) + ax.set_xlabel("logical site x") + ax.set_ylabel("logical site y") + ax.set_aspect("equal", adjustable="datalim") + ax.margins(0.12) + else: + finish_schematic_axes( + ax, + title=title if show_title else None, + ) + if show_chain_label and show_mps_order and site_order: + ax.text( + 0.5, + -0.105, + "MPS chain: " + " → ".join(map(str, site_order)), + transform=ax.transAxes, + ha="center", + va="top", + fontsize=9, + color="#41464c", + ) + if show_axes and coords: + y_values = [point[1] for point in coords.values()] + if max(y_values) - min(y_values) > 0.0: + ax.set_aspect("equal", adjustable="datalim") + return fig, ax + + plot_layout = plot diff --git a/src/pepsy/optimizers/mps/optimizer.py b/src/pepsy/optimizers/mps/optimizer.py index 64aa4d1..9cf3d6f 100644 --- a/src/pepsy/optimizers/mps/optimizer.py +++ b/src/pepsy/optimizers/mps/optimizer.py @@ -1040,6 +1040,28 @@ def current_gate_stream_layout(self, *, sites=None, L=None, **kwargs): return self.layout_finder(sites=sites, L=L).run(**kwargs) + def plot_layout( + self, + plan=None, + *, + sites=None, + L=None, + layout_kwargs=None, + **plot_kwargs, + ): + """Plot the current gate-stream layout and selected MPS order. + + This is a convenience wrapper around + :meth:`MpsGateStreamLayoutFinder.plot`. It returns ``(fig, ax)`` and + does not mutate the optimizer or install the plotted layout. When + ``plan`` is omitted, the finder computes its default quality plan; + pass ``layout_kwargs`` to customize that search. + """ + finder = self.layout_finder(sites=sites, L=L) + if plan is None: + plan = finder.run(**dict(layout_kwargs or {})) + return finder.plot(plan, **plot_kwargs) + def __init__( # pylint: disable=too-many-arguments,too-many-positional-arguments self, p, diff --git a/src/pepsy/optimizers/tree/layout.py b/src/pepsy/optimizers/tree/layout.py index 0cc0435..d9f2e8b 100644 --- a/src/pepsy/optimizers/tree/layout.py +++ b/src/pepsy/optimizers/tree/layout.py @@ -40,12 +40,23 @@ _normalize_weight_mode, ) from ..mps.optimizer import _control_event_parts as _mps_control_event_parts +from .._layout_visualization import ( + add_order_colorbar, + coordinate_lattice_edge_keys, + coordinate_lattice_edges, + event_color, + finish_schematic_axes, + matplotlib_modules, + resolve_site_coords, + scale_color, +) __all__ = ["TreePlan", "TreeLayoutFinder"] _DEFAULT_MAX_ARITY = object() _DEFAULT_CHI = object() _DEFAULT_SEARCH_OPTION = object() +_DEFAULT_SCALE_MARKERS = ("o",) def _looks_like_tree_tensor_network(value): @@ -311,6 +322,25 @@ def _chi_cut_fields(plan, chi): return fields +def _tree_node_scales(plan): + """Return hierarchical scales, with leaves at scale zero.""" + scales = {} + + def visit(node): + if node in scales: + return scales[node] + children = tuple(plan.children.get(node, ())) + if not children: + scale = 0 + else: + scale = 1 + max(visit(child) for child in children) + scales[node] = scale + return scale + + visit(plan.root) + return scales + + class TreePlan: """A rooted tree over ``n`` qubits (any internal-node arity). @@ -2501,3 +2531,910 @@ def report(self, plan=None, *, include_edge_loads=True): "selected_candidate": getattr(self, "_selected_candidate", "interaction"), "candidate_scores": getattr(self, "_last_candidate_scores", {}), } + + def _plot_gate_routes( + self, + plan=None, + *, + site_coords=None, + ax=None, + figsize=(10, 8), + cmap="turbo", + color_by="gate", + scale_cmap="viridis", + scale_markers=_DEFAULT_SCALE_MARKERS, + lattice=True, + show_gate_connectivity=True, + show_gate_paths=False, + show_node_ids=False, + show_site_labels=False, + show_event_labels=False, + colorbar=False, + show_axes=False, + show_title=False, + rubberband=False, + node_size=58, + event_linewidth=2.0, + event_alpha=0.5, + tree_edge_alpha=0.38, + gate_path_curvature=0.08, + ): + """Plot a tree plan over the physical lattice and gate connectivity. + + By default this draws only the explicit TTN geometry over the optional + physical background. Pass ``show_gate_paths=True`` to add gate-stream + route overlays as a separate diagnostic layer; those routes are not + tensor legs. Pass ``rubberband=True`` for the + physical-lattice rubberband view, where each non-root tree cluster is + wrapped by a rounded translucent band. In either view, + ``color_by="scale"`` uses colors independent of gate-stream length; + the explicit tree view also uses circle markers by default (custom + marker cycles can be supplied with ``scale_markers``). When enabled, + gate-path edges are kept visually distinct: structural edges are + straight grey segments, while colored gate routes are offset by + small deterministic arcs (controlled by ``gate_path_curvature``). + No stream-order colorbar or title is shown by default. Pass + ``site_coords={qubit: (x, y)}`` to place the physical leaves on an + existing lattice; internal tree nodes are then placed above the + supplied leaves. Without coordinates, leaves use their deterministic + tree order and the plot becomes a clean rooted-tree view. The default + presentation is axis-free, following quimb's schematic drawing style; + set ``show_axes=True`` to retain Matplotlib axes. + + Returns + ------- + (matplotlib.figure.Figure, matplotlib.axes.Axes) + The figure and axes, ready for further customization or saving. + """ + plt, colormaps, ScalarMappable, Normalize, FancyArrowPatch = ( + matplotlib_modules() + ) + if plan is None: + plan = self.run() + if not isinstance(plan, TreePlan): + raise TypeError("plan must be a TreePlan returned by run().") + color_by = str(color_by).replace("-", "_").strip().lower() + color_by = {"stream": "gate", "event": "gate", "level": "scale"}.get( + color_by, color_by + ) + if color_by not in {"gate", "scale"}: + raise ValueError("color_by must be 'gate' or 'scale'.") + try: + scale_markers = tuple(scale_markers) + except TypeError as exc: + raise TypeError("scale_markers must be a non-empty sequence.") from exc + if not scale_markers: + raise ValueError("scale_markers must be a non-empty sequence.") + if rubberband: + return self.plot_rubberband( + plan, + site_coords=site_coords, + ax=ax, + figsize=figsize, + cmap=cmap, + color_by=color_by, + scale_cmap=scale_cmap, + lattice=lattice, + show_gate_connectivity=show_gate_connectivity, + show_site_nodes=True, + colorbar=colorbar, + show_axes=show_axes, + show_title=show_title, + band_alpha=event_alpha, + band_linewidth=event_linewidth, + node_size=node_size, + ) + created_ax = ax is None + if created_ax: + _, ax = plt.subplots(figsize=figsize) + if not show_axes: + ax.figure.subplots_adjust(left=0, right=1, bottom=0, top=1) + fig = ax.figure + + qubits = tuple(range(plan.n)) + supplied_coords = site_coords is not None + logical_coords = resolve_site_coords(qubits, site_coords) + leaf_order = tuple(sorted(plan.qubit_of_leaf)) + leaf_position = {node: index for index, node in enumerate(leaf_order)} + positions = {} + + if supplied_coords: + for node, qubit in plan.qubit_of_leaf.items(): + positions[node] = logical_coords[qubit] + if plan.root_qubit is not None: + # The root physical site shares the root node, so it is shown + # at the root's eventual position below. + positions[plan.root] = logical_coords[plan.root_qubit] + else: + for node, index in leaf_position.items(): + positions[node] = (float(index), 0.0) + + def place_internal(node): + if node in positions: + return positions[node] + child_points = [place_internal(child) for child in plan.children[node]] + x = sum(point[0] for point in child_points) / len(child_points) + y = max(point[1] for point in child_points) + 1.0 + positions[node] = (x, y) + return positions[node] + + place_internal(plan.root) + node_scales = _tree_node_scales(plan) + n_scales = max(node_scales.values(), default=0) + 1 + # If a root qubit was supplied, its physical coordinate should not + # flatten the structural root into the lattice. Keep the tree center + # while still recording the logical root-site label at that node. + if plan.root_qubit is not None and not supplied_coords: + positions[plan.root] = ( + positions[plan.root][0], + positions[plan.root][1], + ) + + if lattice: + for left, right in coordinate_lattice_edges(logical_coords): + x0, y0 = logical_coords[left] + x1, y1 = logical_coords[right] + ax.plot( + (x0, x1), + (y0, y1), + color="#d5d9de", + linewidth=1.0, + alpha=0.78, + zorder=1, + ) + + if show_gate_connectivity: + lattice_pairs = ( + coordinate_lattice_edge_keys(logical_coords) + if lattice + else set() + ) + for support in self.supports: + unique = tuple(dict.fromkeys(support)) + for left, right in zip(unique, unique[1:]): + if frozenset((left, right)) in lattice_pairs: + continue + x0, y0 = logical_coords[left] + x1, y1 = logical_coords[right] + ax.plot( + (x0, x1), + (y0, y1), + color="#7e8995", + linewidth=0.72, + linestyle="-", + alpha=0.62, + zorder=1, + ) + + # Draw the rooted tree underneath the gate ribbons. + for parent, children in plan.children.items(): + for child in children: + x0, y0 = positions[parent] + x1, y1 = positions[child] + ax.plot( + (x0, x1), + (y0, y1), + color="#aeb6bf", + linewidth=1.05, + alpha=tree_edge_alpha, + zorder=2, + ) + + internal = [node for node in plan.nodes() if not plan.is_leaf(node)] + leaves = list(plan.leaves()) + if color_by == "scale": + def draw_scale_nodes(nodes, size): + for scale in sorted({node_scales[node] for node in nodes}): + scale_nodes = [ + node for node in nodes if node_scales[node] == scale + ] + marker = scale_markers[scale % len(scale_markers)] + ax.scatter( + [positions[node][0] for node in scale_nodes], + [positions[node][1] for node in scale_nodes], + s=size, + marker=marker, + color=scale_color( + colormaps, scale_cmap, scale, n_scales + ), + edgecolors="#41464c", + linewidths=0.7, + zorder=5, + ) + + draw_scale_nodes(internal, node_size * 0.82) + draw_scale_nodes(leaves, node_size) + else: + if internal: + ax.scatter( + [positions[node][0] for node in internal], + [positions[node][1] for node in internal], + s=node_size * 0.82, + color="#7b8188", + edgecolors="#41464c", + linewidths=0.7, + zorder=5, + ) + if leaves: + ax.scatter( + [positions[node][0] for node in leaves], + [positions[node][1] for node in leaves], + s=node_size, + c=[plan.qubit_of_leaf[node] for node in leaves], + cmap=colormaps.get_cmap(cmap), + vmin=0, + vmax=max(1, plan.n - 1), + edgecolors="#41464c", + linewidths=0.7, + zorder=5, + ) + + n_events = len(self.supports) + if show_gate_paths: + event_weights = tuple(self.event_weights) + max_weight = max(event_weights, default=1.0) + for event_index, (support, weight) in enumerate( + zip(self.supports, event_weights) + ): + support = tuple(dict.fromkeys(support)) + event_color_value = event_color( + colormaps, cmap, event_index, n_events + ) + width = event_linewidth * ( + 0.75 + + 0.75 * (float(weight) / max(max_weight, 1.0)) ** 0.5 + ) + if gate_path_curvature: + side = 1.0 if event_index % 2 == 0 else -1.0 + magnitude = 1.0 + float((event_index // 2) % 3) + route_curvature = ( + side * float(gate_path_curvature) * magnitude + ) + else: + route_curvature = 0.0 + paths = [] + for left, right in zip(support, support[1:]): + path = plan.node_path( + plan.node_of_qubit[left], plan.node_of_qubit[right] + ) + paths.append(path) + segments = set() + for path in paths: + for left, right in zip(path, path[1:]): + edge = (left, right) if left < right else (right, left) + if edge in segments: + continue + segments.add(edge) + x0, y0 = positions[left] + x1, y1 = positions[right] + if color_by == "scale": + segment_color = scale_color( + colormaps, + scale_cmap, + max(node_scales[left], node_scales[right]), + n_scales, + ) + else: + segment_color = event_color_value + ax.add_patch( + FancyArrowPatch( + (x0, y0), + (x1, y1), + arrowstyle="-", + connectionstyle=( + f"arc3,rad={route_curvature:.4g}" + ), + linewidth=width, + color=segment_color, + alpha=event_alpha, + zorder=3, + ) + ) + if color_by == "gate": + for qubit in support: + x, y = positions[plan.node_of_qubit[qubit]] + ax.scatter( + [x], [y], s=node_size * 1.25, + color=[event_color_value], alpha=event_alpha, + edgecolors="white", linewidths=0.5, zorder=7, + ) + if show_event_labels and support: + node = plan.node_of_qubit[support[0]] + x, y = positions[node] + ax.text( + x, + y, + str(event_index), + color=( + event_color_value + if color_by == "gate" + else "#59636e" + ), + fontsize=8, + ha="center", + va="center", + zorder=8, + ) + + if show_site_labels: + for qubit in qubits: + node = plan.node_of_qubit[qubit] + x, y = positions[node] + ax.annotate( + f"q{qubit}", + (x, y), + xytext=(0, 7), + textcoords="offset points", + ha="center", + fontsize=8, + color="#374151", + zorder=9, + ) + if show_node_ids: + for node in plan.nodes(): + x, y = positions[node] + ax.annotate( + f"n{node}", + (x, y), + xytext=(0, -10), + textcoords="offset points", + ha="center", + fontsize=7, + color="#5b6168", + zorder=9, + ) + + colorbar_count = ( + n_events if color_by == "gate" and show_gate_paths else n_scales + ) + if colorbar and colorbar_count: + add_order_colorbar( + fig, + ax, + colormaps, + ScalarMappable, + Normalize, + cmap if color_by == "gate" else scale_cmap, + colorbar_count, + label=( + "gate stream order" + if color_by == "gate" + else "tree scale (leaf = 0)" + ), + ) + title = ( + "Tree layout finder — " + + ("colored gate paths" if color_by == "gate" else "scale-colored tree") + ) + if show_axes: + if show_title: + ax.set_title(title) + ax.set_xlabel("layout x") + ax.set_ylabel("layout y") + ax.margins(0.14) + ax.set_aspect("equal", adjustable="datalim") + else: + finish_schematic_axes( + ax, + title=title if show_title else None, + margins=0.14, + ) + return fig, ax + + def plot_tent( + self, + plan=None, + *, + site_coords=None, + ax=None, + figsize=(8, 7), + cmap="turbo", + edge_cmap="turbo", + node_cmap="YlOrRd", + color_by="scale", + edge_color="#2f80a0", + show_edge_arrows=True, + arrow_size=8.0, + order=True, + lattice=True, + show_gate_connectivity=True, + show_node_ids=False, + show_site_labels=False, + colorbar=False, + show_axes=False, + show_title=False, + node_size=38, + edge_linewidth=1.35, + edge_alpha=1.0, + vertical_spacing=None, + ): + """Plot the hierarchy as a Cotengra-style tent over the raw graph. + + Physical sites and gate connectivity stay in the lower, grey raw + graph. Internal TTN nodes are lifted above the mean position of their + descendant sites, and each parent-child hierarchy edge uses one + uniform solid color by default. Pass ``edge_color=None`` to color + edges by ``edge_cmap`` and ``color_by``. Small arrows at edge + midpoints show the parent-to-child direction by default; disable them + with ``show_edge_arrows=False``. This is deliberately a structural + visualization: + gate-by-gate route overlays are not drawn. Set ``order=True`` to place + hierarchy nodes by a deterministic post-order traversal, matching the + ordering option in Cotengra's tent plots. Use ``color_by="order"`` if + the same traversal should also control the colors. + """ + plt, colormaps, ScalarMappable, Normalize, _FancyArrowPatch = ( + matplotlib_modules() + ) + if plan is None: + plan = self.run() + if not isinstance(plan, TreePlan): + raise TypeError("plan must be a TreePlan returned by run().") + color_by = str(color_by).replace("-", "_").strip().lower() + color_by = {"level": "scale", "size": "scale"}.get( + color_by, color_by + ) + if color_by not in {"scale", "order"}: + raise ValueError("color_by must be 'scale' or 'order'.") + try: + arrow_size = float(arrow_size) + except (TypeError, ValueError) as exc: + raise TypeError( + "arrow_size must be a positive real number." + ) from exc + if not np.isfinite(arrow_size) or arrow_size <= 0.0: + raise ValueError("arrow_size must be a positive real number.") + created_ax = ax is None + if created_ax: + _, ax = plt.subplots(figsize=figsize) + if not show_axes: + ax.figure.subplots_adjust(left=0, right=1, bottom=0, top=1) + fig = ax.figure + + qubits = tuple(range(plan.n)) + coords = resolve_site_coords(qubits, site_coords) + node_scales = _tree_node_scales(plan) + n_scales = max(node_scales.values(), default=0) + 1 + + if lattice: + for left, right in coordinate_lattice_edges(coords): + ax.plot( + (coords[left][0], coords[right][0]), + (coords[left][1], coords[right][1]), + color="#d5d9de", + linewidth=1.0, + alpha=0.78, + zorder=1, + ) + + if show_gate_connectivity: + lattice_pairs = ( + coordinate_lattice_edge_keys(coords) + if lattice + else set() + ) + for support in self.supports: + unique = tuple(dict.fromkeys(support)) + for left, right in zip(unique, unique[1:]): + if frozenset((left, right)) in lattice_pairs: + continue + ax.plot( + (coords[left][0], coords[right][0]), + (coords[left][1], coords[right][1]), + color="#7e8995", + linewidth=0.72, + linestyle="-", + alpha=0.62, + zorder=1, + ) + + subtree_qubits = {} + + def gather_qubits(node): + if node in subtree_qubits: + return subtree_qubits[node] + result = [] + if node in plan.qubit_of_leaf: + result.append(plan.qubit_of_leaf[node]) + if node == plan.root and plan.root_qubit is not None: + result.append(plan.root_qubit) + for child in plan.children.get(node, ()): + result.extend(gather_qubits(child)) + subtree_qubits[node] = tuple(result) + return subtree_qubits[node] + + for node in plan.nodes(): + gather_qubits(node) + + x_span = max( + max(point[0] for point in coords.values()) + - min(point[0] for point in coords.values()), + 1.0, + ) + y_max = max(point[1] for point in coords.values()) + if vertical_spacing is None: + vertical_spacing = max(0.7, 0.28 * x_span) + vertical_spacing = float(vertical_spacing) + if vertical_spacing <= 0.0: + raise ValueError("vertical_spacing must be positive.") + + positions = { + node: coords[qubit] + for node, qubit in plan.qubit_of_leaf.items() + } + if plan.root_qubit is not None: + positions[plan.root] = coords[plan.root_qubit] + + internal = [node for node in plan.nodes() if not plan.is_leaf(node)] + if order or color_by == "order": + postorder = [] + + def visit(node): + for child in plan.children.get(node, ()): + visit(child) + postorder.append(node) + + visit(plan.root) + order_values = {node: i for i, node in enumerate(postorder)} + order_count = max(1, len(postorder)) + internal_order = { + node: index + for index, node in enumerate( + node for node in postorder if node in internal + ) + } + order_span = max(1, n_scales - 1) + order_denominator = max(1, len(internal) - 1) + for node in internal: + sites = gather_qubits(node) + x = sum(coords[qubit][0] for qubit in sites) / len(sites) + if order: + # Preserve post-order relationships without giving every + # internal node a separate vertical layer. A separate + # layer for all nodes makes larger 2D circuits needlessly + # tall and narrow without conveying extra geometry. + height = 1.0 + order_span * ( + internal_order[node] / order_denominator + ) + else: + height = 1.0 + order_values[node] / order_count + y = y_max + vertical_spacing * height + positions[node] = (x, y) + n_colors = order_count if color_by == "order" else n_scales + else: + order_values = None + for node in internal: + sites = gather_qubits(node) + x = sum(coords[qubit][0] for qubit in sites) / len(sites) + y = y_max + vertical_spacing * (node_scales[node] + 1.0) + positions[node] = (x, y) + n_colors = n_scales + + def node_color(node): + if color_by == "order": + return event_color( + colormaps, cmap, order_values[node], n_colors + ) + return scale_color( + colormaps, node_cmap, node_scales[node], n_colors + ) + + def hierarchy_edge_color(parent): + if edge_color is not None: + return edge_color + if color_by == "order": + return event_color( + colormaps, cmap, order_values[parent], n_colors + ) + return scale_color( + colormaps, edge_cmap, node_scales[parent], n_colors + ) + + for parent, children in plan.children.items(): + for child in children: + x0, y0 = positions[parent] + x1, y1 = positions[child] + edge_color_value = hierarchy_edge_color(parent) + ax.plot( + (x0, x1), + (y0, y1), + color=edge_color_value, + linewidth=edge_linewidth, + alpha=edge_alpha, + zorder=2, + ) + if show_edge_arrows: + dx = x1 - x0 + dy = y1 - y0 + ax.add_patch( + _FancyArrowPatch( + (x0 + 0.42 * dx, y0 + 0.42 * dy), + (x0 + 0.62 * dx, y0 + 0.62 * dy), + arrowstyle="-|>", + mutation_scale=arrow_size, + linewidth=max(0.6, 0.75 * edge_linewidth), + color=edge_color_value, + shrinkA=0.0, + shrinkB=0.0, + zorder=3, + ) + ) + + for node in plan.nodes(): + x, y = positions[node] + ax.scatter( + [x], + [y], + s=node_size, + marker="o", + color=[node_color(node)], + edgecolors="#41464c", + linewidths=0.65, + zorder=4, + ) + + if show_site_labels: + for qubit in qubits: + node = plan.node_of_qubit[qubit] + x, y = positions[node] + ax.annotate( + f"q{qubit}", + (x, y), + xytext=(0, 7), + textcoords="offset points", + ha="center", + fontsize=8, + color="#374151", + zorder=5, + ) + if show_node_ids: + for node in plan.nodes(): + x, y = positions[node] + ax.annotate( + f"n{node}", + (x, y), + xytext=(0, -10), + textcoords="offset points", + ha="center", + fontsize=7, + color="#5b6168", + zorder=5, + ) + + if colorbar and n_colors: + add_order_colorbar( + fig, + ax, + colormaps, + ScalarMappable, + Normalize, + cmap if color_by == "order" else edge_cmap, + n_colors, + label=( + "tree order" if color_by == "order" else "tree scale" + ), + ) + + title = "Tree tent" + if show_axes: + if show_title: + ax.set_title(title) + ax.set_xlabel("layout x") + ax.set_ylabel("hierarchy height") + ax.set_aspect("equal", adjustable="datalim") + ax.margins(0.14) + else: + finish_schematic_axes( + ax, + title=title if show_title else None, + margins=0.14, + ) + return fig, ax + + # The public default is the structural tent view. Keep the older direct + # route renderer private so the hierarchy cannot be mistaken for a set of + # gate-stream legs. + plot = plot_tent + + def plot_rubberband( + self, + plan=None, + *, + site_coords=None, + ax=None, + figsize=(10, 8), + cmap="turbo", + color_by="scale", + scale_cmap="viridis", + lattice=True, + show_gate_connectivity=True, + show_site_nodes=True, + colorbar=False, + show_axes=False, + show_title=False, + band_alpha=0.68, + band_linewidth=1.35, + band_padding=0.12, + node_size=58, + ): + """Plot hierarchical tree clusters as smooth rubberband regions. + + This is the physical-lattice counterpart to Quimb's contraction-tree + ``plot_rubberband`` view: the lattice and gate connectivity remain + grey, while each non-root tree cluster is wrapped by a rounded, + translucent colored band. ``color_by="gate"`` colors the bands by a + deterministic post-order through the tree; ``color_by="scale"`` uses + one stable color for each tree scale measured from the leaves. + + The default presentation has no axes, site labels, or title. It + returns a normal Matplotlib ``(fig, ax)`` pair for further styling. + """ + plt, colormaps, ScalarMappable, Normalize, _FancyArrowPatch = ( + matplotlib_modules() + ) + from matplotlib.patches import FancyBboxPatch # noqa: PLC0415 + + if plan is None: + plan = self.run() + if not isinstance(plan, TreePlan): + raise TypeError("plan must be a TreePlan returned by run().") + color_by = str(color_by).replace("-", "_").strip().lower() + color_by = {"stream": "gate", "event": "gate", "level": "scale"}.get( + color_by, color_by + ) + if color_by not in {"gate", "scale"}: + raise ValueError("color_by must be 'gate' or 'scale'.") + created_ax = ax is None + if created_ax: + _, ax = plt.subplots(figsize=figsize) + if not show_axes: + ax.figure.subplots_adjust(left=0, right=1, bottom=0, top=1) + fig = ax.figure + + qubits = tuple(range(plan.n)) + coords = resolve_site_coords(qubits, site_coords) + node_scales = _tree_node_scales(plan) + n_scales = max(node_scales.values(), default=0) + 1 + + if lattice: + for left, right in coordinate_lattice_edges(coords): + ax.plot( + (coords[left][0], coords[right][0]), + (coords[left][1], coords[right][1]), + color="#c7cdd3", + linewidth=1.0, + alpha=0.62, + zorder=1, + ) + + if show_gate_connectivity: + lattice_pairs = ( + coordinate_lattice_edge_keys(coords) + if lattice + else set() + ) + for support in self.supports: + unique = tuple(dict.fromkeys(support)) + for left, right in zip(unique, unique[1:]): + if frozenset((left, right)) in lattice_pairs: + continue + ax.plot( + (coords[left][0], coords[right][0]), + (coords[left][1], coords[right][1]), + color="#87919b", + linewidth=0.7, + linestyle="-", + alpha=0.42, + zorder=1, + ) + + subtree_qubits = {} + + def gather_qubits(node): + if node in subtree_qubits: + return subtree_qubits[node] + children = tuple(plan.children.get(node, ())) + if not children: + result = (plan.qubit_of_leaf[node],) + else: + result = tuple( + qubit + for child in children + for qubit in gather_qubits(child) + ) + if node == plan.root and plan.root_qubit is not None: + result += (plan.root_qubit,) + subtree_qubits[node] = result + return result + + band_nodes = [] + + def visit(node): + for child in plan.children.get(node, ()): + visit(child) + if plan.children.get(node): + band_nodes.append(node) + + visit(plan.root) + n_bands = max(1, len(band_nodes)) + for band_index, node in enumerate(band_nodes): + sites = tuple(dict.fromkeys(gather_qubits(node))) + if len(sites) < 2: + continue + points = [coords[qubit] for qubit in sites] + xmin = min(point[0] for point in points) + xmax = max(point[0] for point in points) + ymin = min(point[1] for point in points) + ymax = max(point[1] for point in points) + padding = band_padding + 0.012 * band_index + width = max(xmax - xmin, 0.16) + 2.0 * padding + height = max(ymax - ymin, 0.16) + 2.0 * padding + rounding = min(0.28, 0.45 * min(width, height)) + if color_by == "scale": + color = scale_color( + colormaps, + scale_cmap, + node_scales[node], + n_scales, + ) + else: + color = event_color(colormaps, cmap, band_index, n_bands) + ax.add_patch( + FancyBboxPatch( + (xmin - padding, ymin - padding), + width, + height, + boxstyle=f"round,pad=0,rounding_size={rounding}", + fill=False, + edgecolor=color, + linewidth=band_linewidth, + alpha=band_alpha, + zorder=3, + ) + ) + + if show_site_nodes: + ax.scatter( + [coords[qubit][0] for qubit in qubits], + [coords[qubit][1] for qubit in qubits], + s=node_size, + marker="o", + color="#858b91", + edgecolors="#3f454b", + linewidths=0.75, + zorder=5, + ) + + if colorbar and (n_bands if color_by == "gate" else n_scales): + add_order_colorbar( + fig, + ax, + colormaps, + ScalarMappable, + Normalize, + cmap if color_by == "gate" else scale_cmap, + n_bands if color_by == "gate" else n_scales, + label=( + "rubberband order" + if color_by == "gate" + else "tree scale (leaf = 0)" + ), + ) + + title = "Tree rubberband" + if show_axes: + if show_title: + ax.set_title(title) + ax.set_xlabel("logical site x") + ax.set_ylabel("logical site y") + ax.set_aspect("equal", adjustable="datalim") + ax.margins(0.14) + else: + finish_schematic_axes( + ax, + title=title if show_title else None, + margins=0.14, + ) + return fig, ax + + plot_layout = plot diff --git a/src/pepsy/optimizers/tree/optimizer.py b/src/pepsy/optimizers/tree/optimizer.py index a655ac8..0fb352a 100644 --- a/src/pepsy/optimizers/tree/optimizer.py +++ b/src/pepsy/optimizers/tree/optimizer.py @@ -1423,6 +1423,84 @@ def layout_report(self): } return self.layout_finder.report(self.plan) + def plot_layout(self, plan=None, *, layout_kwargs=None, **plot_kwargs): + """Plot the tree layout as a Cotengra-style tent. + + The method returns ``(fig, ax)`` and does not alter the live TTN. For + an optimizer constructed with an explicit :class:`TreePlan`, the + temporary finder needed for gate diagnostics is built from the queued + stream without changing that plan. + """ + finder = self.layout_finder + if finder is None: + finder = TreeLayoutFinder( + gates=self._layout_gate_stream(), + n=self.n, + structure=self.structure, + max_arity=self.max_arity, + community_frac=self.community_frac, + star_frac=self.star_frac, + objective=self.layout_objective, + weight_mode=self.layout_weight_mode, + chi=self.chi, + max_operator_qubits=self.max_operator_qubits, + root_qubit=self.plan.root_qubit, + ) + if plan is None: + if layout_kwargs: + plan = finder.run(**dict(layout_kwargs)) + else: + plan = self.plan + return finder.plot(plan, **plot_kwargs) + + def plot_rubberband(self, plan=None, *, layout_kwargs=None, **plot_kwargs): + """Plot the tree's physical clusters as translucent rubberbands.""" + finder = self.layout_finder + if finder is None: + finder = TreeLayoutFinder( + gates=self._layout_gate_stream(), + n=self.n, + structure=self.structure, + max_arity=self.max_arity, + community_frac=self.community_frac, + star_frac=self.star_frac, + objective=self.layout_objective, + weight_mode=self.layout_weight_mode, + chi=self.chi, + max_operator_qubits=self.max_operator_qubits, + root_qubit=self.plan.root_qubit, + ) + if plan is None: + if layout_kwargs: + plan = finder.run(**dict(layout_kwargs)) + else: + plan = self.plan + return finder.plot_rubberband(plan, **plot_kwargs) + + def plot_tent(self, plan=None, *, layout_kwargs=None, **plot_kwargs): + """Plot the selected tree as a Cotengra-style tent over the lattice.""" + finder = self.layout_finder + if finder is None: + finder = TreeLayoutFinder( + gates=self._layout_gate_stream(), + n=self.n, + structure=self.structure, + max_arity=self.max_arity, + community_frac=self.community_frac, + star_frac=self.star_frac, + objective=self.layout_objective, + weight_mode=self.layout_weight_mode, + chi=self.chi, + max_operator_qubits=self.max_operator_qubits, + root_qubit=self.plan.root_qubit, + ) + if plan is None: + if layout_kwargs: + plan = finder.run(**dict(layout_kwargs)) + else: + plan = self.plan + return finder.plot_tent(plan, **plot_kwargs) + def canonize_subtree(self, nodes, *, span=False): """Canonicalise the state around the connected subtree ``nodes``. diff --git a/tests/test_optimize_mps.py b/tests/test_optimize_mps.py index 77de9e9..d08d551 100644 --- a/tests/test_optimize_mps.py +++ b/tests/test_optimize_mps.py @@ -682,6 +682,59 @@ def test_mps_optimizer_gate_stream_layout_remaps_long_range_path(): ) +def test_mps_layout_finder_plot_draws_lattice_and_gate_order(): + """The MPS plot exposes the lattice, gate graph, and colored chain.""" + matplotlib = pytest.importorskip("matplotlib") + matplotlib.use("Agg", force=True) + import matplotlib.pyplot as plt + + gates = [ + (qu.CNOT(), (0, 3)), + (qu.CNOT(), (3, 1)), + (qu.CNOT(), (1, 2)), + ] + finder = py.MpsOptimizer.LayoutFinder(gates, L=4) + plan = finder.run(order="input") + fig, ax = finder.plot( + plan, + site_coords={0: (0, 0), 1: (1, 0), 2: (0, 1), 3: (1, 1)}, + ) + + assert fig is ax.figure + assert ax.get_title() == "" + assert len(ax.patches) == len(plan["site_order"]) - 1 + assert len(fig.axes) == 1 # no stream-order colorbar by default + assert not ax.axison # schematic-style presentation by default + assert any(text.get_text() == "0" for text in ax.texts) + assert any(text.get_text() == "3" for text in ax.texts) + assert any(collection.get_offsets().shape[0] for collection in ax.collections) + plt.close(fig) + + +def test_mps_optimizer_plot_layout_is_non_mutating(): + """The optimizer plotting wrapper does not install or alter a layout.""" + matplotlib = pytest.importorskip("matplotlib") + matplotlib.use("Agg", force=True) + import matplotlib.pyplot as plt + + p0 = qtn.MPS_computational_state("0000", dtype="complex128") + opt = py.MpsOptimizer( + p0, + gates=[(qu.CNOT(), (0, 3))], + chi=8, + mode="svd", + ) + before = tuple(opt.logical_order) + fig, _ = opt.plot_layout( + layout_kwargs={"order": "input"}, + site_coords={q: (q, 0) for q in range(4)}, + ) + + assert tuple(opt.logical_order) == before + assert opt._persistent_layout_plan is None + plt.close(fig) + + def test_mps_optimizer_gate_stream_layout_accepts_weight_fn(): """User event weights should feed the weighted graph and report.""" gates = [ diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index ad4fd56..8122d75 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -1484,6 +1484,225 @@ def test_layout_report_summarizes_quality(): assert rep["score"] <= rep["balanced_score"] + 1e-9 +def test_tree_layout_finder_plot_defaults_to_tent(): + """The public plot shows the structural tent, not gate-route overlays.""" + matplotlib = pytest.importorskip("matplotlib") + matplotlib.use("Agg", force=True) + import matplotlib.pyplot as plt + + gates = [ + (pepsy.cnot(), (0, 3)), + (pepsy.cnot(), (3, 1)), + ] + finder = TreeLayoutFinder(gates, n=4, max_arity=2) + plan = finder.run() + assert plan.is_binary() + assert len(plan.children[plan.root]) == 2 + fig, ax = finder.plot( + plan, + site_coords={0: (0, 0), 1: (1, 0), 2: (0, 1), 3: (1, 1)}, + ) + + assert fig is ax.figure + assert ax.get_title() == "" + assert len(ax.patches) == len(plan.nodes()) - 1 + assert len(fig.axes) == 1 + assert not ax.axison # schematic-style presentation by default + assert not ax.texts + assert len(ax.collections) == len(plan.nodes()) + plt.close(fig) + + +def test_tree_layout_finder_can_hide_gate_paths_for_structural_view(): + """The structural view makes the binary TTN edges unambiguous.""" + matplotlib = pytest.importorskip("matplotlib") + matplotlib.use("Agg", force=True) + import matplotlib.pyplot as plt + + gates = [ + (pepsy.cnot(), (0, 3)), + (pepsy.cnot(), (3, 1)), + ] + finder = TreeLayoutFinder(gates, n=4, max_arity=2) + plan = finder.run() + fig, ax = finder.plot( + plan, + lattice=False, + show_gate_connectivity=False, + show_edge_arrows=False, + ) + + assert len(ax.lines) == len(plan.nodes()) - 1 + assert not ax.patches + assert not ax.texts + assert not ax.axison + plt.close(fig) + + +def test_tree_layout_finder_plot_tent_draws_hierarchy_over_raw_graph(): + """Tent plotting separates the binary hierarchy from raw connectivity.""" + matplotlib = pytest.importorskip("matplotlib") + matplotlib.use("Agg", force=True) + import matplotlib.pyplot as plt + + gates = [ + (pepsy.cnot(), (0, 3)), + (pepsy.cnot(), (3, 1)), + ] + finder = TreeLayoutFinder(gates, n=4, max_arity=2) + plan = finder.run() + fig, ax = finder.plot_tent( + plan, + site_coords={0: (0, 0), 1: (1, 0), 2: (0, 1), 3: (1, 1)}, + ) + + assert plan.is_binary() + assert fig is ax.figure + assert len(ax.patches) == len(plan.nodes()) - 1 + assert not ax.texts + assert not ax.axison + assert len(ax.lines) >= len(plan.nodes()) - 1 + assert len(ax.collections) == len(plan.nodes()) + plt.close(fig) + + +def test_tree_layout_scale_colors_do_not_depend_on_gate_stream_length(): + """Scale coloring remains fixed when the gate stream changes.""" + matplotlib = pytest.importorskip("matplotlib") + matplotlib.use("Agg", force=True) + import matplotlib.pyplot as plt + + plan = TreePlan.from_order(range(8), structure="balanced", max_arity=2) + streams = [ + [(pepsy.cnot(), (0, 7))], + [ + (pepsy.cnot(), (0, 7)), + (pepsy.cnot(), (1, 6)), + (pepsy.cnot(), (2, 5)), + (pepsy.cnot(), (3, 4)), + ], + ] + structural_colors = [] + scale_node_colors = [] + for gates in streams: + finder = TreeLayoutFinder(gates, n=8, max_arity=2) + fig, ax = finder.plot_tent( + plan, + color_by="scale", + edge_color=None, + show_edge_arrows=False, + ) + # With the default one-dimensional coordinates, the first lines are + # the lattice and only the non-lattice gate-connectivity background; + # nearest-neighbor gates are already represented by the lattice. + lattice_pairs = { + frozenset((site, site + 1)) + for site in range(plan.n - 1) + } + nonlattice_gates = sum( + frozenset(where) not in lattice_pairs for _, where in gates + ) + background_lines = len(plan.leaves()) - 1 + nonlattice_gates + structural_colors.append( + tuple( + tuple(line.get_color()) + for line in ax.lines[background_lines:] + ) + ) + scale_node_colors.append( + tuple( + tuple(collection.get_facecolors()[0]) + for collection in ax.collections + ) + ) + assert len(fig.axes) == 1 + assert len(ax.collections) == len(plan.nodes()) + plt.close(fig) + + assert structural_colors[0] == structural_colors[1] + assert scale_node_colors[0] == scale_node_colors[1] + assert len(set(structural_colors[0])) > 1 + assert len(set(scale_node_colors[0])) > 1 + + +def test_tree_layout_tent_edges_are_uniform_by_default(): + """Tent hierarchy edges use one solid color unless explicitly varied.""" + matplotlib = pytest.importorskip("matplotlib") + matplotlib.use("Agg", force=True) + import matplotlib.pyplot as plt + + gates = [(pepsy.cnot(), (0, 3)), (pepsy.cnot(), (1, 2))] + finder = TreeLayoutFinder(gates, n=4, max_arity=2) + plan = finder.run() + fig, ax = finder.plot_tent(plan, color_by="scale") + + lattice_pairs = { + frozenset((site, site + 1)) for site in range(plan.n - 1) + } + background_lines = len(plan.leaves()) - 1 + sum( + frozenset(where) not in lattice_pairs for _, where in gates + ) + hierarchy_colors = { + line.get_color() for line in ax.lines[background_lines:] + } + assert hierarchy_colors == {"#2f80a0"} + assert len(ax.patches) == len(plan.nodes()) - 1 + plt.close(fig) + + +def test_tree_layout_tent_validates_arrow_size(): + """Arrow marker sizing rejects values Matplotlib cannot render usefully.""" + finder = TreeLayoutFinder( + [(pepsy.cnot(), (0, 1))], n=2, max_arity=2 + ) + plan = finder.run() + + with pytest.raises(ValueError, match="arrow_size"): + finder.plot_tent(plan, arrow_size=0.0) + + +def test_tree_layout_finder_plot_rubberband_is_axis_free_and_unlabeled(): + """Rubberband plots show clusters without plot text or axes.""" + matplotlib = pytest.importorskip("matplotlib") + matplotlib.use("Agg", force=True) + import matplotlib.pyplot as plt + + gates = [ + (pepsy.cnot(), (0, 3)), + (pepsy.cnot(), (3, 1)), + (pepsy.cnot(), (1, 2)), + ] + finder = TreeLayoutFinder(gates, n=4, max_arity=2) + plan = finder.run() + fig, ax = finder.plot_rubberband( + plan, + site_coords={0: (0, 0), 1: (1, 0), 2: (0, 1), 3: (1, 1)}, + ) + + assert fig is ax.figure + assert ax.get_title() == "" + assert not ax.axison + assert not ax.texts + assert len(ax.patches) >= 1 + plt.close(fig) + + +def test_tree_optimizer_plot_layout_with_explicit_plan_is_non_mutating(): + """The tree optimizer wrapper plots an explicit plan without replay.""" + matplotlib = pytest.importorskip("matplotlib") + matplotlib.use("Agg", force=True) + import matplotlib.pyplot as plt + + gates = [(pepsy.cnot(), (0, 3))] + plan = TreeLayoutFinder(gates, n=4, max_arity=2).run() + opt = TreeOptimizer(gates, tree=plan, run=False) + before = opt.to_dense().copy() + fig, _ = opt.plot_layout(site_coords={q: (q, 0) for q in range(4)}) + + assert np.allclose(opt.to_dense(), before) + plt.close(fig) + + def test_bond_report_reflects_chi(): """bond_report caps at chi and counts the tree tensors.""" rng = np.random.default_rng(12) From f9ea61db787ecd1dd24cf156898482ba611bfc43 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Thu, 30 Jul 2026 16:12:44 -0700 Subject: [PATCH 33/70] improve MPS and tree layout diagnostics --- docs/api/optimizers/mps.md | 27 ++ docs/api/optimizers/tree.md | 69 ++++- src/pepsy/optimizers/mps/layout.py | 406 ++++++++++++++++++++++--- src/pepsy/optimizers/mps/optimizer.py | 152 ++++++++- src/pepsy/optimizers/tree/layout.py | 260 +++++++++++++--- src/pepsy/optimizers/tree/optimizer.py | 148 ++++++++- tests/test_optimize_mps.py | 45 +++ tests/test_optimize_tree.py | 113 ++++++- 8 files changed, 1122 insertions(+), 98 deletions(-) diff --git a/docs/api/optimizers/mps.md b/docs/api/optimizers/mps.md index 2f47f95..df15e63 100644 --- a/docs/api/optimizers/mps.md +++ b/docs/api/optimizers/mps.md @@ -152,6 +152,14 @@ default to `weight_mode="auto"`: angle metadata when present, otherwise a cheap operator-Schmidt proxy for small dense two-site gates, falling back to count weights. Pass `weight_fn(payload, support, event_type)` for explicit weights. +For compression-oriented selection, pass `objective="compression"`. This +uses operator-Schmidt load over every MPS cut crossed by each support, with +support span retained as a replay-cost tie-breaker. Exact small dense ranks +are used when available; opaque, native, and wide operators use a conservative +operator-space rank bound and are marked in `rank_bound_reasons` rather than +silently being treated as rank two. The default `objective="locality"` keeps +the faster span/congestion heuristic for backwards compatibility. + The layout score depends on gate supports and optional gate/event weights, not on the initial MPS tensor values. The plan does not rewrite the gate stream. To use a layout during replay, call `opt.run(use_layout_finder=True)` or pass a @@ -160,6 +168,25 @@ temporarily permutes the working MPS and restores the returned MPS to the original site order. Layout-aware replay prints a concise report by default; pass `layout_report=False` to silence it. +When the current state matters, use the explicit pilot selector: + +```python +plan = opt.select_layout_for_compression( + pilot_candidates=4, + pilot_steps=64, +) +opt.apply_layout(plan, layout_report=False) +``` + +The selector replays the best static candidates on independent copies using +the real MPS mode, `chi`, cutoff, backend, and dtype. It enables the +infidelity trace for the pilot and chooses by measured compression +infidelity, final bond dimension, and elapsed time, and returns +per-candidate records under `plan["pilot"]`. The original state, queue, and +layout are unchanged. Perform this before installing a persistent layout; +reordering an already-entangled MPS remains explicitly guarded because the +reorder itself can be lossy or expensive. + The layout can be inspected graphically without changing the optimizer. The finder returns a Matplotlib `(fig, ax)` pair. The original lattice and gate connectivity remain a light grey background, while the colored arrow chain diff --git a/docs/api/optimizers/tree.md b/docs/api/optimizers/tree.md index 9e20133..0f6b656 100644 --- a/docs/api/optimizers/tree.md +++ b/docs/api/optimizers/tree.md @@ -397,12 +397,25 @@ it combines normalized path score, maximum edge load, and total edge load with `weight_mode` / `layout_weight_mode` option accepts `count`, `auto`, `angle`, or `operator_schmidt` for interaction-graph weighting. +For compression-first selection, use `objective="compression"` (or +`layout_objective="compression"`). It prioritizes peak and total predicted +operator-Schmidt load, then penalizes the estimated local tensor size at the +configured `chi`, and only then uses path length. This differs from the +default fast `"path"` objective: path optimizes routing locality, while +compression also accounts for bond pressure and wider-node cost. + `weight_mode="operator_schmidt"` is a cheap **two-qubit entangling-strength proxy** used to form the spectral qubit order; it is not itself the exact operator-Schmidt rank. Use `objective="congestion"` when selecting a tree: its edge-load calculation uses the actual rank across each candidate tree cut (or an MPO bond bound), which is the quantity that predicts TTN bond growth. +Rank diagnostics are explicit. Small dense qubit operators use exact +operator-Schmidt ranks; opaque native arrays, MPO bond products, and supports +larger than `max_operator_qubits` use conservative operator-space bounds and +are counted in `rank_bounded_events` with a reason in `rank_bound_reasons`. +They are not silently assigned rank two. + The structure is **not restricted to binary trees**. Internal nodes may have any arity, controlled by two knobs on `TreeLayoutFinder` / `TreePlan.from_order` / `TreeOptimizer`: @@ -567,12 +580,31 @@ chi=chi)` forwards its own `chi` into the finder it builds -- so the everyday exact at `chi`. A bare finder with no `chi` searches `chi`-blind. Set `max_operator_qubits` to bound dense rank diagnostics and operator allocation; wider native MPO events can still replay without dense -materialization. `TreeLayoutFinder(..., max_operator_qubits=...)` uses a conservative rank -proxy above that width. `report(plan, include_edge_loads=False)` skips the -event-by-edge congestion calculation for path-only diagnostics; when loads are +materialization. `TreeLayoutFinder(..., max_operator_qubits=...)` uses a +conservative rank bound above that width and reports the bounded events. The +public `score` remains the path score for compatibility; inspect +`objective_key`, `max_edge_load`, `peak_bond_growth`, and the tensor-cost +fields for compression decisions. `report(plan, include_edge_loads=False)` +skips the event-by-edge calculation for path-only diagnostics; when loads are included, `peak_bond_growth_log2` remains finite even when the human-readable `peak_bond_growth` would overflow floating point. +For a state-aware choice between static candidates, call: + +```python +choice = opt.select_layout_for_compression( + pilot_candidates=4, + pilot_steps=64, +) +opt = py.TreeOptimizer(gate_stream, tree=choice["plan"], chi=chi) +``` + +The pilot replays candidates on independent copies with the real tree update +kernels and returns measured infidelity, final bond, truncation count, and +runtime under `choice["pilot"]`. The original optimizer is unchanged unless +`install=True` is passed. Installation is restricted to product initial states; +an entangled TTN cannot generally be relaid out exactly. + Both helpers are also available from the package-level API: ```python @@ -660,9 +692,11 @@ The same diagnostics are available as a Cotengra-style tent plot. `plot_tent(plan)`): it keeps the raw graph at the bottom and lifts the selected hierarchy above its descendant sites: the raw lattice and gate connectivity are gray, while circular nodes use stable scale colors. Hierarchy -edges use one uniform solid color by default; pass `edge_color=None` to restore -scale/order-colored edges. Midpoint arrows are enabled by default and indicate -the parent-to-child direction; pass `show_edge_arrows=False` to hide them. +edges use one uniform solid color by default. Pass `edge_color=None` to make +each incoming edge exactly match the node it terminates at; `node_cmap` then +controls both. Arrows are disabled by default, matching Cotengra's +structural tent view; pass `show_edge_arrows=True` only when parent-to-child +direction is needed. Nearest-neighbor gate edges are not duplicated over the lattice. Supplying `site_coords={qubit: (x, y)}` places the physical sites on an existing lattice. It returns `(fig, ax)` and does not mutate the plan or live TTN: @@ -675,10 +709,10 @@ fig, ax = finder.plot( site_coords=logical_lattice_coords, color_by="scale", edge_color=None, - edge_cmap="turbo", + edge_cmap="GnBu", node_cmap="YlOrRd", - order=False, - show_edge_arrows=True, + order=True, + show_edge_arrows=False, ) # For a live optimizer, the same plot is available without changing its state. @@ -721,15 +755,18 @@ fig, ax = finder.plot( site_coords=logical_lattice_coords, color_by="scale", edge_color=None, - edge_cmap="turbo", + edge_cmap="GnBu", node_cmap="YlOrRd", ) ``` Here leaves are scale zero and nodes use stable colors for their hierarchical -scale. To also color hierarchy edges by scale, pass `edge_color=None`. -Midpoint arrows show the direction from each parent to its children. The -mapping is independent of the number or order of gates; +scale. With `edge_color=None`, each incoming hierarchy edge uses the exact +same node-palette color as the node it terminates at, so scale layers remain +easy to follow. Set a literal `edge_color` when a uniform structural edge +color is preferred. The scale colorbar, when enabled, follows `node_cmap`. +Midpoint arrows can show the direction from each parent to its children when +`show_edge_arrows=True`. The mapping is independent of the number or order of gates; `colorbar=True` then labels tree scale rather than gate-stream order. The plot has no title by default; pass `show_title=True` if a title is wanted. @@ -747,8 +784,10 @@ fig, ax = opt.plot_rubberband(site_coords=logical_lattice_coords) ``` This keeps the lattice sites and gate connectivity grey and wraps each -non-root tree cluster in a rounded, translucent band. Use `color_by="scale"` -for one stable band color per tree scale. +non-root tree cluster in a rounded, translucent band. The default is a +Cotengra-style `Spectral` post-order progression, giving each nested band a +distinct color. Use `color_by="scale"` for one stable band color per tree +scale. - `TreeLayoutFinder.report(plan=None)` summarises the physical-node geodesic lengths over the interaction graph (`score`, `max_path`, `mean_path`, diff --git a/src/pepsy/optimizers/mps/layout.py b/src/pepsy/optimizers/mps/layout.py index 60d25c9..be1e788 100644 --- a/src/pepsy/optimizers/mps/layout.py +++ b/src/pepsy/optimizers/mps/layout.py @@ -281,6 +281,138 @@ def _operator_schmidt_weight(payload, support, *, schmidt_max_dim=4): return float(powers[1:].sum() / total) +def _operator_schmidt_rank_bound(support, left_support, local_dims=None): + """Return the maximum operator-Schmidt rank for a support cut. + + For a product of local operator spaces the rank is bounded by the smaller + operator-space dimension on either side. The default qubit dimensions + keep this useful even when a payload is opaque, too wide to inspect, or a + native symmetric array cannot be lowered to dense NumPy data. + """ + support = tuple(support) + left = set(left_support) + if not left or left == set(support): + return 1 + if local_dims is None: + local_dims = (2,) * len(support) + local_dims = tuple(int(dim) for dim in local_dims) + if len(local_dims) != len(support) or any(dim < 1 for dim in local_dims): + local_dims = (2,) * len(support) + left_dim = 1 + right_dim = 1 + for i, site in enumerate(support): + if site in left: + left_dim *= local_dims[i] ** 2 + else: + right_dim *= local_dims[i] ** 2 + return max(1, min(left_dim, right_dim)) + + +def _operator_schmidt_rank_info( + payload, + support, + left_support, + *, + max_operator_qubits=None, +): + """Return an exact rank or an honest conservative rank bound. + + The layout finder must remain usable with native/symmetric payloads, but + silently treating an unknown operator as rank two is unsafe: a wide + operator can have a much larger operator-Schmidt rank. This helper keeps + the numeric ``rank`` field for scoring and records whether it was exact. + """ + support = tuple(support) + left_support = tuple(left_support) + default_bound = _operator_schmidt_rank_bound(support, left_support) + if not left_support or set(left_support) == set(support): + return {"rank": 1, "exact": True, "reason": "trivial_cut"} + if ( + max_operator_qubits is not None + and len(support) > int(max_operator_qubits) + ): + return { + "rank": default_bound, + "exact": False, + "reason": "max_operator_qubits", + } + + raw = getattr(payload, "data", payload) + try: + array = np.asarray(raw) + except Exception: + return {"rank": default_bound, "exact": False, "reason": "opaque"} + if array.size == 0 or not np.issubdtype(array.dtype, np.number): + return {"rank": default_bound, "exact": False, "reason": "opaque"} + + local_dims = None + if array.ndim == 2 and array.shape[0] == array.shape[1]: + dimension = int(array.shape[0]) + local_dim = int(round(dimension ** (1.0 / len(support)))) + if local_dim ** len(support) == dimension: + local_dims = (local_dim,) * len(support) + elif array.ndim == 2 * len(support): + output_dims = tuple(int(dim) for dim in array.shape[:len(support)]) + input_dims = tuple(int(dim) for dim in array.shape[len(support):]) + if output_dims == input_dims: + local_dims = output_dims + + if local_dims is None: + return {"rank": default_bound, "exact": False, "reason": "shape"} + + positions = {site: pos for pos, site in enumerate(support)} + try: + if array.ndim == 2: + array = array.reshape(local_dims + local_dims) + left_positions = [positions[site] for site in left_support] + right_positions = [ + positions[site] for site in support if site not in set(left_support) + ] + axes = ( + left_positions + + [len(support) + pos for pos in left_positions] + + right_positions + + [len(support) + pos for pos in right_positions] + ) + left_dim = 1 + right_dim = 1 + for pos in left_positions: + left_dim *= local_dims[pos] ** 2 + for pos in right_positions: + right_dim *= local_dims[pos] ** 2 + matrix = array.transpose(axes).reshape(left_dim, right_dim) + rank = max(1, int(np.linalg.matrix_rank(matrix))) + except (IndexError, TypeError, ValueError, np.linalg.LinAlgError): + return { + "rank": _operator_schmidt_rank_bound( + support, left_support, local_dims + ), + "exact": False, + "reason": "decomposition", + } + return {"rank": rank, "exact": True, "reason": "dense_svd"} + + +def _gate_stream_layout_objective(objective): + """Normalize MPS layout objective names.""" + name = str(objective).replace("-", "_").strip().lower() + aliases = { + "path": "locality", + "span": "locality", + "routing": "locality", + "compress": "compression", + "bond": "compression", + "bond_load": "compression", + } + name = aliases.get(name, name) + if name not in {"locality", "compression"}: + raise ValueError( + f"Unknown MPS layout objective {objective!r}. Expected " + "'locality' or 'compression'." + ) + return name + + def _normalize_weight_mode(weight_mode): """Normalize user-facing gate-stream weight mode names.""" name = str(weight_mode).replace("-", "_").strip().lower() @@ -370,6 +502,48 @@ def _gate_stream_event_weights( return tuple(weights) +def _gate_stream_event_rank_weights( + payloads, + supports, + event_types, + *, + max_operator_qubits=8, +): + """Return log-rank weights for compression-oriented layout search.""" + weights = [] + exact = [] + reasons = [] + for payload, support, event_type in zip(payloads, supports, event_types): + normalized_type = str(event_type).lower() + support = _unique_ordered(support) + if len(support) < 2 or normalized_type in { + "measure", "reset", "measure_reset", "cap" + }: + weights.append(0.0) + exact.append(True) + reasons.append("non_entangling_event") + continue + if payload is None: + info = { + "rank": _operator_schmidt_rank_bound( + support, support[:1] + ), + "exact": False, + "reason": "missing_payload", + } + else: + info = _operator_schmidt_rank_info( + payload, + support, + support[:1], + max_operator_qubits=max_operator_qubits, + ) + weights.append(float(np.log2(max(1, info["rank"])))) + exact.append(bool(info["exact"])) + reasons.append(info["reason"]) + return tuple(weights), tuple(exact), tuple(reasons) + + def _gate_stream_pair_weights(supports, sites, event_weights=None): """Return unordered pair weights induced by a gate/sub-MPO support stream.""" site_rank = {site: pos for pos, site in enumerate(sites)} @@ -525,6 +699,105 @@ def _gate_stream_layout_stats( } +def _gate_stream_compression_stats( + order, + payloads, + supports, + event_types, + *, + event_weights=None, + max_operator_qubits=8, +): + """Estimate MPS cut load from operator-Schmidt ranks over chain cuts. + + This is a static operator-growth bound, not a state-dependent truncation + prediction. It is nevertheless closer to compression pressure than a + pairwise span score because a gate contributes to every chain cut that + separates its support. + """ + order = list(order) + position = {site: pos for pos, site in enumerate(order)} + cut_loads = np.zeros(max(0, len(order) - 1), dtype=float) + total_load = 0.0 + weighted_span = 0.0 + max_span = 0 + exact_events = 0 + bounded_events = 0 + rank_reasons = {} + if event_weights is None: + event_weights = (1.0,) * len(supports) + + for payload, support, _event_type, event_weight in zip( + payloads, supports, event_types, event_weights + ): + support = _unique_ordered(support) + points = [position[site] for site in support if site in position] + if len(points) < 2: + continue + lo, hi = min(points), max(points) + span = hi - lo + max_span = max(max_span, span) + event_weight = max(0.0, float(event_weight)) + weighted_span += event_weight * span + for cut in range(lo, hi): + left = tuple(site for site in support if position[site] <= cut) + right = tuple(site for site in support if position[site] > cut) + if not left or not right: + continue + if payload is None: + info = { + "rank": _operator_schmidt_rank_bound(support, left), + "exact": False, + "reason": "missing_payload", + } + else: + info = _operator_schmidt_rank_info( + payload, + support, + left, + max_operator_qubits=max_operator_qubits, + ) + rank_load = float(np.log2(max(1, info["rank"]))) + rank_load *= event_weight + cut_loads[cut] += rank_load + total_load += rank_load + if info["exact"]: + exact_events += 1 + else: + bounded_events += 1 + rank_reasons[info["reason"]] = ( + rank_reasons.get(info["reason"], 0) + 1 + ) + + max_cut = float(cut_loads.max()) if cut_loads.size else 0.0 + cut_load_l2 = float(np.dot(cut_loads, cut_loads)) if cut_loads.size else 0.0 + mean_cut = float(cut_loads.mean()) if cut_loads.size else 0.0 + # Keep a small span term so two equally loaded layouts still prefer the + # cheaper replay geometry. + loss = float(total_load + cut_load_l2 + 0.05 * weighted_span) + return { + "compression_loss": loss, + "compression_score": loss, + "operator_cut_load": cut_loads, + "max_operator_cut_load": max_cut, + "total_operator_cut_load": float(total_load), + "mean_operator_cut_load": mean_cut, + "operator_cut_load_l2": cut_load_l2, + "weighted_total_span": float(weighted_span), + "max_span": int(max_span), + "rank_exact_events": int(exact_events), + "rank_bounded_events": int(bounded_events), + "rank_exact_cuts": int(exact_events), + "rank_bounded_cuts": int(bounded_events), + "rank_bound_reasons": rank_reasons, + "objective": { + "total_operator_cut_load": 1.0, + "operator_cut_load_l2": 1.0, + "weighted_total_span": 0.05, + }, + } + + def _gate_stream_score_loss(score): """Scalarize the lexicographic layout score for black-box optimizers.""" if isinstance(score, (int, float, np.floating)): @@ -1292,6 +1565,7 @@ def run( self, order="quality", *, + objective="locality", refine_passes=8, refine_numba=True, spectral_dense_max=512, @@ -1304,9 +1578,22 @@ def run( weight_fn=None, weight_mode="auto", schmidt_max_dim=4, + max_operator_qubits=8, ): """Return a layout plan for the stored gate stream.""" order_name = _normalize_gate_stream_layout_order(order) + objective = _gate_stream_layout_objective(objective) + if max_operator_qubits is not None: + try: + max_operator_qubits = int(max_operator_qubits) + except (TypeError, ValueError) as exc: + raise ValueError( + "max_operator_qubits must be a positive integer or None." + ) from exc + if max_operator_qubits < 1: + raise ValueError( + "max_operator_qubits must be a positive integer or None." + ) event_weights = _gate_stream_event_weights( self.payloads, self.supports, @@ -1315,11 +1602,26 @@ def run( weight_mode=weight_mode, schmidt_max_dim=schmidt_max_dim, ) - pair_weights = _gate_stream_pair_weights( + rank_weights, rank_exact, rank_reasons = _gate_stream_event_rank_weights( + self.payloads, self.supports, - self.sites, - event_weights, + self.event_types, + max_operator_qubits=max_operator_qubits, ) + if objective == "compression": + pair_weights = _gate_stream_pair_weights( + self.supports, + self.sites, + rank_weights, + ) + score_event_weights = rank_weights + else: + pair_weights = _gate_stream_pair_weights( + self.supports, + self.sites, + event_weights, + ) + score_event_weights = event_weights include_nevergrad = ( order_name == "auto" or order_name.startswith("nevergrad") ) @@ -1342,16 +1644,30 @@ def run( kahypar_seed=kahypar_seed, ) - candidate_stats = { - name: _gate_stream_layout_stats( + candidate_stats = {} + for name, candidate in candidates.items(): + locality_stats = _gate_stream_layout_stats( candidate, pair_weights, num_events=len(self.supports), supports=self.supports, - event_weights=event_weights, + event_weights=score_event_weights, ) - for name, candidate in candidates.items() - } + stats = dict(locality_stats) + stats["path_loss"] = locality_stats["loss"] + stats["path_score"] = locality_stats["score"] + if objective == "compression": + stats.update(_gate_stream_compression_stats( + candidate, + self.payloads, + self.supports, + self.event_types, + event_weights=score_event_weights, + max_operator_qubits=max_operator_qubits, + )) + stats["loss"] = stats["compression_loss"] + stats["score"] = stats["compression_score"] + candidate_stats[name] = stats if order_name == "auto": selected_order = min( @@ -1383,41 +1699,63 @@ def run( f"{hint}" ) - site_order = tuple(candidates[selected_order]) - site_map = {site: pos for pos, site in enumerate(site_order)} - mapped_where = tuple( - tuple(site_map[site] for site in support) - for support in self.supports - ) - stats = candidate_stats[selected_order] - return { - "kind": "mps_gate_stream_layout", - "selected_order": selected_order, - "qubit_inds": site_order, - "site_order": site_order, - "order": site_order, - "original_sites": self.sites, - "layout": site_map, - "site_map": site_map, - "inverse_site_map": {pos: site for site, pos in site_map.items()}, - "where": self.where, - "mapped_where": mapped_where, - "event_types": self.event_types, - "event_weights": event_weights, - "weight_mode": _normalize_weight_mode(weight_mode), - "stats": stats, - "input_stats": candidate_stats["input"], - "score": stats["score"], + def make_plan(name): + site_order = tuple(candidates[name]) + site_map = {site: pos for pos, site in enumerate(site_order)} + mapped_where = tuple( + tuple(site_map[site] for site in support) + for support in self.supports + ) + stats = candidate_stats[name] + return { + "kind": "mps_gate_stream_layout", + "selected_order": name, + "qubit_inds": site_order, + "site_order": site_order, + "order": site_order, + "original_sites": self.sites, + "layout": site_map, + "site_map": site_map, + "inverse_site_map": {pos: site for site, pos in site_map.items()}, + "where": self.where, + "mapped_where": mapped_where, + "event_types": self.event_types, + "event_weights": event_weights, + "compression_event_weights": rank_weights, + "rank_exact_events": sum(rank_exact), + "rank_bounded_events": len(rank_exact) - sum(rank_exact), + "rank_bound_reasons": { + reason: rank_reasons.count(reason) + for reason in set(rank_reasons) + }, + "weight_mode": _normalize_weight_mode(weight_mode), + "objective": objective, + "max_operator_qubits": max_operator_qubits, + "stats": stats, + "input_stats": candidate_stats["input"], + "score": stats["score"], + } + + candidate_plans = { + name: make_plan(name) for name in candidates + } + selected_plan = dict(candidate_plans[selected_order]) + selected_plan.update({ + "candidate_plans": candidate_plans, "candidate_scores": { name: info["score"] for name, info in candidate_stats.items() }, "candidate_losses": { name: info["loss"] for name, info in candidate_stats.items() }, + "candidate_path_scores": { + name: info["path_score"] for name, info in candidate_stats.items() + }, "candidate_score_tuples": { name: info["score_tuple"] for name, info in candidate_stats.items() }, - } + }) + return selected_plan def map_where(self, where, plan): """Map one original ``where`` through ``plan``.""" diff --git a/src/pepsy/optimizers/mps/optimizer.py b/src/pepsy/optimizers/mps/optimizer.py index 9cf3d6f..78cb939 100644 --- a/src/pepsy/optimizers/mps/optimizer.py +++ b/src/pepsy/optimizers/mps/optimizer.py @@ -56,6 +56,7 @@ from copy import deepcopy from collections.abc import Mapping from numbers import Integral +import time import types import warnings import autoray as ar @@ -934,6 +935,7 @@ def gate_stream_layout( # pylint: disable=too-many-locals sites=None, L=None, order="quality", + objective="locality", refine_passes=8, refine_numba=True, spectral_dense_max=512, @@ -946,6 +948,7 @@ def gate_stream_layout( # pylint: disable=too-many-locals weight_fn=None, weight_mode="auto", schmidt_max_dim=4, + max_operator_qubits=8, ): """Find a good 1D MPS layout for a bundled gate stream. @@ -965,6 +968,10 @@ def gate_stream_layout( # pylint: disable=too-many-locals inferred from first use in ``gate_stream`` unless ``L`` is given. L : int | None Convenience for ``sites=range(L)``. + objective : {"locality", "compression"} + ``"locality"`` minimizes support span and cut congestion using + event weights. ``"compression"`` ranks layouts by operator- + Schmidt load over the MPS cuts, with path span as a tie-breaker. order : str One of ``"quality"``/``"auto"``/``"best"``, ``"recursive"``, ``"input"``, ``"degree"``, ``"bfs"``, ``"spectral"``, @@ -998,6 +1005,11 @@ def gate_stream_layout( # pylint: disable=too-many-locals small dense gates, falling back to count weights. schmidt_max_dim : int Maximum local dimension for the optional operator-Schmidt proxy. + max_operator_qubits : int | None + Maximum support size for exact dense rank probes in the + compression objective. Larger or opaque operators use a + conservative operator-space rank bound and are marked as bounded + in the returned diagnostics. Returns ------- @@ -1010,6 +1022,7 @@ def gate_stream_layout( # pylint: disable=too-many-locals finder = cls.LayoutFinder(gate_stream, sites=sites, L=L) return finder.run( order=order, + objective=objective, refine_passes=refine_passes, refine_numba=refine_numba, spectral_dense_max=spectral_dense_max, @@ -1022,6 +1035,7 @@ def gate_stream_layout( # pylint: disable=too-many-locals weight_fn=weight_fn, weight_mode=weight_mode, schmidt_max_dim=schmidt_max_dim, + max_operator_qubits=max_operator_qubits, ) @classmethod @@ -1040,6 +1054,128 @@ def current_gate_stream_layout(self, *, sites=None, L=None, **kwargs): return self.layout_finder(sites=sites, L=L).run(**kwargs) + def select_layout_for_compression( + self, + *, + sites=None, + L=None, + layout_kwargs=None, + pilot_candidates=4, + pilot_steps=None, + cutoff=1e-12, + cutoff_mode="rsum2", + run_kwargs=None, + ): + """Select an MPS layout using a bounded, state-aware pilot replay. + + The finder first produces cheap static candidates with + ``objective="compression"``. The best ``pilot_candidates`` are then + replayed on independent copies of the current MPS using the real + execution mode, ``chi``, cutoff, and backend. The returned plan is + non-mutating and contains ``pilot`` diagnostics for every candidate. + + This method is intentionally separate from :meth:`run`: layout + selection can be expensive and should be explicit in production + workflows. ``pilot_steps`` limits the replay prefix while preserving + the original optimizer and gate queue. + """ + if self.mode == "exact": + raise ValueError( + "compression layout pilots require an MPS compression mode, " + "not mode='exact'." + ) + if self._persistent_layout_plan is not None: + raise ValueError( + "compression layout pilots require an optimizer without a " + "persistent layout; create the pilot before apply_layout()." + ) + try: + pilot_candidates = int(pilot_candidates) + except (TypeError, ValueError) as exc: + raise ValueError("pilot_candidates must be a positive integer.") from exc + if pilot_candidates < 1: + raise ValueError("pilot_candidates must be a positive integer.") + if pilot_steps is not None: + try: + pilot_steps = int(pilot_steps) + except (TypeError, ValueError) as exc: + raise ValueError("pilot_steps must be a positive integer or None.") from exc + if pilot_steps < 1: + raise ValueError("pilot_steps must be a positive integer or None.") + + finder = self.layout_finder(sites=sites, L=L) + kwargs = dict(layout_kwargs or {}) + kwargs["objective"] = "compression" + static_plan = finder.run(**kwargs) + candidates = dict(static_plan.get("candidate_plans", {})) + if not candidates: + candidates = {static_plan["selected_order"]: static_plan} + ranked_names = sorted( + candidates, + key=lambda name: candidates[name]["stats"].get( + "compression_score", candidates[name]["stats"].get("score", 0.0) + ), + )[:pilot_candidates] + + base_run_kwargs = dict(run_kwargs or {}) + base_run_kwargs.setdefault("progbar", False) + base_run_kwargs.setdefault("layout_report", False) + base_run_kwargs.setdefault("cutoff", cutoff) + base_run_kwargs.setdefault("cutoff_mode", cutoff_mode) + # The selector ranks candidates by measured retained fidelity. Ensure + # that diagnostic trace is available even when the source optimizer + # was created with track_infidelity=False. + base_run_kwargs["track_infidelity"] = True + pilot_reports = {} + successful = [] + for name in ranked_names: + trial = self.copy() + if pilot_steps is not None: + trial.G = trial.G[:pilot_steps] + trial.where = trial.where[:pilot_steps] + trial.event_types = trial.event_types[:pilot_steps] + started = time.perf_counter() + try: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + trial.run(layout=candidates[name], **base_run_kwargs) + elapsed = time.perf_counter() - started + infidelity = float(trial.infidelities[-1]) + final_bond = int(trial.p.max_bond()) + report = { + "status": "ok", + "elapsed_seconds": float(elapsed), + "final_bond": final_bond, + "infidelity": infidelity, + "pilot_steps": len(trial.G), + } + successful.append((infidelity, final_bond, elapsed, name)) + except Exception as exc: # pragma: no cover - backend-specific + report = { + "status": "error", + "error": f"{type(exc).__name__}: {exc}", + "elapsed_seconds": float(time.perf_counter() - started), + "pilot_steps": len(trial.G), + } + pilot_reports[name] = report + + if not successful: + raise RuntimeError( + "All MPS compression layout pilot candidates failed. " + f"Diagnostics: {pilot_reports!r}" + ) + selected_name = min(successful)[-1] + selected = dict(candidates[selected_name]) + selected["selected_order"] = selected_name + selected["pilot"] = { + "objective": "compression", + "pilot_candidates": tuple(ranked_names), + "selected_order": selected_name, + "reports": pilot_reports, + } + selected["candidate_plans"] = candidates + return selected + def plot_layout( self, plan=None, @@ -2056,12 +2192,13 @@ def _layout_report_text(cls, plan): selected = plan.get("selected_order", "") site_order = plan.get("site_order", plan.get("qubit_inds", ())) weight_mode = plan.get("weight_mode", "count") + objective = plan.get("objective", "locality") lines = [ ( "MpsOptimizer layout finder: " f"order={selected}, sites={len(site_order)}, " f"events={stats.get('num_events', input_stats.get('num_events', 0))}, " - f"weight_mode={weight_mode}" + f"weight_mode={weight_mode}, objective={objective}" ), ( " long-range events: " @@ -2103,6 +2240,19 @@ def _layout_report_text(cls, plan): ) ), ] + if objective == "compression": + lines.append( + " operator cut load max/total: " + + cls._format_layout_value( + stats.get("max_operator_cut_load", 0.0) + ) + + "/" + + cls._format_layout_value( + stats.get("total_operator_cut_load", 0.0) + ) + + " | bounded cut probes: " + + cls._format_layout_value(stats.get("rank_bounded_cuts", 0)) + ) return "\n".join(lines) def run( # pylint: disable=too-many-arguments,too-many-positional-arguments diff --git a/src/pepsy/optimizers/tree/layout.py b/src/pepsy/optimizers/tree/layout.py index d9f2e8b..f0841d7 100644 --- a/src/pepsy/optimizers/tree/layout.py +++ b/src/pepsy/optimizers/tree/layout.py @@ -37,6 +37,8 @@ _gate_stream_spectral_order, _normalize_layout_gate_queue, _normalize_layout_support, + _operator_schmidt_rank_bound, + _operator_schmidt_rank_info as _mps_operator_schmidt_rank_info, _normalize_weight_mode, ) from ..mps.optimizer import _control_event_parts as _mps_control_event_parts @@ -174,12 +176,15 @@ def _normalize_layout_objective(objective): "bond": "congestion", "bond_load": "congestion", "combined": "hybrid", + "compress": "compression", + "accuracy": "compression", + "bond_growth": "compression", } name = aliases.get(name, name) - if name not in {"path", "congestion", "hybrid"}: + if name not in {"path", "congestion", "hybrid", "compression"}: raise ValueError( f"Unknown tree layout objective {objective!r}. " - "Expected 'path', 'congestion', or 'hybrid'." + "Expected 'path', 'congestion', 'compression', or 'hybrid'." ) return name @@ -191,12 +196,13 @@ def _operator_schmidt_rank(payload, support, left_support): left_set = set(left_support) if not left_set or left_set == set(support): return 1 + default_bound = _operator_schmidt_rank_bound(support, left_support) try: array = ar.to_numpy(payload) except Exception: - return 2 + return default_bound if array.size != 4 ** len(support): - return 2 + return default_bound try: array = array.reshape((2,) * (2 * len(support))) positions = {site: pos for pos, site in enumerate(support)} @@ -216,7 +222,7 @@ def _operator_schmidt_rank(payload, support, left_support): ) return max(1, int(np.linalg.matrix_rank(matrix))) except (TypeError, ValueError, np.linalg.LinAlgError): - return 2 + return default_bound def _submpo_schmidt_rank_bound(payload, support, left_support): @@ -1101,7 +1107,7 @@ class TreeLayoutFinder: (see :meth:`TreePlan.from_order`). dense_max : int Maximum subsystem size for dense spectral reordering. - objective : {"path", "congestion", "hybrid"} + objective : {"path", "congestion", "compression", "hybrid"} Layout objective. `"path"` preserves the co-occurrence/path-length heuristic; `"congestion"` selects among layout candidates using the predicted operator-Schmidt load on tree edges. `"hybrid"` combines @@ -1250,6 +1256,7 @@ def __init__(self, gates=None, n=None, *, supports=None, structure="quality", # fresh dictionaries from the public diagnostic methods below. self._plan_cache = {} self._edge_load_cache = {} + self._rank_diagnostics_cache = {} self._schmidt_rank_cache = {} self._similarity_cache = {} self._congestion_weights_cache = None @@ -1492,12 +1499,61 @@ def _hybrid_key(self, plan): int(max_path), ) + def _tensor_cost_key(self, plan): + """Return a chi-scaled proxy for local TTN tensor cost. + + A wider node reduces geodesic distance but increases the number of + virtual legs on one tensor. The exact contraction cost depends on + the realized bond dimensions, so this uses the configured ``chi`` (or + a conservative qubit bond of two) to rank structures without ever + allocating tensors. + """ + chi = max(2, int(self.chi or 2)) + log_chi = float(np.log2(chi)) + degrees = [] + log_sizes = [] + for node, children in plan.children.items(): + if not children: + continue + virtual_degree = len(children) + (1 if node in plan.parent else 0) + physical_legs = 1 if node in plan.qubit_of_node else 0 + degrees.append(virtual_degree) + log_sizes.append(virtual_degree * log_chi + physical_legs) + if not degrees: + return (0.0, 0.0, 0, 0) + max_log_size = max(log_sizes) + # log2(sum(2**log_size)) without overflowing for large chi/arity. + shifted = np.asarray(log_sizes, dtype=float) - max_log_size + total_log_size = max_log_size + float(np.log2(np.exp2(shifted).sum())) + return ( + float(max_log_size), + float(total_log_size), + int(max(degrees)), + int(sum(degrees)), + ) + def _objective_key(self, plan): """Return the selected objective's deterministic comparison key.""" if self.objective == "path": return self._path_score_and_max(plan) if self.objective == "congestion": return self._congestion_key(plan) + if self.objective == "compression": + loads = self.edge_loads(plan) + values = tuple(loads.values()) + tensor_cost = self._tensor_cost_key(plan) + return ( + max(values, default=0.0), + sum(values), + tensor_cost[0], + tensor_cost[1], + self.score(plan), + max( + (plan.tree_distance(a, b) for a in range(self.n) + for b in range(a + 1, self.n)), + default=0, + ), + ) return self._hybrid_key(plan) def _selection_key(self, plan, chi): @@ -1512,7 +1568,7 @@ def _selection_loss(self, plan, chi): key = self._objective_key(plan) if self.objective == "path": value = key[0] - elif self.objective == "congestion": + elif self.objective in {"congestion", "compression"}: value = key[0] + 1.0e-6 * key[1] + 1.0e-12 * key[2] else: value = key[0] @@ -1525,6 +1581,9 @@ def _discard_plan_cache(self, plan): cached = self._edge_load_cache.get(id(plan)) if cached is not None and cached[0] is plan: del self._edge_load_cache[id(plan)] + cached = self._rank_diagnostics_cache.get(id(plan)) + if cached is not None and cached[0] is plan: + del self._rank_diagnostics_cache[id(plan)] def _refine_plan_greedy(self, plan, *, chi, budget, progbar=False): """Greedily improve a fixed topology through adjacent leaf swaps.""" @@ -1763,15 +1822,26 @@ def _build_plan(self, weights, *, structure=None, return plan def _schmidt_rank(self, payload, support, left_support): - """Return a cached operator-Schmidt rank for layout diagnostics.""" + """Return a cached numeric operator-Schmidt rank or bound.""" + return self._schmidt_rank_info(payload, support, left_support)["rank"] + + def _schmidt_rank_info(self, payload, support, left_support): + """Return rank metadata used by compression diagnostics. + + ``exact=False`` is deliberate for opaque native arrays, MPO bond + bounds, and supports larger than ``max_operator_qubits``. The numeric + rank is then a conservative operator-space bound, never an optimistic + hard-coded rank-two fallback. + """ if ( self.max_operator_qubits is not None and len(support) > self.max_operator_qubits ): - # Keep layout search bounded. This is a conservative rank proxy; - # callers that need exact wide-operator layout costs can opt out - # with max_operator_qubits=None. - return 2 + return { + "rank": _operator_schmidt_rank_bound(support, left_support), + "exact": False, + "reason": "max_operator_qubits", + } # For an ordinary dense gate, its Schmidt rank depends on the operator # data and *wire positions* in ``support``, not on the global qubit # labels. Reusing the same CNOT/CZ/parameterized matrix across many @@ -1794,10 +1864,21 @@ def _schmidt_rank(self, payload, support, left_support): if cached is not None and cached[0] is payload: return cached[1] rank = _submpo_schmidt_rank_bound(payload, support, left_support) - if rank is None: - rank = _operator_schmidt_rank(payload, support, left_support) - self._schmidt_rank_cache[key] = (payload, rank) - return rank + if rank is not None: + info = { + "rank": int(rank), + "exact": False, + "reason": "mpo_bond_bound", + } + else: + info = _mps_operator_schmidt_rank_info( + payload, + support, + left_support, + max_operator_qubits=self.max_operator_qubits, + ) + self._schmidt_rank_cache[key] = (payload, info) + return info def _candidate_plans(self, max_arity): """Build the candidate plans considered by the selected objective.""" @@ -2019,6 +2100,14 @@ def recommend_layered( "max_path": report["max_path"], "max_edge_load": report["max_edge_load"], "peak_bond_growth": report["peak_bond_growth"], + "max_virtual_degree": report["max_virtual_degree"], + "total_virtual_degree": report["total_virtual_degree"], + "estimated_max_tensor_log2": report[ + "estimated_max_tensor_log2" + ], + "estimated_total_tensor_log2": report[ + "estimated_total_tensor_log2" + ], **_chi_cut_fields(plan, chi), "order": self._leaf_order(plan), "planning": planning, @@ -2119,6 +2208,38 @@ def run( self._selected_candidate = selected return candidates[selected] + def candidate_plans(self, *, chi=_DEFAULT_CHI): + """Return immutable candidate plans for optional pilot replay. + + The normal :meth:`run` path remains static and cheap. This method + exposes the interaction, congestion, balanced, and arity candidates + that a state-aware pilot can compare without rebuilding the finder. + Candidate names are stable strings such as + ``"congestion:arity=2"``. + """ + if chi is _DEFAULT_CHI: + chi = self.chi + else: + chi = _validate_chi(chi) + arities = ( + tuple(self.arity_candidates) + if self.arity_candidates is not None + else (self.max_arity,) + ) + result = {} + for arity in arities: + plans = self._candidate_plans(arity) + for name, plan in plans.items(): + key = f"{name}:arity={arity}" + result[key] = { + "plan": plan, + "objective_key": self._selection_key(plan, chi), + "path_score": self.score(plan), + "tensor_cost": self._tensor_cost_key(plan), + "edge_loads": self.edge_loads(plan), + } + return result + def recommend_arities( self, max_arities=(2, 3, 4), @@ -2217,10 +2338,21 @@ def recommend_arities( ), default=0, ), + "total_virtual_degree": sum( + len(children) + (1 if node in plan.parent else 0) + for node, children in plan.children.items() + if children + ), "score": report["score"], "max_path": report["max_path"], "max_edge_load": report["max_edge_load"], "peak_bond_growth": report["peak_bond_growth"], + "estimated_max_tensor_log2": report[ + "estimated_max_tensor_log2" + ], + "estimated_total_tensor_log2": report[ + "estimated_total_tensor_log2" + ], **_chi_cut_fields(plan, chi), "order": self._leaf_order(plan), "planning": planning, @@ -2309,6 +2441,11 @@ def edge_loads(self, plan=None): for parent, children in plan.children.items() for child in children } + rank_diagnostics = { + "exact_events": 0, + "bounded_events": 0, + "reasons": {}, + } for payload, support, event_type in zip( self.payloads, self.supports, self.event_types ): @@ -2354,14 +2491,24 @@ def edge_loads(self, plan=None): left = tuple( site for site in support if left_mask & (1 << site) ) - rank = ( - self._schmidt_rank(payload, support, left) - if payload is not None else 2 - ) - loads[edge] += float(np.log2(rank)) + info = self._schmidt_rank_info(payload, support, left) + rank = int(info["rank"]) + loads[edge] += float(np.log2(max(1, rank))) + if info["exact"]: + rank_diagnostics["exact_events"] += 1 + else: + rank_diagnostics["bounded_events"] += 1 + reason = info["reason"] + rank_diagnostics["reasons"][reason] = ( + rank_diagnostics["reasons"].get(reason, 0) + 1 + ) # Retain the plan alongside its id so a future id reuse cannot return # diagnostics for an unrelated short-lived plan. self._edge_load_cache[cache_key] = (plan, dict(loads)) + self._rank_diagnostics_cache[cache_key] = ( + plan, + rank_diagnostics, + ) return dict(loads) def _congestion_key(self, plan): @@ -2457,9 +2604,11 @@ def report(self, plan=None, *, include_edge_loads=True): if include_edge_loads: loads = self.edge_loads(plan) balanced_loads = self.edge_loads(balanced) + rank_info = self._rank_diagnostics_cache.get(id(plan), (plan, {}))[1] else: loads = None balanced_loads = None + rank_info = {} max_load = max(loads.values(), default=0.0) if loads is not None else None total_load = sum(loads.values()) if loads is not None else None balanced_max_load = ( @@ -2479,6 +2628,8 @@ def report(self, plan=None, *, include_edge_loads=True): arity_histogram[len(children)] = ( arity_histogram.get(len(children), 0) + 1 ) + tensor_cost = self._tensor_cost_key(plan) + objective_key = self._objective_key(plan) return { "n_qubits": self.n, "n_interacting_pairs": n_pairs, @@ -2488,6 +2639,12 @@ def report(self, plan=None, *, include_edge_loads=True): self.hybrid_weights if self.objective == "hybrid" else None ), "hybrid_cost": hybrid_cost, + "objective_key": objective_key, + "path_score": float(weighted_sum), + "compression_score": ( + float(objective_key[0] + objective_key[1]) + if self.objective == "compression" else None + ), "root": plan.root, "root_qubit": plan.root_qubit, "is_binary": plan.is_binary(), @@ -2528,6 +2685,13 @@ def report(self, plan=None, *, include_edge_loads=True): float(balanced_max_load) if balanced_max_load is not None else None ), + "rank_exact_events": int(rank_info.get("exact_events", 0)), + "rank_bounded_events": int(rank_info.get("bounded_events", 0)), + "rank_bound_reasons": dict(rank_info.get("reasons", {})), + "max_virtual_degree": tensor_cost[2], + "total_virtual_degree": tensor_cost[3], + "estimated_max_tensor_log2": tensor_cost[0], + "estimated_total_tensor_log2": tensor_cost[1], "selected_candidate": getattr(self, "_selected_candidate", "interaction"), "candidate_scores": getattr(self, "_last_candidate_scores", {}), } @@ -2929,11 +3093,11 @@ def plot_tent( ax=None, figsize=(8, 7), cmap="turbo", - edge_cmap="turbo", + edge_cmap="GnBu", node_cmap="YlOrRd", color_by="scale", edge_color="#2f80a0", - show_edge_arrows=True, + show_edge_arrows=False, arrow_size=8.0, order=True, lattice=True, @@ -2953,10 +3117,12 @@ def plot_tent( Physical sites and gate connectivity stay in the lower, grey raw graph. Internal TTN nodes are lifted above the mean position of their descendant sites, and each parent-child hierarchy edge uses one - uniform solid color by default. Pass ``edge_color=None`` to color - edges by ``edge_cmap`` and ``color_by``. Small arrows at edge - midpoints show the parent-to-child direction by default; disable them - with ``show_edge_arrows=False``. This is deliberately a structural + uniform solid color by default. Pass ``edge_color=None`` to match + each incoming edge to the node it terminates at (so ``node_cmap`` + controls both). The default has no arrows, matching Cotengra's + structural tent view; pass + ``show_edge_arrows=True`` only when parent-to-child direction is + needed. This is deliberately a structural visualization: gate-by-gate route overlays are not drawn. Set ``order=True`` to place hierarchy nodes by a deterministic post-order traversal, matching the @@ -3053,7 +3219,10 @@ def gather_qubits(node): ) y_max = max(point[1] for point in coords.values()) if vertical_spacing is None: - vertical_spacing = max(0.7, 0.28 * x_span) + # Keep the tent compact for square 2-D lattices. The previous + # spacing made a 6x6 lattice grow into a very tall strip even + # though ``figsize`` only changes the canvas, not the geometry. + vertical_spacing = max(0.55, 0.16 * x_span) vertical_spacing = float(vertical_spacing) if vertical_spacing <= 0.0: raise ValueError("vertical_spacing must be positive.") @@ -3119,22 +3288,21 @@ def node_color(node): colormaps, node_cmap, node_scales[node], n_colors ) - def hierarchy_edge_color(parent): + def hierarchy_edge_color(parent, child): if edge_color is not None: return edge_color - if color_by == "order": - return event_color( - colormaps, cmap, order_values[parent], n_colors - ) - return scale_color( - colormaps, edge_cmap, node_scales[parent], n_colors - ) + # ``None`` means "follow the node palette": this is intentionally + # the node color itself rather than a separate edge colormap, so + # an incoming edge and its child are visually identical. + return node_color(child) for parent, children in plan.children.items(): for child in children: x0, y0 = positions[parent] x1, y1 = positions[child] - edge_color_value = hierarchy_edge_color(parent) + # The incoming edge is colored like the node it terminates at, + # making each scale/order layer visually self-consistent. + edge_color_value = hierarchy_edge_color(parent, child) ax.plot( (x0, x1), (y0, y1), @@ -3208,7 +3376,7 @@ def hierarchy_edge_color(parent): colormaps, ScalarMappable, Normalize, - cmap if color_by == "order" else edge_cmap, + cmap if color_by == "order" else node_cmap, n_colors, label=( "tree order" if color_by == "order" else "tree scale" @@ -3243,8 +3411,8 @@ def plot_rubberband( site_coords=None, ax=None, figsize=(10, 8), - cmap="turbo", - color_by="scale", + cmap="Spectral", + color_by="gate", scale_cmap="viridis", lattice=True, show_gate_connectivity=True, @@ -3262,9 +3430,10 @@ def plot_rubberband( This is the physical-lattice counterpart to Quimb's contraction-tree ``plot_rubberband`` view: the lattice and gate connectivity remain grey, while each non-root tree cluster is wrapped by a rounded, - translucent colored band. ``color_by="gate"`` colors the bands by a - deterministic post-order through the tree; ``color_by="scale"`` uses - one stable color for each tree scale measured from the leaves. + translucent colored band. The default ``color_by="gate"`` uses a + ``Spectral`` post-order progression, matching Cotengra's many-color + rubberband view. ``color_by="scale"`` is available when one stable + color is wanted for each tree scale measured from the leaves. The default presentation has no axes, site labels, or title. It returns a normal Matplotlib ``(fig, ax)`` pair for further styling. @@ -3389,7 +3558,10 @@ def visit(node): edgecolor=color, linewidth=band_linewidth, alpha=band_alpha, - zorder=3, + # Draw inner/earlier contractions above outer/later + # bands, as in Cotengra, so overlapping bands remain + # individually legible. + zorder=3.0 + (n_bands - band_index) / n_bands, ) ) diff --git a/src/pepsy/optimizers/tree/optimizer.py b/src/pepsy/optimizers/tree/optimizer.py index 0fb352a..d652e48 100644 --- a/src/pepsy/optimizers/tree/optimizer.py +++ b/src/pepsy/optimizers/tree/optimizer.py @@ -41,6 +41,7 @@ from copy import deepcopy import heapq from numbers import Integral +import time import warnings import autoray as ar @@ -323,10 +324,12 @@ class TreeOptimizer: optimizer's ``chi`` (structures that stay exact at ``chi`` are preferred). Pass ``max_arity=2`` to force a fixed binary tree. Ignored when an explicit ``tree`` is supplied. - layout_objective : {"path", "congestion", "hybrid"} + layout_objective : {"path", "congestion", "compression", "hybrid"} Objective used when building an automatic tree. ``"path"`` is the backward-compatible interaction-path heuristic; ``"congestion"`` selects a candidate using predicted operator-Schmidt edge load; + ``"compression"`` additionally penalizes peak/total load and the + estimated local tensor cost at ``chi``; ``"hybrid"`` combines normalized path, peak-load, and total-load costs. Pass a configured :class:`TreeLayoutFinder` through ``layout=`` to customize its hybrid weights or enable pre-simulation refinement. @@ -1423,6 +1426,149 @@ def layout_report(self): } return self.layout_finder.report(self.plan) + def select_layout_for_compression( + self, + *, + pilot_candidates=4, + pilot_steps=None, + install=False, + progbar=False, + ): + """Select a tree layout using state-aware pilot replay. + + Static compression candidates are generated with + ``objective="compression"`` and then replayed on independent copies + of the current state. The pilot uses the real tree update kernels, + ``chi``, cutoff, backend, and queued gate stream. The original state + is unchanged. By default the selected plan is returned for explicit + hand-off; ``install=True`` is allowed only for a product state and + remounts that state exactly on the selected geometry. + """ + try: + pilot_candidates = int(pilot_candidates) + except (TypeError, ValueError) as exc: + raise ValueError("pilot_candidates must be a positive integer.") from exc + if pilot_candidates < 1: + raise ValueError("pilot_candidates must be a positive integer.") + if pilot_steps is not None: + try: + pilot_steps = int(pilot_steps) + except (TypeError, ValueError) as exc: + raise ValueError("pilot_steps must be a positive integer or None.") from exc + if pilot_steps < 1: + raise ValueError("pilot_steps must be a positive integer or None.") + + finder = TreeLayoutFinder( + gates=self._layout_gate_stream(), + n=self.n, + structure=self.structure, + max_arity=self.max_arity, + community_frac=self.community_frac, + star_frac=self.star_frac, + objective="compression", + weight_mode=self.layout_weight_mode, + chi=self.chi, + max_operator_qubits=self.max_operator_qubits, + root_qubit=self.plan.root_qubit, + ) + candidates = finder.candidate_plans(chi=self.chi) + ranked = sorted( + candidates, + key=lambda name: candidates[name]["objective_key"], + )[:pilot_candidates] + reports = {} + successful = [] + for name in ranked: + plan = candidates[name]["plan"] + if not _is_product_tensor_network(self.tn): + raise ValueError( + "Tree compression pilots require a product initial state " + "when comparing different tree geometries. Convert the " + "entangled state explicitly onto each candidate plan first." + ) + started = time.perf_counter() + trial = type(self)( + None, + n=self.n, + chi=self.chi, + cutoff=self.cutoff, + cutoff_mode=self.cutoff_mode, + mode=self.mode, + structure=self.structure, + max_arity=self.max_arity, + community_frac=self.community_frac, + star_frac=self.star_frac, + tree=plan, + dtype=self.dtype, + threads=self.threads, + track_truncation=True, + track_infidelity=True, + max_intermediate_bond=self.max_intermediate_bond, + max_operator_qubits=self.max_operator_qubits, + max_subtree_nodes=self.max_subtree_nodes, + record_history=self.record_history, + run=False, + tn=self.tn, + ) + trial.G = list(self.G) + trial.where = list(self.where) + trial.event_types = list(self.event_types) + if pilot_steps is not None: + trial.G = trial.G[:pilot_steps] + trial.where = trial.where[:pilot_steps] + trial.event_types = trial.event_types[:pilot_steps] + try: + trial.run(progbar=progbar) + elapsed = time.perf_counter() - started + infidelity = float(trial.infidelities[-1]) + final_bond = int(trial.max_bond()) + truncated_edges = int(sum( + event.get("truncated", False) + for event in trial.truncation_history + )) + reports[name] = { + "status": "ok", + "elapsed_seconds": float(elapsed), + "infidelity": infidelity, + "final_bond": final_bond, + "truncated_edges": truncated_edges, + "pilot_steps": len(trial.G), + } + successful.append((infidelity, truncated_edges, final_bond, elapsed, name)) + except Exception as exc: # pragma: no cover - backend-specific + reports[name] = { + "status": "error", + "error": f"{type(exc).__name__}: {exc}", + "elapsed_seconds": float(time.perf_counter() - started), + "pilot_steps": len(trial.G), + } + + if not successful: + raise RuntimeError( + "All Tree compression layout pilot candidates failed. " + f"Diagnostics: {reports!r}" + ) + selected_name = min(successful)[-1] + selected_plan = candidates[selected_name]["plan"] + if install: + self.plan = selected_plan + self.tn = self._remount_product_state(self.tn) + self.center = self.plan.root + self.layout_finder = finder + self.layout_objective = "compression" + return { + "plan": selected_plan, + "selected_candidate": selected_name, + "candidates": candidates, + "pilot": { + "objective": "compression", + "pilot_candidates": tuple(ranked), + "selected_candidate": selected_name, + "reports": reports, + "installed": bool(install), + }, + } + def plot_layout(self, plan=None, *, layout_kwargs=None, **plot_kwargs): """Plot the tree layout as a Cotengra-style tent. diff --git a/tests/test_optimize_mps.py b/tests/test_optimize_mps.py index d08d551..8fc91f4 100644 --- a/tests/test_optimize_mps.py +++ b/tests/test_optimize_mps.py @@ -682,6 +682,51 @@ def test_mps_optimizer_gate_stream_layout_remaps_long_range_path(): ) +def test_mps_compression_layout_reports_operator_cut_load(): + """Compression objective exposes cut-load diagnostics and rank bounds.""" + gate = np.eye(8, dtype=complex) + plan = py.MpsOptimizer.gate_stream_layout( + [(gate, (0, 1, 2))], + L=3, + objective="compression", + max_operator_qubits=2, + ) + + assert plan["objective"] == "compression" + assert plan["stats"]["compression_score"] == plan["score"] + assert plan["rank_bounded_events"] > 0 + assert plan["rank_bound_reasons"]["max_operator_qubits"] > 0 + assert plan["candidate_plans"] + + exact = py.MpsOptimizer.gate_stream_layout( + [(qu.CNOT(), (0, 1))], + L=2, + objective="compression", + ) + assert exact["stats"]["rank_exact_events"] == 1 + assert exact["stats"]["total_operator_cut_load"] == pytest.approx(1.0) + + +def test_mps_compression_layout_pilot_is_non_mutating(): + """Pilot selection uses copied state and does not install a layout.""" + p0 = qtn.MPS_computational_state("0000", dtype="complex128") + gates = [(qu.CNOT(), (0, 3)), (qu.CNOT(), (3, 1))] + opt = py.MpsOptimizer( + p0, gates=gates, chi=2, mode="svd", track_infidelity=False + ) + before = opt.to_dense() + + selected = opt.select_layout_for_compression( + pilot_candidates=1, + pilot_steps=1, + ) + + assert selected["pilot"]["selected_order"] + assert selected["pilot"]["reports"] + assert opt._persistent_layout_plan is None + assert np.allclose(opt.to_dense(), before) + + def test_mps_layout_finder_plot_draws_lattice_and_gate_order(): """The MPS plot exposes the lattice, gate graph, and colored chain.""" matplotlib = pytest.importorskip("matplotlib") diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index 8122d75..86f43ea 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -974,6 +974,44 @@ def test_congestion_layout_uses_operator_schmidt_edge_load(): assert report["peak_bond_growth"] == pytest.approx(8.0) +def test_compression_layout_reports_rank_bounds_and_tensor_cost(): + """Compression selection is explicit and honest for wide operators.""" + gate = np.eye(8, dtype=complex) + finder = TreeLayoutFinder( + [(gate, (0, 1, 2))], + n=3, + objective="compression", + max_arity=(2, 3), + max_operator_qubits=2, + chi=2, + ) + plan = finder.run() + report = finder.report(plan) + + assert report["objective"] == "compression" + assert report["rank_bounded_events"] > 0 + assert report["rank_bound_reasons"]["max_operator_qubits"] > 0 + assert report["estimated_max_tensor_log2"] >= 0.0 + assert len(report["objective_key"]) >= 5 + + +def test_tree_compression_layout_pilot_is_non_mutating(): + """Tree pilot selection compares copied product states only.""" + gates = [(pepsy.cnot(), (0, 3)), (pepsy.cnot(), (3, 1))] + opt = TreeOptimizer(gates, n=4, chi=2, run=False) + original_plan = opt.plan + + selected = opt.select_layout_for_compression( + pilot_candidates=1, + pilot_steps=1, + ) + + assert selected["selected_candidate"] + assert selected["pilot"]["reports"] + assert opt.plan is original_plan + assert opt.max_bond() == 1 + + def test_tree_edge_loads_match_full_edge_reference(): """Steiner-only edge scanning preserves the full congestion calculation.""" rng = np.random.default_rng(109) @@ -1505,7 +1543,7 @@ def test_tree_layout_finder_plot_defaults_to_tent(): assert fig is ax.figure assert ax.get_title() == "" - assert len(ax.patches) == len(plan.nodes()) - 1 + assert not ax.patches assert len(fig.axes) == 1 assert not ax.axison # schematic-style presentation by default assert not ax.texts @@ -1558,7 +1596,7 @@ def test_tree_layout_finder_plot_tent_draws_hierarchy_over_raw_graph(): assert plan.is_binary() assert fig is ax.figure - assert len(ax.patches) == len(plan.nodes()) - 1 + assert not ax.patches assert not ax.texts assert not ax.axison assert len(ax.lines) >= len(plan.nodes()) - 1 @@ -1646,7 +1684,49 @@ def test_tree_layout_tent_edges_are_uniform_by_default(): line.get_color() for line in ax.lines[background_lines:] } assert hierarchy_colors == {"#2f80a0"} - assert len(ax.patches) == len(plan.nodes()) - 1 + assert not ax.patches + plt.close(fig) + + +def test_tree_layout_tent_colored_edges_match_child_nodes(): + """Colored incoming edges use the same scale color as their child node.""" + matplotlib = pytest.importorskip("matplotlib") + matplotlib.use("Agg", force=True) + import matplotlib.pyplot as plt + + finder = TreeLayoutFinder( + [(pepsy.cnot(), (0, 3)), (pepsy.cnot(), (1, 2))], + n=4, + max_arity=2, + ) + plan = finder.run() + fig, ax = finder.plot_tent( + plan, + color_by="scale", + edge_color=None, + show_edge_arrows=False, + ) + + lattice_pairs = { + frozenset((site, site + 1)) for site in range(plan.n - 1) + } + background_lines = plan.n - 1 + sum( + frozenset(where) not in lattice_pairs + for _, where in [(pepsy.cnot(), (0, 3)), (pepsy.cnot(), (1, 2))] + ) + node_colors = { + node: tuple(collection.get_facecolors()[0]) + for node, collection in zip(plan.nodes(), ax.collections) + } + hierarchy_lines = ax.lines[background_lines:] + line_index = 0 + for parent, children in plan.children.items(): + for child in children: + assert tuple(hierarchy_lines[line_index].get_color()) == pytest.approx( + node_colors[child] + ) + line_index += 1 + assert line_index == len(hierarchy_lines) plt.close(fig) @@ -1687,6 +1767,33 @@ def test_tree_layout_finder_plot_rubberband_is_axis_free_and_unlabeled(): plt.close(fig) +def test_tree_layout_rubberband_defaults_to_cotengra_ordered_colors(): + """Default rubberbands use distinct post-order Spectral colors.""" + matplotlib = pytest.importorskip("matplotlib") + matplotlib.use("Agg", force=True) + import matplotlib.pyplot as plt + + finder = TreeLayoutFinder( + [(pepsy.cnot(), (0, 3)), (pepsy.cnot(), (1, 2))], + n=4, + max_arity=2, + ) + fig, ax = finder.plot_rubberband( + finder.run(), + site_coords={0: (0, 0), 1: (1, 0), 2: (0, 1), 3: (1, 1)}, + ) + + expected = matplotlib.colormaps["Spectral"] + assert np.allclose( + ax.patches[0].get_edgecolor()[:3], expected(0.0)[:3] + ) + assert np.allclose( + ax.patches[-1].get_edgecolor()[:3], expected(1.0)[:3] + ) + assert ax.patches[0].get_zorder() > ax.patches[-1].get_zorder() + plt.close(fig) + + def test_tree_optimizer_plot_layout_with_explicit_plan_is_non_mutating(): """The tree optimizer wrapper plots an explicit plan without replay.""" matplotlib = pytest.importorskip("matplotlib") From 6a8a29829caafa2b666a98eb033ba2c6d5266f29 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Thu, 30 Jul 2026 16:24:29 -0700 Subject: [PATCH 34/70] add Tree quality layout order --- docs/api/optimizers/tree.md | 8 +++++ src/pepsy/optimizers/tree/layout.py | 49 ++++++++++++++++++++++++++++- tests/test_optimize_tree.py | 30 ++++++++++++++++++ 3 files changed, 86 insertions(+), 1 deletion(-) diff --git a/docs/api/optimizers/tree.md b/docs/api/optimizers/tree.md index 0f6b656..547c9a2 100644 --- a/docs/api/optimizers/tree.md +++ b/docs/api/optimizers/tree.md @@ -22,6 +22,7 @@ from pepsy.optimizers.tree import TreeLayoutFinder, TreeOptimizer finder = TreeLayoutFinder(gates, n=5, root_qubit=4, max_arity=2) plan = finder.run( + order="quality", refine="greedy", refine_budget=64, search="nevergrad", @@ -397,6 +398,12 @@ it combines normalized path score, maximum edge load, and total edge load with `weight_mode` / `layout_weight_mode` option accepts `count`, `auto`, `angle`, or `operator_schmidt` for interaction-graph weighting. +Use `order="quality"` with `finder.run()` (or set it on the finder) for the +MPS-style higher-quality offline search. It enables bounded greedy leaf +refinement and opportunistic Nevergrad refinement when Nevergrad is installed; +otherwise it falls back to greedy refinement. The zero-argument `run()` path +remains the fast deterministic candidate selection. + For compression-first selection, use `objective="compression"` (or `layout_objective="compression"`). It prioritizes peak and total predicted operator-Schmidt load, then penalizes the estimated local tensor size at the @@ -552,6 +559,7 @@ layout API: ```python tree_plan = finder.run( + order="quality", refine="greedy", refine_budget=64, search="nevergrad", diff --git a/src/pepsy/optimizers/tree/layout.py b/src/pepsy/optimizers/tree/layout.py index f0841d7..ce0d414 100644 --- a/src/pepsy/optimizers/tree/layout.py +++ b/src/pepsy/optimizers/tree/layout.py @@ -57,6 +57,7 @@ _DEFAULT_MAX_ARITY = object() _DEFAULT_CHI = object() +_DEFAULT_ORDER = object() _DEFAULT_SEARCH_OPTION = object() _DEFAULT_SCALE_MARKERS = ("o",) @@ -147,6 +148,31 @@ def _normalize_layout_search(search): return name +def _normalize_layout_order(order): + """Normalize the optional high-quality layout mode.""" + if order is None: + return None + name = str(order).replace("-", "_").strip().lower() + aliases = { + "auto": "quality", + "best": "quality", + "best_quality": "quality", + } + name = aliases.get(name, name) + if name != "quality": + raise ValueError("order must be None or 'quality'.") + return name + + +def _nevergrad_available(): + """Return whether the optional Nevergrad dependency can be imported.""" + try: + import nevergrad # pylint: disable=import-outside-toplevel,unused-import + except ImportError: + return False + return True + + def _validate_search_budget(value, name): """Validate a positive bounded layout-search evaluation budget.""" try: @@ -1113,6 +1139,10 @@ class TreeLayoutFinder: predicted operator-Schmidt load on tree edges. `"hybrid"` combines normalized path, peak-edge-load, and total-edge-load costs using ``hybrid_weights``. + order : {None, "quality"}, optional + Optional high-quality offline mode. `"quality"` enables bounded + greedy refinement and opportunistic Nevergrad refinement; omitted + keeps the fast deterministic candidate selection. hybrid_weights : mapping or sequence of three floats, optional Weights for the hybrid path, maximum edge load, and total edge load. The default is ``(1.0, 1.0, 0.25)``. @@ -1142,7 +1172,7 @@ def __init__(self, gates=None, n=None, *, supports=None, structure="quality", dense_max=512, objective="path", weight_mode="count", chi=None, max_operator_qubits=8, hybrid_weights=None, refine=None, refine_budget=None, search=None, search_budget=128, seed=0, - nevergrad_optimizer="OnePlusOne", root_qubit=None): + nevergrad_optimizer="OnePlusOne", order=None, root_qubit=None): if ( _looks_like_tree_tensor_network(gates) or _looks_like_tree_tensor_network(supports) @@ -1224,6 +1254,7 @@ def __init__(self, gates=None, n=None, *, supports=None, structure="quality", self.objective = _normalize_layout_objective(objective) self.hybrid_weights = _normalize_hybrid_weights(hybrid_weights) self.weight_mode = _normalize_weight_mode(weight_mode) + self.order = _normalize_layout_order(order) self.refine = _normalize_layout_refinement(refine) if refine_budget is not None: refine_budget = _validate_search_budget(refine_budget, "refine_budget") @@ -2135,6 +2166,7 @@ def candidate_key(candidate): def run( self, *, + order=_DEFAULT_ORDER, chi=_DEFAULT_CHI, refine=_DEFAULT_SEARCH_OPTION, refine_budget=_DEFAULT_SEARCH_OPTION, @@ -2156,7 +2188,22 @@ def run( overridden for this call. Pass ``progbar=True`` to display greedy and Nevergrad search progress. Omitted values inherit the corresponding finder settings, so the original zero-argument behavior is unchanged. + + ``order="quality"`` is a convenience mode matching the MPS layout + API: it enables bounded greedy refinement and opportunistic Nevergrad + refinement when the optional dependency is installed. If Nevergrad is + unavailable, quality mode falls back to greedy refinement. Pass + ``search=None`` or ``refine=None`` explicitly to disable either stage. """ + if order is _DEFAULT_ORDER: + order = self.order + else: + order = _normalize_layout_order(order) + if order == "quality": + if refine is _DEFAULT_SEARCH_OPTION: + refine = "greedy" + if search is _DEFAULT_SEARCH_OPTION: + search = "nevergrad" if _nevergrad_available() else None if chi is _DEFAULT_CHI: chi = self.chi else: diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index 86f43ea..14eaca1 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -1767,6 +1767,36 @@ def test_tree_layout_finder_plot_rubberband_is_axis_free_and_unlabeled(): plt.close(fig) +def test_tree_layout_quality_order_enables_bounded_refinement(monkeypatch): + """Tree order='quality' mirrors the MPS high-quality mode.""" + monkeypatch.setitem(sys.modules, "nevergrad", None) + finder = TreeLayoutFinder( + [(pepsy.cnot(), (0, 3)), (pepsy.cnot(), (1, 2))], + n=4, + max_arity=2, + order="quality", + ) + captured = {} + + def fake_improve(plan, *, chi, settings, progbar=False): + captured.update(settings) + return plan, {"method": "test"} + + monkeypatch.setattr(finder, "_improve_plan", fake_improve) + plan = finder.run() + + assert plan.n == 4 + assert captured["refine"] == "greedy" + assert captured["search"] is None + + +def test_tree_layout_order_rejects_non_quality_modes(): + """Tree layouts expose quality mode rather than 1-D order names.""" + finder = TreeLayoutFinder([], n=4, max_arity=2) + with pytest.raises(ValueError, match="order"): + finder.run(order="input") + + def test_tree_layout_rubberband_defaults_to_cotengra_ordered_colors(): """Default rubberbands use distinct post-order Spectral colors.""" matplotlib = pytest.importorskip("matplotlib") From 858a11773acd55ad8de407f79229a61381d98aa1 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Thu, 30 Jul 2026 17:49:28 -0700 Subject: [PATCH 35/70] tree layout: add direct hypergraph search --- docs/api/optimizers/tree.md | 48 ++- src/pepsy/optimizers/tree/layout.py | 413 +++++++++++++++++++++++-- src/pepsy/optimizers/tree/optimizer.py | 72 ++++- tests/test_optimize_tree.py | 135 ++++++++ 4 files changed, 625 insertions(+), 43 deletions(-) diff --git a/docs/api/optimizers/tree.md b/docs/api/optimizers/tree.md index 547c9a2..7820fac 100644 --- a/docs/api/optimizers/tree.md +++ b/docs/api/optimizers/tree.md @@ -398,11 +398,35 @@ it combines normalized path score, maximum edge load, and total edge load with `weight_mode` / `layout_weight_mode` option accepts `count`, `auto`, `angle`, or `operator_schmidt` for interaction-graph weighting. +For a genuinely multi-site layout objective, use +`objective="hypergraph"` (or `layout_objective="hypergraph"`). Each original +gate support is kept as one hyperedge, and the finder scores its actual +operator-Schmidt load on every crossed tree edge rather than selecting from a +pairwise proxy alone. This mode starts from inexpensive pairwise-derived seed +trees, then automatically performs bounded direct greedy leaf swaps and binary +NNI topology moves using the full hyperedge score. Pass +`refine=None, topology_refine=None` to inspect the unrefined direct score, or +set explicit budgets for a larger search. Dense operators wider than +`max_operator_qubits` still use the documented conservative rank bound. + Use `order="quality"` with `finder.run()` (or set it on the finder) for the MPS-style higher-quality offline search. It enables bounded greedy leaf -refinement and opportunistic Nevergrad refinement when Nevergrad is installed; -otherwise it falls back to greedy refinement. The zero-argument `run()` path -remains the fast deterministic candidate selection. +refinement, bounded binary-tree nearest-neighbor-interchange (NNI) topology +refinement, and opportunistic Nevergrad refinement when Nevergrad is installed; +otherwise it uses the deterministic NNI and leaf stages. NNI changes the +internal grouping itself, so quality mode can improve a tree even when the +best leaf labels are already fixed. The zero-argument `run()` path remains the +fast deterministic candidate selection. Disable the topology stage explicitly +with `topology_refine=None`, or bound it with `topology_budget=`. + +For a stream whose locality changes over time, pass `time_decay=` and/or +`time_window=` to `TreeLayoutFinder`. A decay in `(0, 1]` weights an event by +`time_decay ** age` (the newest event has age zero), while a window keeps only +the final events. The same factors are used for interaction paths, congestion +candidate construction, and per-edge operator-Schmidt load estimates, so the +diagnostics and selected plan use one consistent time model. The defaults are +unchanged. `TreeOptimizer` exposes these as `layout_time_decay=` and +`layout_time_window=`. For compression-first selection, use `objective="compression"` (or `layout_objective="compression"`). It prioritizes peak and total predicted @@ -532,8 +556,10 @@ choice = finder.recommend_layered( tree_plan = choice["plan"] ``` -`refine="greedy"` is deterministic and bounded; it is opt-in so existing -fast/default layout construction remains unchanged. A balanced TTN turns a +`refine="greedy"` is deterministic and bounded; it is opt-in for the existing +fast/default objectives. The explicit `objective="hypergraph"` mode enables +greedy and NNI refinement by default because otherwise its full-support score +would only rank a few pairwise-derived seed trees. A balanced TTN turns a well-aligned physical span `r` into a path with `O(log r)` tree hops, so the hybrid score uses path length as a replay-cost proxy while edge loads estimate the accuracy/bond-dimension cost. @@ -560,6 +586,8 @@ layout API: ```python tree_plan = finder.run( order="quality", + topology_refine="nni", + topology_budget=64, refine="greedy", refine_budget=64, search="nevergrad", @@ -609,9 +637,13 @@ opt = py.TreeOptimizer(gate_stream, tree=choice["plan"], chi=chi) The pilot replays candidates on independent copies with the real tree update kernels and returns measured infidelity, final bond, truncation count, and -runtime under `choice["pilot"]`. The original optimizer is unchanged unless -`install=True` is passed. Installation is restricted to product initial states; -an entangled TTN cannot generally be relaid out exactly. +runtime under `choice["pilot"]`. By default one bounded `order="quality"` +candidate (greedy leaf refinement plus binary NNI topology refinement) is +reserved a pilot slot, so it cannot be rejected before state-aware replay. Use +`include_quality=False` for the static-only candidate set. The original +optimizer is unchanged unless `install=True` is passed. Installation is +restricted to product initial states; an entangled TTN cannot generally be +relaid out exactly. Both helpers are also available from the package-level API: diff --git a/src/pepsy/optimizers/tree/layout.py b/src/pepsy/optimizers/tree/layout.py index ce0d414..7b8b7fa 100644 --- a/src/pepsy/optimizers/tree/layout.py +++ b/src/pepsy/optimizers/tree/layout.py @@ -136,6 +136,23 @@ def _normalize_layout_refinement(refine): return name +def _normalize_topology_refinement(refine): + """Normalize the optional joint topology refinement mode.""" + if refine is None or refine is False: + return None + name = str(refine).replace("-", "_").strip().lower() + aliases = { + "topology": "nni", + "joint": "nni", + "joint_greedy": "nni", + "greedy_topology": "nni", + } + name = aliases.get(name, name) + if name != "nni": + raise ValueError("topology_refine must be None or 'nni'.") + return name + + def _normalize_layout_search(search): """Normalize an optional offline fixed-plan search mode.""" if search is None or search is False: @@ -184,6 +201,49 @@ def _validate_search_budget(value, name): return value +def _normalize_time_decay(value): + """Validate an optional newest-event temporal decay factor.""" + if value is None: + return None + try: + value = float(value) + except (TypeError, ValueError) as exc: + raise ValueError("time_decay must be in (0, 1] or None.") from exc + if not np.isfinite(value) or value <= 0.0 or value > 1.0: + raise ValueError("time_decay must be in (0, 1] or None.") + return value + + +def _normalize_time_window(value): + """Validate an optional trailing event window.""" + if value is None: + return None + if isinstance(value, bool): + raise ValueError("time_window must be a positive integer or None.") + try: + value = int(value) + except (TypeError, ValueError) as exc: + raise ValueError( + "time_window must be a positive integer or None." + ) from exc + if value < 1: + raise ValueError("time_window must be a positive integer or None.") + return value + + +def _temporal_event_factors(num_events, *, time_decay=None, time_window=None): + """Return one newest-event-normalized factor for each stream event.""" + if num_events < 1: + return () + factors = np.ones(int(num_events), dtype=float) + if time_window is not None and time_window < num_events: + factors[: num_events - time_window] = 0.0 + if time_decay is not None and time_decay != 1.0: + ages = np.arange(num_events - 1, -1, -1, dtype=float) + factors *= np.power(time_decay, ages) + return tuple(float(factor) for factor in factors) + + def _safe_exp2(value): """Return ``2**value`` without emitting overflow warnings.""" if value > np.log2(np.finfo(float).max): @@ -205,12 +265,20 @@ def _normalize_layout_objective(objective): "compress": "compression", "accuracy": "compression", "bond_growth": "compression", + "hyperedge": "hypergraph", + "hyperedges": "hypergraph", + "hypergraph_load": "hypergraph", + "per_edge": "hypergraph", + "per_edge_load": "hypergraph", } name = aliases.get(name, name) - if name not in {"path", "congestion", "hybrid", "compression"}: + if name not in { + "path", "congestion", "hybrid", "compression", "hypergraph" + }: raise ValueError( f"Unknown tree layout objective {objective!r}. " - "Expected 'path', 'congestion', 'compression', or 'hybrid'." + "Expected 'path', 'congestion', 'compression', 'hypergraph', " + "or 'hybrid'." ) return name @@ -1133,12 +1201,16 @@ class TreeLayoutFinder: (see :meth:`TreePlan.from_order`). dense_max : int Maximum subsystem size for dense spectral reordering. - objective : {"path", "congestion", "compression", "hybrid"} + objective : {"path", "congestion", "compression", "hypergraph", "hybrid"} Layout objective. `"path"` preserves the co-occurrence/path-length heuristic; `"congestion"` selects among layout candidates using the predicted operator-Schmidt load on tree edges. `"hybrid"` combines normalized path, peak-edge-load, and total-edge-load costs using ``hybrid_weights``. + `"compression"` adds a local tensor-size proxy to the edge-load + objective. `"hypergraph"` is the direct multi-site mode: it ranks + plans from the full support hyperedges and per-edge Schmidt loads, + then applies bounded leaf and binary-topology refinement by default. order : {None, "quality"}, optional Optional high-quality offline mode. `"quality"` enables bounded greedy refinement and opportunistic Nevergrad refinement; omitted @@ -1150,9 +1222,17 @@ class TreeLayoutFinder: Optional fixed-plan local search used by :meth:`run` and recommendation methods. `"greedy"` tries adjacent leaf-label swaps before simulation; it never changes a live :class:`TreeOptimizer` tree. + topology_refine : {None, "nni"} + Optional joint topology refinement for binary candidates. `"nni"` + tries bounded nearest-neighbor interchange moves on internal edges, + retaining only objective-improving trees. It never changes a live + :class:`TreeOptimizer` tree. refine_budget : int, optional Maximum greedy swap proposals per candidate plan. Defaults to at most 64 proposals when refinement is enabled. + topology_budget : int, optional + Maximum NNI proposals per candidate plan. Defaults to at most 64 + proposals when topology refinement is enabled. search : {None, "nevergrad"} Optional offline derivative-free refinement. It is never run unless requested and requires the optional ``nevergrad`` package. @@ -1165,14 +1245,22 @@ class TreeLayoutFinder: weight_mode : {"count", "auto", "angle", "operator_schmidt"} Event weighting used for the interaction graph. `"count"` is the backward-compatible default. + time_decay : float, optional + If supplied, multiply an event's weight by ``time_decay ** age`` where + the newest event has age zero. Values must be in ``(0, 1]``. + time_window : int, optional + If supplied, only the final ``time_window`` stream events contribute to + layout scoring and predicted edge load. """ def __init__(self, gates=None, n=None, *, supports=None, structure="quality", max_arity=(2, 3, 4), community_frac=0.35, star_frac=0.75, dense_max=512, objective="path", weight_mode="count", chi=None, max_operator_qubits=8, hybrid_weights=None, refine=None, - refine_budget=None, search=None, search_budget=128, seed=0, - nevergrad_optimizer="OnePlusOne", order=None, root_qubit=None): + refine_budget=None, topology_refine=None, topology_budget=None, + search=None, search_budget=128, seed=0, + nevergrad_optimizer="OnePlusOne", order=None, root_qubit=None, + time_decay=None, time_window=None): if ( _looks_like_tree_tensor_network(gates) or _looks_like_tree_tensor_network(supports) @@ -1259,6 +1347,12 @@ def __init__(self, gates=None, n=None, *, supports=None, structure="quality", if refine_budget is not None: refine_budget = _validate_search_budget(refine_budget, "refine_budget") self.refine_budget = refine_budget + self.topology_refine = _normalize_topology_refinement(topology_refine) + if topology_budget is not None: + topology_budget = _validate_search_budget( + topology_budget, "topology_budget" + ) + self.topology_budget = topology_budget self.search = _normalize_layout_search(search) self.search_budget = _validate_search_budget(search_budget, "search_budget") try: @@ -1280,6 +1374,8 @@ def __init__(self, gates=None, n=None, *, supports=None, structure="quality", "max_operator_qubits must be a positive integer or None." ) self.max_operator_qubits = max_operator_qubits + self.time_decay = _normalize_time_decay(time_decay) + self.time_window = _normalize_time_window(time_window) # Layout search asks for the same structural quantities several times # (once per candidate arity and once per diagnostic). Keep these @@ -1294,7 +1390,7 @@ def __init__(self, gates=None, n=None, *, supports=None, structure="quality", self._balanced_plan_cache = None sites = list(range(self.n)) - self.event_weights = tuple( + base_event_weights = tuple( _gate_stream_event_weights( self.payloads, self.supports, @@ -1302,11 +1398,18 @@ def __init__(self, gates=None, n=None, *, supports=None, structure="quality", weight_mode=self.weight_mode, ) ) + self.temporal_factors = _temporal_event_factors( + len(base_event_weights), + time_decay=self.time_decay, + time_window=self.time_window, + ) self.event_weights = tuple( 0.0 if str(event_type).lower() in { "measure", "reset", "measure_reset", "cap" - } else weight - for weight, event_type in zip(self.event_weights, self.event_types) + } else weight * temporal_factor + for weight, temporal_factor, event_type in zip( + base_event_weights, self.temporal_factors, self.event_types + ) ) self.pair_weights = _gate_stream_pair_weights( supports, sites, self.event_weights @@ -1374,6 +1477,8 @@ def _resolve_search_settings( *, refine=_DEFAULT_SEARCH_OPTION, refine_budget=_DEFAULT_SEARCH_OPTION, + topology_refine=_DEFAULT_SEARCH_OPTION, + topology_budget=_DEFAULT_SEARCH_OPTION, search=_DEFAULT_SEARCH_OPTION, search_budget=_DEFAULT_SEARCH_OPTION, seed=_DEFAULT_SEARCH_OPTION, @@ -1382,6 +1487,12 @@ def _resolve_search_settings( """Resolve method overrides against finder-owned search defaults.""" if refine is _DEFAULT_SEARCH_OPTION: refine = self.refine + if self.objective == "hypergraph" and refine is None: + # A direct hypergraph score is only useful as a layout search + # objective when the candidate is allowed to move. Keep the + # old objectives fast, but make the explicitly requested + # hypergraph mode perform its bounded local search by default. + refine = "greedy" else: refine = _normalize_layout_refinement(refine) if refine_budget is _DEFAULT_SEARCH_OPTION: @@ -1393,6 +1504,21 @@ def _resolve_search_settings( if refine is not None and refine_budget is None: refine_budget = max(1, min(len(self.leaf_qubits) - 1, 64)) + if topology_refine is _DEFAULT_SEARCH_OPTION: + topology_refine = self.topology_refine + if self.objective == "hypergraph" and topology_refine is None: + topology_refine = "nni" + else: + topology_refine = _normalize_topology_refinement(topology_refine) + if topology_budget is _DEFAULT_SEARCH_OPTION: + topology_budget = self.topology_budget + elif topology_budget is not None: + topology_budget = _validate_search_budget( + topology_budget, "topology_budget" + ) + if topology_refine is not None and topology_budget is None: + topology_budget = max(1, min(max(1, len(self.leaf_qubits) - 2), 64)) + if search is _DEFAULT_SEARCH_OPTION: search = self.search else: @@ -1417,6 +1543,8 @@ def _resolve_search_settings( return { "refine": refine, "refine_budget": refine_budget, + "topology_refine": topology_refine, + "topology_budget": topology_budget, "search": search, "search_budget": search_budget, "seed": seed, @@ -1465,6 +1593,52 @@ def _plan_with_leaf_swap(self, plan, left_leaf, right_leaf): root_qubit=plan.root_qubit, ) + def _plan_with_nni(self, plan, parent, child, variant): + """Return one binary nearest-neighbor interchange of ``parent-child``. + + The rooted local pattern is ``parent -> (child, sibling)`` and + ``child -> (a, b)``. Each NNI variant keeps ``child`` below ``parent`` + while moving either ``a`` or ``b`` across the internal edge. Node ids + and leaf labels are retained so this is a topology move, not a hidden + relabeling. + """ + parent_children = tuple(plan.children[parent]) + child_children = tuple(plan.children[child]) + if ( + len(parent_children) != 2 + or len(child_children) != 2 + or child not in parent_children + ): + raise ValueError("NNI requires a binary internal parent-child edge.") + sibling = next(node for node in parent_children if node != child) + if variant not in (0, 1): + raise ValueError("NNI variant must be 0 or 1.") + a, b = child_children + moved = b if variant == 0 else a + retained = a if variant == 0 else b + children = { + node: tuple(child_ids) for node, child_ids in plan.children.items() + } + children[parent] = (child, moved) + children[child] = (retained, sibling) + return TreePlan.from_children( + children, + plan.qubit_of_leaf, + root=plan.root, + root_qubit=plan.root_qubit, + ) + + @staticmethod + def _nni_edges(plan): + """Return deterministic binary internal edges eligible for NNI.""" + return tuple( + (parent, child) + for parent, children in sorted(plan.children.items()) + if len(children) == 2 + for child in children + if len(plan.children.get(child, ())) == 2 + ) + def _path_score_and_max(self, plan): """Return the weighted interaction path sum and longest active path.""" score = 0.0 @@ -1569,7 +1743,7 @@ def _objective_key(self, plan): return self._path_score_and_max(plan) if self.objective == "congestion": return self._congestion_key(plan) - if self.objective == "compression": + if self.objective in {"compression", "hypergraph"}: loads = self.edge_loads(plan) values = tuple(loads.values()) tensor_cost = self._tensor_cost_key(plan) @@ -1599,7 +1773,7 @@ def _selection_loss(self, plan, chi): key = self._objective_key(plan) if self.objective == "path": value = key[0] - elif self.objective in {"congestion", "compression"}: + elif self.objective in {"congestion", "compression", "hypergraph"}: value = key[0] + 1.0e-6 * key[1] + 1.0e-12 * key[2] else: value = key[0] @@ -1691,6 +1865,71 @@ def _refine_plan_greedy(self, plan, *, chi, budget, progbar=False): "final_key": current_key, } + def _refine_plan_topology(self, plan, *, chi, budget, progbar=False): + """Greedily improve binary topology through bounded NNI moves.""" + initial_key = self._selection_key(plan, chi) + if budget < 1 or not plan.is_binary(): + return plan, { + "method": "nni", + "evaluations": 0, + "accepted_moves": 0, + "initial_key": initial_key, + "final_key": initial_key, + } + + current = plan + current_key = initial_key + evaluations = 0 + accepted_moves = 0 + progress = None + if progbar: + from tqdm import tqdm # pylint: disable=import-outside-toplevel + + progress = tqdm( + total=budget, + desc="tree layout topology", + leave=False, + ) + + while evaluations < budget: + best = current + best_key = current_key + for parent, child in self._nni_edges(current): + for variant in (0, 1): + if evaluations >= budget: + break + evaluations += 1 + if progress is not None: + progress.update() + candidate = self._plan_with_nni( + current, parent, child, variant + ) + candidate_key = self._selection_key(candidate, chi) + if candidate_key < best_key: + if best is not current: + self._discard_plan_cache(best) + best = candidate + best_key = candidate_key + else: + self._discard_plan_cache(candidate) + if evaluations >= budget: + break + if best is current: + break + current = best + current_key = best_key + accepted_moves += 1 + + if progress is not None: + progress.close() + return current, { + "method": "nni", + "evaluations": evaluations, + "accepted_moves": accepted_moves, + "initial_key": initial_key, + "final_key": current_key, + } + def _refine_plan_nevergrad( self, plan, *, chi, budget, seed, optimizer_name, progbar=False ): @@ -1803,9 +2042,17 @@ def _improve_plan(self, plan, *, chi, settings, progbar=False): info = { "initial_order": initial_order, "initial_key": initial_key, + "topology_refinement": None, "refinement": None, "search": None, } + if settings["topology_refine"] == "nni": + plan, info["topology_refinement"] = self._refine_plan_topology( + plan, + chi=chi, + budget=settings["topology_budget"], + progbar=progbar, + ) if settings["refine"] == "greedy": plan, info["refinement"] = self._refine_plan_greedy( plan, @@ -1912,7 +2159,14 @@ def _schmidt_rank_info(self, payload, support, left_support): return info def _candidate_plans(self, max_arity): - """Build the candidate plans considered by the selected objective.""" + """Build deterministic seed plans for the selected objective. + + The ``hypergraph`` objective deliberately keeps several inexpensive + pairwise-derived seeds, but all final ranking and its default local + refinement use :meth:`edge_loads`, which scans each original + multi-site support across the candidate tree. The pairwise seeds are + therefore only an initialization strategy, not the objective itself. + """ interaction_plan = self._build_plan(self._similarity_weights()) if max_arity != self.max_arity: interaction_plan = self._build_plan( @@ -2015,6 +2269,8 @@ def recommend_layered( chi=_DEFAULT_CHI, refine=_DEFAULT_SEARCH_OPTION, refine_budget=_DEFAULT_SEARCH_OPTION, + topology_refine=_DEFAULT_SEARCH_OPTION, + topology_budget=_DEFAULT_SEARCH_OPTION, search=_DEFAULT_SEARCH_OPTION, search_budget=_DEFAULT_SEARCH_OPTION, seed=_DEFAULT_SEARCH_OPTION, @@ -2053,9 +2309,16 @@ def recommend_layered( refine : {None, "greedy"}, optional Override the finder refinement setting. `"greedy"` performs a bounded adjacent leaf-swap search on each candidate tree. + topology_refine : {None, "nni"}, optional + Override the optional binary-tree topology refinement. `"nni"` + performs bounded nearest-neighbor interchange proposals. It is a + no-op for the non-binary layered structure. refine_budget : int, optional Maximum greedy proposals per candidate. When omitted, an enabled greedy search uses at most ``min(n - 1, 64)`` proposals. + topology_budget : int, optional + Maximum NNI proposals per candidate. When omitted, an enabled NNI + search uses at most 64 proposals. search : {None, "nevergrad"}, optional Override the finder offline search setting. Nevergrad optimizes only the returned fixed plan; it never mutates a live TTN. @@ -2079,6 +2342,8 @@ def recommend_layered( settings = self._resolve_search_settings( refine=refine, refine_budget=refine_budget, + topology_refine=topology_refine, + topology_budget=topology_budget, search=search, search_budget=search_budget, seed=seed, @@ -2158,6 +2423,7 @@ def candidate_key(candidate): "order": recommended["order"], "chi": chi, "refine": settings["refine"], + "topology_refine": settings["topology_refine"], "search": settings["search"], "plan": recommended["plan"], "candidates": candidates, @@ -2170,6 +2436,8 @@ def run( chi=_DEFAULT_CHI, refine=_DEFAULT_SEARCH_OPTION, refine_budget=_DEFAULT_SEARCH_OPTION, + topology_refine=_DEFAULT_SEARCH_OPTION, + topology_budget=_DEFAULT_SEARCH_OPTION, search=_DEFAULT_SEARCH_OPTION, search_budget=_DEFAULT_SEARCH_OPTION, seed=_DEFAULT_SEARCH_OPTION, @@ -2188,6 +2456,10 @@ def run( overridden for this call. Pass ``progbar=True`` to display greedy and Nevergrad search progress. Omitted values inherit the corresponding finder settings, so the original zero-argument behavior is unchanged. + The explicit ``objective="hypergraph"`` mode is the one exception: + when no refinement controls are supplied, it enables bounded greedy + and binary-NNI stages so the full support hyperedges directly + influence the returned layout. ``order="quality"`` is a convenience mode matching the MPS layout API: it enables bounded greedy refinement and opportunistic Nevergrad @@ -2202,6 +2474,8 @@ def run( if order == "quality": if refine is _DEFAULT_SEARCH_OPTION: refine = "greedy" + if topology_refine is _DEFAULT_SEARCH_OPTION: + topology_refine = "nni" if search is _DEFAULT_SEARCH_OPTION: search = "nevergrad" if _nevergrad_available() else None if chi is _DEFAULT_CHI: @@ -2211,6 +2485,8 @@ def run( settings = self._resolve_search_settings( refine=refine, refine_budget=refine_budget, + topology_refine=topology_refine, + topology_budget=topology_budget, search=search, search_budget=search_budget, seed=seed, @@ -2233,7 +2509,11 @@ def run( } return rec["plan"] candidates = self._candidate_plans(self.max_arity) - if settings["refine"] is not None or settings["search"] is not None: + if ( + settings["topology_refine"] is not None + or settings["refine"] is not None + or settings["search"] is not None + ): candidates = { name: self._improve_plan( plan, @@ -2255,7 +2535,14 @@ def run( self._selected_candidate = selected return candidates[selected] - def candidate_plans(self, *, chi=_DEFAULT_CHI): + def candidate_plans( + self, + *, + chi=_DEFAULT_CHI, + include_quality=False, + quality_refine_budget=None, + quality_topology_budget=None, + ): """Return immutable candidate plans for optional pilot replay. The normal :meth:`run` path remains static and cheap. This method @@ -2263,6 +2550,20 @@ def candidate_plans(self, *, chi=_DEFAULT_CHI): that a state-aware pilot can compare without rebuilding the finder. Candidate names are stable strings such as ``"congestion:arity=2"``. + + Parameters + ---------- + chi : int, optional + Bond-dimension budget used in candidate ranking. + include_quality : bool, optional + Also add one ``"quality:arity=..."`` candidate per arity. These + candidates start from the static objective candidates and apply + bounded greedy leaf and binary NNI topology refinement. This is + deliberately opt-in because it is more expensive than the static + candidate list. + quality_refine_budget, quality_topology_budget : int, optional + Bounds for the quality candidate's leaf-swap and NNI proposals. + Each defaults to the normal bounded quality-mode budget. """ if chi is _DEFAULT_CHI: chi = self.chi @@ -2274,6 +2575,15 @@ def candidate_plans(self, *, chi=_DEFAULT_CHI): else (self.max_arity,) ) result = {} + quality_settings = None + if include_quality: + quality_settings = self._resolve_search_settings( + refine="greedy", + refine_budget=quality_refine_budget, + topology_refine="nni", + topology_budget=quality_topology_budget, + search=None, + ) for arity in arities: plans = self._candidate_plans(arity) for name, plan in plans.items(): @@ -2285,6 +2595,33 @@ def candidate_plans(self, *, chi=_DEFAULT_CHI): "tensor_cost": self._tensor_cost_key(plan), "edge_loads": self.edge_loads(plan), } + if quality_settings is not None: + refined_candidates = [] + for name, plan in plans.items(): + refined, planning = self._improve_plan( + plan, + chi=chi, + settings=quality_settings, + ) + refined_candidates.append(( + self._selection_key(refined, chi), + name, + refined, + planning, + )) + _key, source_name, quality_plan, planning = min( + refined_candidates, + key=lambda item: (item[0], item[1]), + ) + result[f"quality:arity={arity}"] = { + "plan": quality_plan, + "objective_key": self._selection_key(quality_plan, chi), + "path_score": self.score(quality_plan), + "tensor_cost": self._tensor_cost_key(quality_plan), + "edge_loads": self.edge_loads(quality_plan), + "selected_from": source_name, + "planning": planning, + } return result def recommend_arities( @@ -2294,6 +2631,8 @@ def recommend_arities( chi=_DEFAULT_CHI, refine=_DEFAULT_SEARCH_OPTION, refine_budget=_DEFAULT_SEARCH_OPTION, + topology_refine=_DEFAULT_SEARCH_OPTION, + topology_budget=_DEFAULT_SEARCH_OPTION, search=_DEFAULT_SEARCH_OPTION, search_budget=_DEFAULT_SEARCH_OPTION, seed=_DEFAULT_SEARCH_OPTION, @@ -2321,7 +2660,8 @@ def recommend_arities( ``max_bond_cut``, ``chi_overflow``, and ``exact_at_chi``. When omitted, uses the ``chi`` supplied to the finder; pass ``chi=None`` explicitly for a chi-blind comparison. - refine, refine_budget, search, search_budget, seed, nevergrad_optimizer + refine, refine_budget, topology_refine, topology_budget, search, + search_budget, seed, nevergrad_optimizer Optional fixed-plan search controls with the same meaning as in :meth:`recommend_layered`. They are applied to each arity candidate before selecting one final immutable plan. @@ -2335,6 +2675,8 @@ def recommend_arities( settings = self._resolve_search_settings( refine=refine, refine_budget=refine_budget, + topology_refine=topology_refine, + topology_budget=topology_budget, search=search, search_budget=search_budget, seed=seed, @@ -2417,6 +2759,7 @@ def candidate_key(candidate): "recommended_max_arity": recommended["max_arity"], "chi": chi, "refine": settings["refine"], + "topology_refine": settings["topology_refine"], "search": settings["search"], "plan": recommended["plan"], "candidates": candidates, @@ -2432,8 +2775,11 @@ def _congestion_pair_weights(self): if cached is not None: return cached event_weights = [] - for payload, support, event_type in zip( - self.payloads, self.supports, self.event_types + for payload, support, event_type, temporal_factor in zip( + self.payloads, + self.supports, + self.event_types, + self.temporal_factors, ): if len(support) < 2 or str(event_type).lower() in { "measure", "reset", "measure_reset", "cap" @@ -2441,7 +2787,7 @@ def _congestion_pair_weights(self): event_weights.append(0.0) continue if payload is None: - event_weights.append(1.0) + event_weights.append(float(temporal_factor)) continue support = tuple(dict.fromkeys(support)) logs = [] @@ -2457,7 +2803,7 @@ def _congestion_pair_weights(self): for site in support: rank = self._schmidt_rank(payload, support, (site,)) logs.append(float(np.log2(rank))) - event_weights.append(max(logs, default=1.0)) + event_weights.append(float(temporal_factor) * max(logs, default=1.0)) self._congestion_weights_cache = _gate_stream_pair_weights( self.supports, range(self.n), @@ -2493,13 +2839,20 @@ def edge_loads(self, plan=None): "bounded_events": 0, "reasons": {}, } - for payload, support, event_type in zip( - self.payloads, self.supports, self.event_types + for payload, support, event_type, temporal_factor in zip( + self.payloads, + self.supports, + self.event_types, + self.temporal_factors, ): support = tuple(dict.fromkeys(support)) - if len(support) < 2 or str(event_type).lower() in { - "measure", "reset", "measure_reset", "cap" - }: + if ( + temporal_factor <= 0.0 + or len(support) < 2 + or str(event_type).lower() in { + "measure", "reset", "measure_reset", "cap" + } + ): continue support_mask = 0 for site in support: @@ -2540,7 +2893,9 @@ def edge_loads(self, plan=None): ) info = self._schmidt_rank_info(payload, support, left) rank = int(info["rank"]) - loads[edge] += float(np.log2(max(1, rank))) + loads[edge] += float(temporal_factor) * float( + np.log2(max(1, rank)) + ) if info["exact"]: rank_diagnostics["exact_events"] += 1 else: @@ -2682,6 +3037,9 @@ def report(self, plan=None, *, include_edge_loads=True): "n_interacting_pairs": n_pairs, "objective": self.objective, "weight_mode": self.weight_mode, + "time_decay": self.time_decay, + "time_window": self.time_window, + "active_events": int(sum(factor > 0.0 for factor in self.temporal_factors)), "hybrid_weights": ( self.hybrid_weights if self.objective == "hybrid" else None ), @@ -2692,6 +3050,13 @@ def report(self, plan=None, *, include_edge_loads=True): float(objective_key[0] + objective_key[1]) if self.objective == "compression" else None ), + "hypergraph_score": ( + { + "max_edge_load": float(max_load), + "total_edge_load": float(total_load), + } + if self.objective == "hypergraph" and loads is not None else None + ), "root": plan.root, "root_qubit": plan.root_qubit, "is_binary": plan.is_binary(), diff --git a/src/pepsy/optimizers/tree/optimizer.py b/src/pepsy/optimizers/tree/optimizer.py index d652e48..4cff216 100644 --- a/src/pepsy/optimizers/tree/optimizer.py +++ b/src/pepsy/optimizers/tree/optimizer.py @@ -63,6 +63,8 @@ from .layout import ( TreeLayoutFinder, TreePlan, + _normalize_time_decay, + _normalize_time_window, _submpo_schmidt_rank_bound, ) from .ttn import TreeTensorNetwork @@ -324,17 +326,24 @@ class TreeOptimizer: optimizer's ``chi`` (structures that stay exact at ``chi`` are preferred). Pass ``max_arity=2`` to force a fixed binary tree. Ignored when an explicit ``tree`` is supplied. - layout_objective : {"path", "congestion", "compression", "hybrid"} + layout_objective : {"path", "congestion", "compression", "hypergraph", "hybrid"} Objective used when building an automatic tree. ``"path"`` is the backward-compatible interaction-path heuristic; ``"congestion"`` selects a candidate using predicted operator-Schmidt edge load; ``"compression"`` additionally penalizes peak/total load and the - estimated local tensor cost at ``chi``; + estimated local tensor cost at ``chi``; ``"hypergraph"`` directly + scores every original multi-qubit support across every crossed tree + edge and enables bounded leaf/NNI refinement by default; ``"hybrid"`` combines normalized path, peak-load, and total-load costs. Pass a configured :class:`TreeLayoutFinder` through ``layout=`` to customize its hybrid weights or enable pre-simulation refinement. layout_weight_mode : {"count", "auto", "angle", "operator_schmidt"} Event weighting used by the automatic layout interaction graph. + layout_time_decay : float, optional + Optional newest-event decay passed to :class:`TreeLayoutFinder`. + Values are in ``(0, 1]``; omitted means no temporal weighting. + layout_time_window : int, optional + Optional trailing gate-event window passed to the layout finder. layout : TreeLayoutFinder or TreePlan, optional A precomputed layout finder or its resulting plan. This is an alias layer over ``tree=`` and is useful when the finder also provides @@ -435,7 +444,8 @@ def __init__(self, gates=None, n=None, *, chi=64, two_site_mode=None, structure="quality", max_arity=(2, 3, 4), community_frac=0.35, star_frac=0.75, layout_objective="path", - layout_weight_mode="count", layout=None, tree=None, + layout_weight_mode="count", layout_time_decay=None, + layout_time_window=None, layout=None, tree=None, root_qubit=None, dtype=complex, threads=1, seed=None, run=True, tn=None, state=None, track_truncation=False, track_infidelity=True, @@ -584,6 +594,8 @@ def __init__(self, gates=None, n=None, *, chi=64, self.star_frac = float(star_frac) self.layout_objective = str(layout_objective) self.layout_weight_mode = str(layout_weight_mode) + self.layout_time_decay = _normalize_time_decay(layout_time_decay) + self.layout_time_window = _normalize_time_window(layout_time_window) self.dtype = dtype self.threads = None if threads is None else int(threads) if self.threads is not None and self.threads < 1: @@ -622,6 +634,8 @@ def __init__(self, gates=None, n=None, *, chi=64, star_frac=self.star_frac, objective=self.layout_objective, weight_mode=self.layout_weight_mode, + time_decay=self.layout_time_decay, + time_window=self.layout_time_window, chi=self.chi, max_operator_qubits=self.max_operator_qubits, root_qubit=root_qubit, @@ -1431,6 +1445,7 @@ def select_layout_for_compression( *, pilot_candidates=4, pilot_steps=None, + include_quality=True, install=False, progbar=False, ): @@ -1438,11 +1453,15 @@ def select_layout_for_compression( Static compression candidates are generated with ``objective="compression"`` and then replayed on independent copies - of the current state. The pilot uses the real tree update kernels, - ``chi``, cutoff, backend, and queued gate stream. The original state - is unchanged. By default the selected plan is returned for explicit - hand-off; ``install=True`` is allowed only for a product state and - remounts that state exactly on the selected geometry. + of the current state. When ``include_quality=True`` (the default), one + bounded greedy/NNI quality candidate is reserved a pilot slot so it + cannot be excluded by static surrogate ranking. The pilot uses the + real tree update kernels, ``chi``, cutoff, backend, and queued gate + stream. The original state is unchanged. By default the selected plan + is returned for explicit hand-off; ``install=True`` is allowed only + for a product state and remounts that state exactly on the selected + geometry. Pass ``include_quality=False`` for the previous static-only + candidate set. """ try: pilot_candidates = int(pilot_candidates) @@ -1467,15 +1486,34 @@ def select_layout_for_compression( star_frac=self.star_frac, objective="compression", weight_mode=self.layout_weight_mode, + time_decay=self.layout_time_decay, + time_window=self.layout_time_window, chi=self.chi, max_operator_qubits=self.max_operator_qubits, root_qubit=self.plan.root_qubit, ) - candidates = finder.candidate_plans(chi=self.chi) - ranked = sorted( + candidates = finder.candidate_plans( + chi=self.chi, + include_quality=bool(include_quality), + ) + ranked_static = sorted( candidates, key=lambda name: candidates[name]["objective_key"], - )[:pilot_candidates] + ) + if include_quality: + quality_names = [ + name for name in ranked_static if name.startswith("quality:") + ] + non_quality_names = [ + name for name in ranked_static if not name.startswith("quality:") + ] + reserved_quality = quality_names[:1] + ranked = ( + reserved_quality + + non_quality_names[: max(0, pilot_candidates - 1)] + ) + else: + ranked = ranked_static[:pilot_candidates] reports = {} successful = [] for name in ranked: @@ -1562,6 +1600,7 @@ def select_layout_for_compression( "candidates": candidates, "pilot": { "objective": "compression", + "include_quality": bool(include_quality), "pilot_candidates": tuple(ranked), "selected_candidate": selected_name, "reports": reports, @@ -1588,6 +1627,8 @@ def plot_layout(self, plan=None, *, layout_kwargs=None, **plot_kwargs): star_frac=self.star_frac, objective=self.layout_objective, weight_mode=self.layout_weight_mode, + time_decay=self.layout_time_decay, + time_window=self.layout_time_window, chi=self.chi, max_operator_qubits=self.max_operator_qubits, root_qubit=self.plan.root_qubit, @@ -1612,6 +1653,8 @@ def plot_rubberband(self, plan=None, *, layout_kwargs=None, **plot_kwargs): star_frac=self.star_frac, objective=self.layout_objective, weight_mode=self.layout_weight_mode, + time_decay=self.layout_time_decay, + time_window=self.layout_time_window, chi=self.chi, max_operator_qubits=self.max_operator_qubits, root_qubit=self.plan.root_qubit, @@ -1636,6 +1679,8 @@ def plot_tent(self, plan=None, *, layout_kwargs=None, **plot_kwargs): star_frac=self.star_frac, objective=self.layout_objective, weight_mode=self.layout_weight_mode, + time_decay=self.layout_time_decay, + time_window=self.layout_time_window, chi=self.chi, max_operator_qubits=self.max_operator_qubits, root_qubit=self.plan.root_qubit, @@ -4538,6 +4583,8 @@ def copy(self): threads=self.threads, layout_objective=self.layout_objective, layout_weight_mode=self.layout_weight_mode, + layout_time_decay=self.layout_time_decay, + layout_time_window=self.layout_time_window, track_truncation=self.track_truncation, track_infidelity=self.track_infidelity, max_intermediate_bond=self.max_intermediate_bond, @@ -4600,6 +4647,7 @@ def find_tree_layout(cls, gates, n=None, *, structure="quality", max_arity=(2, 3, 4), community_frac=0.35, star_frac=0.75, layout_objective="path", layout_weight_mode="count", + layout_time_decay=None, layout_time_window=None, root_qubit=None, max_operator_qubits=_DEFAULT_MAX_OPERATOR_QUBITS): """Return the :class:`TreePlan` a :class:`TreeLayoutFinder` would use.""" @@ -4608,6 +4656,8 @@ def find_tree_layout(cls, gates, n=None, *, structure="quality", max_arity=max_arity, community_frac=community_frac, star_frac=star_frac, objective=layout_objective, weight_mode=layout_weight_mode, + time_decay=layout_time_decay, + time_window=layout_time_window, root_qubit=root_qubit, max_operator_qubits=max_operator_qubits, ).run() diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index 14eaca1..c0f15fa 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -995,6 +995,48 @@ def test_compression_layout_reports_rank_bounds_and_tensor_cost(): assert len(report["objective_key"]) >= 5 +def test_hypergraph_layout_scores_full_multisite_supports(): + """Direct mode ranks original hyperedges on every crossed tree cut.""" + rng = np.random.default_rng(104) + gate = _rand_unitary(3, rng) + supports = ((0, 1, 2), (2, 3, 4)) + finder = TreeLayoutFinder( + [(gate, supports[0]), (gate, supports[1])], + n=5, + max_arity=2, + objective="hypergraph", + ) + plan = TreePlan.from_order(range(5), structure="balanced", max_arity=2) + + loads = finder.edge_loads(plan) + below = plan.subtree_qubit_masks() + expected = {edge: 0.0 for edge in loads} + for payload, support in zip(finder.payloads, finder.supports): + support_mask = sum(1 << q for q in support) + for edge in expected: + _parent, child = edge + left_mask = support_mask & below[child] + if not left_mask or left_mask == support_mask: + continue + left = tuple(q for q in support if left_mask & (1 << q)) + expected[edge] += np.log2(finder._schmidt_rank(payload, support, left)) + + assert loads == pytest.approx(expected) + report = finder.report(plan) + assert report["objective"] == "hypergraph" + assert report["hypergraph_score"] == { + "max_edge_load": max(loads.values()), + "total_edge_load": sum(loads.values()), + } + + recommendation = finder.recommend_arities((2,), chi=None) + assert recommendation["refine"] == "greedy" + assert recommendation["topology_refine"] == "nni" + assert recommendation["candidates"][0]["planning"][ + "topology_refinement" + ]["method"] == "nni" + + def test_tree_compression_layout_pilot_is_non_mutating(): """Tree pilot selection compares copied product states only.""" gates = [(pepsy.cnot(), (0, 3)), (pepsy.cnot(), (3, 1))] @@ -1008,10 +1050,40 @@ def test_tree_compression_layout_pilot_is_non_mutating(): assert selected["selected_candidate"] assert selected["pilot"]["reports"] + assert any( + name.startswith("quality:") + for name in selected["pilot"]["pilot_candidates"] + ) assert opt.plan is original_plan assert opt.max_bond() == 1 +def test_tree_candidate_plans_include_quality_for_state_aware_pilots(): + """Quality refinement is exposed as an explicit pilot candidate.""" + finder = TreeLayoutFinder( + [(pepsy.cnot(), (0, 3)), (pepsy.cnot(), (3, 1))], + n=4, + max_arity=2, + objective="compression", + ) + + candidates = finder.candidate_plans( + chi=2, + include_quality=True, + quality_refine_budget=2, + quality_topology_budget=2, + ) + + quality = candidates["quality:arity=2"] + assert quality["planning"]["topology_refinement"]["method"] == "nni" + assert quality["planning"]["refinement"]["method"] == "greedy" + assert quality["plan"].is_binary() + assert not any( + name.startswith("quality:") + for name in finder.candidate_plans(chi=2) + ) + + def test_tree_edge_loads_match_full_edge_reference(): """Steiner-only edge scanning preserves the full congestion calculation.""" rng = np.random.default_rng(109) @@ -1363,6 +1435,8 @@ def capture(max_arities, **kwargs): "chi": None, "refine": "greedy", "refine_budget": 2, + "topology_refine": None, + "topology_budget": None, "search": None, "search_budget": 7, "seed": 11, @@ -1787,9 +1861,70 @@ def fake_improve(plan, *, chi, settings, progbar=False): assert plan.n == 4 assert captured["refine"] == "greedy" + assert captured["topology_refine"] == "nni" assert captured["search"] is None +def test_tree_layout_nni_refinement_changes_binary_topology(): + """Quality refinement can move a correlated subtree, not only labels.""" + cnot = pepsy.cnot() + finder = TreeLayoutFinder( + [(cnot, (0, 2)), (cnot, (0, 3))], + n=4, + max_arity=2, + objective="path", + ) + initial = TreePlan.from_order(range(4), structure="balanced", max_arity=2) + + refined, planning = finder._refine_plan_topology( + initial, + chi=None, + budget=4, + ) + + assert refined.is_binary() + assert refined.children != initial.children + assert finder.score(refined) < finder.score(initial) + assert planning["accepted_moves"] >= 1 + + +def test_tree_layout_temporal_weights_apply_to_paths_and_edge_loads(): + """Recent-event weighting affects both locality and Schmidt-load scoring.""" + cnot = pepsy.cnot() + gates = [(cnot, (0, 1)), (cnot, (2, 3))] + full = TreeLayoutFinder(gates, n=4, max_arity=2, objective="congestion") + recent = TreeLayoutFinder( + gates, + n=4, + max_arity=2, + objective="congestion", + time_decay=0.5, + time_window=1, + ) + + assert full.temporal_factors == (1.0, 1.0) + assert recent.temporal_factors == (0.0, 1.0) + assert sum(recent.event_weights) == pytest.approx(1.0) + assert sum(recent.edge_loads(recent.run()).values()) < sum( + full.edge_loads(full.run()).values() + ) + report = recent.report() + assert report["time_decay"] == pytest.approx(0.5) + assert report["time_window"] == 1 + assert report["active_events"] == 1 + + opt = TreeOptimizer( + gates, + n=4, + max_arity=2, + layout_time_decay=0.5, + layout_time_window=1, + run=False, + ) + assert opt.layout_finder.time_window == 1 + assert opt.layout_finder.time_decay == pytest.approx(0.5) + + def test_tree_layout_order_rejects_non_quality_modes(): """Tree layouts expose quality mode rather than 1-D order names.""" finder = TreeLayoutFinder([], n=4, max_arity=2) From 2924cb7f9faaee3e97b916c7b4fe14920edfa56e Mon Sep 17 00:00:00 2001 From: rezaquant Date: Thu, 30 Jul 2026 19:17:10 -0600 Subject: [PATCH 36/70] Add fermionic open loop-series BP expansion --- .github/skills/belief-propagation/SKILL.md | 28 + docs/api/bp.md | 125 ++ .../references/belief_propagation.md | 2 + src/pepsy/bp/__init__.py | 16 + src/pepsy/bp/_symmray.py | 79 +- src/pepsy/bp/series.py | 1916 ++++++++++++++++- tests/test_bp_open_series.py | 464 ++++ tests/test_bp_symmray.py | 638 ++++++ 8 files changed, 3197 insertions(+), 71 deletions(-) create mode 100644 tests/test_bp_open_series.py diff --git a/.github/skills/belief-propagation/SKILL.md b/.github/skills/belief-propagation/SKILL.md index 081797d..5e3626a 100644 --- a/.github/skills/belief-propagation/SKILL.md +++ b/.github/skills/belief-propagation/SKILL.md @@ -148,6 +148,34 @@ not treat `optimize="auto-hq"` as a fermionic sign fix: it is a Cotengra path preset, just like a reusable `PathOptimizer`, and a path-dependent sign means the input metadata is invalid or incomplete. +The missing-`dummy_modes` defect and the cyclic open-series defect are related, +but they are not the same bug. The former is a metadata-repair failure: a +labelled odd array created with `new_with(...)` has lost the implicit modes +needed to preserve its global fermionic phase. The latter can occur even when +every local array is parity-even and has `dummy_modes=()`. In that case the +failure is the native Symmray representation of a cyclic open correction: the +unexcited `P` projectors and excited `Q` projectors can require incompatible +mixed bra/ket orientations, so relabelling an open `Q` to make pairwise +contractions run is not an algebraically safe fermionic contraction. It can +produce a non-convergent series even though BP has converged and the dense +shadow is correct. + +When a native fermionic open scalar or rho calculation has a cyclic graph, +use the graded loop-cluster-compatible route. It keeps the rho native, inserts +the gate in the graded ket/bra contraction for scalar observables, and avoids +using `trace(rho @ gate)` as an oracle. A useful diagnosis is to check +`parity`, `dummy_modes`, and `label` first: if all arrays are even with no +dummy modes, do not blame the missing-dummy repair. Compare against the native +exact oracle and the native loop-cluster result instead. Keep the explicit open +edge-series route for dense/ordinary cases and fermionic trees, where the +mixed-orientation cyclic obstruction is absent. If contraction budgets are +provided, apply them to the cluster contractions as well and inspect the +reported FLOP/peak-memory decisions; never use the budgeted call as a reason +to fall back to the unsafe mixed-orientation route. Use the route-independent +diagnostics `open_*_edge_term_costs` and +`open_*_cluster_region_costs` when consuming cost metadata; the older +`open_*_term_costs` fields are route-specific compatibility aliases. + ### Fermionic local observables: required construction and oracle Do **not** evaluate a native fermionic observable as `trace(rho @ gate)`, even diff --git a/docs/api/bp.md b/docs/api/bp.md index 95b9a2a..058c406 100644 --- a/docs/api/bp.md +++ b/docs/api/bp.md @@ -119,6 +119,131 @@ fermionic scalar correction is currently restricted to one-site gates; the multi-site graded-Q contraction is rejected explicitly while its block routing is completed. +For separated sites, use `partial_trace_open_loop_series_expand` when the +explicit configuration family should include Q paths between the retained +sites. Its integer `gloops` is a maximum number of excited virtual edges. A +configuration is retained when degree-one Q vertices occur only at the +selected rho sites, so the sum contains open paths, closed loops, and +path-plus-loop combinations: + +```python +from pepsy.bp import ( + OpenLoopSeriesCache, + partial_trace_open_loop_series_expand, + two_norm_bp, +) + +bp = two_norm_bp(peps.tn, max_iterations=1000, tol=1e-10) +rho = partial_trace_open_loop_series_expand( + peps.tn, + where=((0, 0), (0, 7)), + gloops=8, + messages=bp.messages, + run_bp=False, +) +``` + +This path performs an explicit configuration sum and normalizes only after +the sum; it does not apply the scalar disconnected-loop resummation used by +`partial_trace_edge_loop_series_expand`. + +For a cutoff sweep, reuse both the converged messages and the two caches. The +same `info` dictionary keeps already-contracted rho terms, while the +`OpenLoopSeriesCache` keeps the eligible edge configurations: + +```python +cache = OpenLoopSeriesCache() +info = {} +for cutoff in (2, 4, 6, 8): + rho = partial_trace_open_loop_series_expand( + peps.tn, + where=((0, 0), (0, 7)), + gloops=cutoff, + messages=bp.messages, + run_bp=False, + cache=cache, + info=info, + ) +``` + +For convergence diagnostics, inspect `info["open_rho_family_counts"]` and +`info["open_rho_family_weights"]`. The families are `open_path`, +`closed_loop`, and `path_plus_loop`; `open_rho_base_weight` is the unexcited +BP contribution. For native fermionic PEPS, treat this rho as a diagnostic +(trace and charge-block check), and evaluate fermionic operators through the +graded scalar APIs such as `compute_local_expectation_open_loop_series`. +On cyclic native fermionic graphs, those scalar and rho APIs use the +equivalent graded loop-cluster contraction internally because Symmray cannot +currently contract arbitrary mixed open ``P/Q`` configurations; the returned +rho remains native and the gate is still inserted in the ket/bra network. +The same one-BP/many-support workflow is runnable in the downstream example +`../pepsy_examples/symmetric_tensors/peps/bp_open_rho_series.py`; the +long-range native doublon comparison is in +`../pepsy_examples/symmetric_tensors/peps/bp_long_doublon_4x4.py`. + +To measure a long-range operator without materializing the diagnostic rho, use +the scalar companion. It keeps the same open-path, closed-loop, and +path-plus-loop family bookkeeping, inserts the observable into the physical +native contraction, and normalizes the accumulated numerator by the +accumulated denominator: + +```python +from pepsy import build_contraction +from pepsy.bp import compute_local_expectation_open_loop_series + +contraction_opt = build_contraction( + max_time=2.0, + max_repeats=8, + parallel=False, +) +value = compute_local_expectation_open_loop_series( + peps.tn, + {((0, 0), (0, 7)): fermion.hopping_operator()}, + gloops=8, + normalized=True, + optimize=contraction_opt, +) +``` + +The `optimize` object is forwarded to every Quimb contraction, so a reusable +`build_contraction` optimizer can cache Cotengra path searches across the +explicit loop terms. + +For a finite contraction budget, pass `max_flops_log10` and +`max_peak_memory_log2`. Each explicit term is contracted only when both +Cotengra tree diagnostics pass; on cyclic native fermionic graphs the same +limits are applied to the equivalent graded cluster contractions. +`info["open_scalar_term_costs"]` and `info["open_scalar_skipped_terms"]` +record the decision. The corresponding open-rho API exposes the same controls +and diagnostics. For a route-independent schema, use +`open_scalar_edge_term_costs` / `open_scalar_edge_skipped_terms` and +`open_scalar_cluster_region_costs` / +`open_scalar_cluster_region_skipped_terms`; the rho API provides the analogous +`open_rho_*` fields. The older `open_*_term_costs` names remain as +route-specific compatibility aliases. + +For native Symmray fermions, the scalar route keeps the physical gate in the +native ket/bra contraction and uses the graded open-bond Q projector when the +gate has local off-diagonal fermion action (for example hopping or pairing). +Diagonal density operators use the unphased open-bond projector. This +preserves the gate's graded ordering; a dense `trace(rho @ gate)` is not an +equivalent fermionic observable contraction. + +For reusable application code, `partial_trace_open_loop_series_sweep` wraps +the same pattern and accepts one-site, two-site, or larger retained supports: + +```python +from pepsy.bp import partial_trace_open_loop_series_sweep + +result = partial_trace_open_loop_series_sweep( + peps.tn, + supports=(((0, 0), (0, 7)), ((0, 0), (0, 1), (0, 7))), + cutoffs=(2, 4, 6), +) +rho = result.get_rho(((0, 0), (0, 7)), 4) +families = result.diagnostics[((0, 0), (0, 7))][4]["family_counts"] +``` + `partial_trace_loop_cluster_expand` and `compute_local_expectation_loop_cluster` provide the parallel D2BP generalized-loop-cluster route. Its default `combine="sum"` uses the usual diff --git a/docs/development/references/belief_propagation.md b/docs/development/references/belief_propagation.md index 6edf4f1..0db3165 100644 --- a/docs/development/references/belief_propagation.md +++ b/docs/development/references/belief_propagation.md @@ -23,6 +23,8 @@ side · **[roots]** foundational / prior art. | `loop_series_expand`, `LoopSeriesTerm`, `LoopSeriesCache` | edge-resolved `P + Q` loop series for D1BP and D2BP; retains excited-bond degree and distinct embeddings/chord subsets | Evenbly et al. 2409.03108 | | `partial_trace_loop_series_expand`, `compute_local_expectation_loop_series` | D2BP local reduced-density-matrix and scalar `P + Q` loop series; keeps physical output legs open and uses native Symmray virtual projectors | Evenbly et al. 2409.03108; quimb local loop-series API | | `partial_trace_edge_loop_series_expand`, `compute_local_expectation_edge_loop_series` | D2BP local RDM and graded scalar observable expansion over canonical explicit Q-edge terms; does not reinterpret Quimb's local-region cutoff | Evenbly et al. 2409.03108; Pepsy API | +| `partial_trace_open_loop_series_expand`, `partial_trace_open_loop_series_sweep` | Explicit D2BP rho configuration sum over open Q paths, closed loops, and attached or disconnected path-plus-loop terms; degree-one vertices are allowed only on retained rho sites; the sweep wrapper reuses one BP solve and cache across supports and cutoffs; cyclic native fermionic graphs use the equivalent native graded cluster contraction when mixed open P/Q contractions are unsupported | Evenbly et al. 2409.03108; Pepsy API | +| `compute_local_expectation_open_loop_series` | Scalar companion for long-range gates: reuses the open-loop family bookkeeping, inserts native gates through the graded open-bond projector route (or the equivalent graded cluster route on cyclic native fermionic graphs), normalizes numerator/denominator after the configuration estimate, and optionally filters explicit terms or cyclic cluster contractions by Cotengra log10-FLOP/log2-peak-size limits | Evenbly et al. 2409.03108; Pepsy API | | `partial_trace_loop_cluster_expand`, `compute_local_expectation_loop_cluster` | D2BP local reduced-density-matrix and scalar generalized-loop cluster expansion; combines BP-closed regions with inclusion--exclusion counts | Gray et al. 2510.05647; quimb local cluster API | | `loop_expand` | explicit selector between the edge loop series and region loop-cluster expansion; preserves each method's cutoff and result metadata | Pepsy API | | `partitioned_expand`, `pne_expand`, `PNEExpansionResult` | linear and combinatorial partitioned network expansions for D1BP/D2BP, with optional residue, explicit projectors, open outputs, and fixed recursive schedules | Evenbly, Gray & Chan 2512.10910 | diff --git a/src/pepsy/bp/__init__.py b/src/pepsy/bp/__init__.py index 71403fe..c867656 100644 --- a/src/pepsy/bp/__init__.py +++ b/src/pepsy/bp/__init__.py @@ -14,6 +14,12 @@ density matrix and its scalar expectation companion, * :func:`partial_trace_loop_series_expand` -- a D2BP reduced-density-matrix loop series and scalar local-observable companion, +* :func:`partial_trace_open_loop_series_expand` -- an explicit open-edge + rho series retaining long-range excitation paths and closed loops, +* :func:`partial_trace_open_loop_series_sweep` -- one-BP multi-support, + multi-cutoff open-rho diagnostics, +* :func:`compute_local_expectation_open_loop_series` -- direct gate-inserted + scalar expectations from the same open paths and loops, * :func:`loop_expand` -- an explicit selector between the correction families. * :func:`partitioned_expand` -- the partitioned network expansion (PNE, arXiv:2512.10910), with :func:`recursive_partitioned_expand` for fixed @@ -39,13 +45,18 @@ select_bp_candidate, ) from .series import ( + OpenLoopSeriesCache, + OpenLoopSeriesSweepResult, LoopSeriesCache, LoopSeriesResult, LoopSeriesTerm, compute_local_expectation_edge_loop_series, + compute_local_expectation_open_loop_series, compute_local_expectation_loop_cluster, compute_local_expectation_loop_series, partial_trace_edge_loop_series_expand, + partial_trace_open_loop_series_expand, + partial_trace_open_loop_series_sweep, partial_trace_loop_cluster_expand, partial_trace_loop_series_expand, loop_series_expand, @@ -129,12 +140,17 @@ "LoopClusterResult", "LoopClusterTerm", "LoopSeriesCache", + "OpenLoopSeriesCache", + "OpenLoopSeriesSweepResult", "LoopSeriesResult", "LoopSeriesTerm", "compute_local_expectation_edge_loop_series", + "compute_local_expectation_open_loop_series", "compute_local_expectation_loop_cluster", "compute_local_expectation_loop_series", "partial_trace_edge_loop_series_expand", + "partial_trace_open_loop_series_expand", + "partial_trace_open_loop_series_sweep", "partial_trace_loop_cluster_expand", "partial_trace_loop_series_expand", "RelayGaugeOptions", diff --git a/src/pepsy/bp/_symmray.py b/src/pepsy/bp/_symmray.py index 9d9d54f..928cd53 100644 --- a/src/pepsy/bp/_symmray.py +++ b/src/pepsy/bp/_symmray.py @@ -122,6 +122,56 @@ def dense_index_map(chargemap): return result +def _charge_parity(charge): + """Return the fermion parity of one Abelian charge.""" + if isinstance(charge, tuple): + return sum(int(component) for component in charge) % 2 + return int(charge) % 2 + + +def _dense_index_parities(index): + """Expand a Symmray index into the parity of each dense basis state.""" + parities = [] + for charge, size in index.chargemap.items(): + parities.extend([_charge_parity(charge)] * int(size)) + return parities + + +def _fermionic_open_q_phase(tn, index, dense): + """Apply the graded cup/cap phase to an open-bond Q operator. + + An open D2 bond has two bra legs followed by two ket legs. Splitting a + fermionic virtual bond into those independent copies changes the graded + ordering relative to the native tensor contraction. The basis-change + phase is ``-(-1) ** (p0*p1 + p0*p2 + p1*p2)`` for axes ordered as + ``(left-bra, left-ket, right-bra, right-ket)``. + """ + _, _, _, _, left_index, right_index = _bond_endpoint_data(tn, index) + left_parity = _dense_index_parities(left_index) + right_parity = _dense_index_parities(right_index) + parities = (left_parity, left_parity, right_parity, right_parity) + + original_shape = dense.shape + if dense.ndim == 2: + dimension = int(np.sqrt(dense.shape[0])) + dense = dense.reshape( + dimension, dimension, dimension, dimension + ) + if dense.ndim != 4: + raise ValueError("fermionic open Q operators must have rank four") + + left_bra = np.asarray(parities[0])[:, None, None, None] + left_ket = np.asarray(parities[1])[None, :, None, None] + right_bra = np.asarray(parities[2])[None, None, :, None] + exponent = ( + left_bra * left_ket + + left_bra * right_bra + + left_ket * right_bra + ) + phase = -((-1) ** exponent) + return (dense * phase).reshape(original_shape) + + def zero_charge(chargemap): """Return the neutral charge matching a scalar or product symmetry.""" charge = next(iter(chargemap), 0) @@ -258,8 +308,15 @@ def rank4_operator_from_dense(tn, index, operator, *, layout="pne"): not right_index.dual, right_index.dual, ) + elif layout == "open": + duals = ( + left_index.dual, + not left_index.dual, + right_index.dual, + not right_index.dual, + ) else: - raise ValueError("layout must be 'pne' or 'series'") + raise ValueError("layout must be 'pne', 'series', or 'open'") return array_cls.from_dense( dense, @@ -286,8 +343,22 @@ def rank_one_d2_projector( ) -def d2_operator(tn, index, operator, *, complement=False, layout="pne"): - """Normalize a D2 projector/operator and preserve native Symmray data.""" +def d2_operator( + tn, + index, + operator, + *, + complement=False, + layout="pne", + fermionic=False, +): + """Normalize a D2 operator and preserve native Symmray data. + + ``fermionic=True`` applies the graded open-bond cup/cap phase to a + complementary Q operator. It is intentionally opt-in because the phase + belongs to the physical open-observable ordering, not to ordinary D2BP + projectors. + """ dense = to_dense(operator) left, _, left_data, _, left_index, _ = _bond_endpoint_data(tn, index) del left @@ -302,4 +373,6 @@ def d2_operator(tn, index, operator, *, complement=False, layout="pne"): ) if complement: dense = np.eye(dimension * dimension, dtype=dense.dtype) - dense + if fermionic and layout == "open": + dense = _fermionic_open_q_phase(tn, index, dense) return rank4_operator_from_dense(tn, index, dense, layout=layout) diff --git a/src/pepsy/bp/series.py b/src/pepsy/bp/series.py index cb69730..0fee4c2 100644 --- a/src/pepsy/bp/series.py +++ b/src/pepsy/bp/series.py @@ -21,7 +21,7 @@ from __future__ import annotations -from collections import deque +from collections import Counter, deque from dataclasses import dataclass, field import functools from itertools import combinations @@ -49,13 +49,17 @@ ) __all__ = [ + "OpenLoopSeriesCache", + "OpenLoopSeriesSweepResult", "LoopSeriesCache", "LoopSeriesResult", "LoopSeriesTerm", "compute_local_expectation_edge_loop_series", + "compute_local_expectation_open_loop_series", "compute_local_expectation_loop_cluster", "partial_trace_loop_cluster_expand", "partial_trace_edge_loop_series_expand", + "partial_trace_open_loop_series_expand", "compute_local_expectation_loop_series", "partial_trace_loop_series_expand", "loop_series_expand", @@ -131,11 +135,118 @@ def terms_for(self, tn, max_degree: int) -> tuple[LoopSeriesTerm, ...]: try: return self.terms_by_max_degree[max_degree] except KeyError: - terms = _enumerate_edge_loops(tn, max_degree) + larger_degrees = [ + degree + for degree in self.terms_by_max_degree + if degree > max_degree + ] + if larger_degrees: + larger = self.terms_by_max_degree[min(larger_degrees)] + terms = tuple( + term for term in larger if term.degree <= max_degree + ) + else: + terms = _enumerate_edge_loops(tn, max_degree) self.terms_by_max_degree[max_degree] = terms return terms +@dataclass +class OpenLoopSeriesCache: + """Cache open-edge rho-series geometry for a fixed TN topology. + + Open rho terms depend on the selected physical support as well as the + tensor-network topology: degree-one Q vertices are allowed only on that + support, and bonds internal to the support are contracted exactly. This + cache keeps those choices in the key so it can safely be reused for a + family of observables on the same network. + """ + + terms_by_key: dict[tuple[Any, frozenset[Any], frozenset[Any]], tuple] = field( + default_factory=dict + ) + _topology_signature: Any = field(default=None, init=False, repr=False) + + def _check_topology(self, tn) -> None: + signature = LoopSeriesCache._signature(tn) + if self._topology_signature is None: + self._topology_signature = signature + elif self._topology_signature != signature: + raise ValueError( + "OpenLoopSeriesCache belongs to a different tensor-network " + "topology or tensor-id layout; create a fresh cache" + ) + + def terms_for( + self, + tn, + max_degree: int, + allowed_tids, + excluded_edges=(), + ) -> tuple[LoopSeriesTerm, ...]: + """Return open generalized-loop terms for one rho support.""" + self._check_topology(tn) + max_degree = _validate_nonnegative_degree(max_degree) + allowed_tids = frozenset(allowed_tids) + excluded_edges = frozenset(excluded_edges) + key = (max_degree, allowed_tids, excluded_edges) + try: + return self.terms_by_key[key] + except KeyError: + larger_keys = [ + known_key + for known_key in self.terms_by_key + if known_key[1:] == (allowed_tids, excluded_edges) + and known_key[0] > max_degree + ] + if larger_keys: + larger = self.terms_by_key[min(larger_keys)] + terms = tuple( + term for term in larger if term.degree <= max_degree + ) + else: + terms = _enumerate_open_edge_loops( + tn, + max_degree, + allowed_tids=allowed_tids, + excluded_edges=excluded_edges, + ) + self.terms_by_key[key] = terms + return terms + + +@dataclass +class OpenLoopSeriesSweepResult: + """One-BP cutoff sweep over one or more open-rho supports. + + ``rhos`` and ``diagnostics`` are keyed by ``tuple(where)`` and then by + integer Q-edge cutoff. The same D2BP message set and geometry cache are + used for every support and cutoff in the sweep. + """ + + rhos: dict[tuple[Any, ...], dict[int, Any]] + diagnostics: dict[tuple[Any, ...], dict[int, dict[str, Any]]] + infos: dict[tuple[Any, ...], dict[str, Any]] + bp: Any + cache: OpenLoopSeriesCache + bp_converged: bool | None + bp_iterations: int | None + bp_max_mdiff: float | None + + @property + def messages(self): + """Return the shared D2BP messages used throughout the sweep.""" + return self.bp.messages + + def get_rho(self, where, cutoff: int): + """Return one stored rho from the sweep.""" + support_key = tuple( + tuple(site) if isinstance(site, (list, tuple)) else site + for site in where + ) + return self.rhos[support_key][int(cutoff)] + + @dataclass class LoopSeriesResult: """Result of an edge-resolved BP loop-series contraction. @@ -162,6 +273,12 @@ class LoopSeriesResult: bp_iterations: int | None bp_max_mdiff: float | None bp: Any + requested_terms: tuple[LoopSeriesTerm, ...] = field(default_factory=tuple) + contraction_costs: dict[tuple[Any, ...], dict[str, float]] = field( + default_factory=dict + ) + skipped_terms: tuple[LoopSeriesTerm, ...] = () + cost_limits: dict[str, float | None] | None = None _normalized: bool = field(default=True, repr=False) _cache: LoopSeriesCache | None = field(default=None, repr=False) _contract_defaults: dict[str, Any] = field(default_factory=dict, repr=False) @@ -184,7 +301,9 @@ def expand( tol_correction: float | None = None, maxiter_correction: int | None = None, strip_exponent: bool = False, - optimize: str = "auto-hq", + optimize: Any = "auto-hq", + max_flops_log10: float | None = None, + max_peak_memory_log2: float | None = None, **contract_opts, ): """Evaluate another edge-loop cutoff using the same BP messages. @@ -199,6 +318,10 @@ def expand( tol_correction = self._contract_defaults["tol_correction"] if maxiter_correction is None: maxiter_correction = self._contract_defaults["maxiter_correction"] + if max_flops_log10 is None and self.cost_limits is not None: + max_flops_log10 = self.cost_limits["max_flops_log10"] + if max_peak_memory_log2 is None and self.cost_limits is not None: + max_peak_memory_log2 = self.cost_limits["max_peak_memory_log2"] terms = _parse_gloops(self.bp.tn, gloops, cache=self._cache) return _contract_loop_series( @@ -211,6 +334,8 @@ def expand( optimize=optimize, contract_opts=contract_opts, normalize=False, + max_flops_log10=max_flops_log10, + max_peak_memory_log2=max_peak_memory_log2, )[0] @@ -220,6 +345,12 @@ def _validate_degree(value: int) -> int: return int(value) +def _validate_nonnegative_degree(value: int) -> int: + if not isinstance(value, (int, np.integer)) or value < 0: + raise ValueError("open loop-series degree must be a non-negative integer") + return int(value) + + def _pairwise_edges(tn, *, norm: str): """Return deterministic pairwise edge records used by both BP families.""" edges = [] @@ -286,6 +417,147 @@ def _enumerate_edge_loops(tn, max_degree: int) -> tuple[LoopSeriesTerm, ...]: ) +def _open_term_from_edges( + tn, + edges, + *, + allowed_tids, + excluded_edges=(), +): + """Validate an edge subset whose only dangling sites are in ``allowed``. + + A local rho keeps its physical sites open. Consequently a Q-edge + configuration can have degree one at those sites, whereas a degree-one + vertex elsewhere is still a vanishing dangling excitation. Unlike the + global loop-series terms, an open configuration can also contain several + closed components attached to the open part, or several disconnected + closed components: those are the path-plus-loop terms used by the + long-range rho expansion. + """ + records = _edge_records(tn) + edges = tuple(edges) + if not edges: + raise ValueError("an open loop-series term must contain at least one edge") + if len(set(edges)) != len(edges): + raise ValueError("an open loop-series term cannot contain duplicate edges") + unknown = set(edges).difference(records) + if unknown: + raise ValueError(f"open loop-series term contains unknown bonds: {unknown!r}") + excluded_edges = set(excluded_edges) + internal = set(edges).intersection(excluded_edges) + if internal: + raise ValueError( + "open loop-series terms cannot excite bonds internal to the " + f"selected rho support: {internal!r}" + ) + + degrees: dict[Any, int] = {} + selected_tids = set() + for index in edges: + left, right = records[index] + selected_tids.update((left, right)) + degrees[left] = degrees.get(left, 0) + 1 + degrees[right] = degrees.get(right, 0) + 1 + + allowed_tids = frozenset(allowed_tids) + dangling = { + tid for tid, degree in degrees.items() if degree == 1 + } + invalid = dangling.difference(allowed_tids) + if invalid: + raise ValueError( + "open loop-series terms may have degree-one excitations only at " + f"the selected rho sites; invalid vertices: {invalid!r}" + ) + + return LoopSeriesTerm( + tuple(sorted(edges, key=repr)), + frozenset(selected_tids), + ) + + +def _enumerate_open_edge_loops( + tn, + max_degree: int, + *, + allowed_tids, + excluded_edges=(), +) -> tuple[LoopSeriesTerm, ...]: + """Enumerate open and closed Q-edge configurations for a rho support. + + The cutoff is the number of excited Q edges. Every non-support tensor + touched by a retained configuration must have at least two excited edges; + selected rho tensors may have one dangling excited edge. All edge + subsets are retained, including paths attached to or disconnected from + closed loops, matching the explicit expansion used by the rho notebook. + """ + max_degree = _validate_nonnegative_degree(max_degree) + if max_degree == 0: + return () + + excluded_edges = frozenset(excluded_edges) + edges = tuple( + edge + for edge in _pairwise_edges(tn, norm="2norm") + if edge[0] not in excluded_edges + ) + max_degree = min(max_degree, len(edges)) + allowed_tids = frozenset(allowed_tids) + remaining = Counter() + for _, left, right in edges: + remaining[left] += 1 + remaining[right] += 1 + + degrees: Counter[Any] = Counter() + selected = [] + terms = [] + + def has_closed_dangling_vertex(): + return any( + remaining[tid] == 0 + and degree == 1 + and tid not in allowed_tids + for tid, degree in degrees.items() + ) + + def visit(edge_pos, selected_count): + if edge_pos == len(edges): + if selected and not has_closed_dangling_vertex(): + terms.append( + LoopSeriesTerm( + tuple(sorted(selected, key=repr)), + frozenset(degrees), + ) + ) + return + + _, left, right = edges[edge_pos] + remaining[left] -= 1 + remaining[right] -= 1 + + if not has_closed_dangling_vertex(): + visit(edge_pos + 1, selected_count) + + if selected_count < max_degree: + selected.append(edges[edge_pos][0]) + degrees[left] += 1 + degrees[right] += 1 + if not has_closed_dangling_vertex(): + visit(edge_pos + 1, selected_count + 1) + degrees[left] -= 1 + degrees[right] -= 1 + selected.pop() + + remaining[left] += 1 + remaining[right] += 1 + + visit(0, 0) + + return tuple( + sorted(terms, key=lambda term: (term.degree, tuple(map(repr, term.edges)))) + ) + + def _edge_records(tn): return { index: (left, right) @@ -293,6 +565,79 @@ def _edge_records(tn): } +def _open_term_family(tn, term): + """Classify an open rho term for cutoff and convergence diagnostics.""" + records = _edge_records(tn) + adjacency: dict[Any, set[Any]] = {} + degrees: dict[Any, int] = {} + for index in term.edges: + left, right = records[index] + adjacency.setdefault(left, set()).add(right) + adjacency.setdefault(right, set()).add(left) + degrees[left] = degrees.get(left, 0) + 1 + degrees[right] = degrees.get(right, 0) + 1 + + dangling = any(degree == 1 for degree in degrees.values()) + components = 0 + unseen = set(adjacency) + while unseen: + components += 1 + stack = [unseen.pop()] + while stack: + tid = stack.pop() + for neighbor in adjacency[tid].intersection(unseen): + unseen.remove(neighbor) + stack.append(neighbor) + + if not dangling: + return "closed_loop" + cycle_rank = len(term.edges) - len(adjacency) + components + if components == 1 and cycle_rank == 0: + return "open_path" + return "path_plus_loop" + + +def _pairwise_graph_has_cycle(tn): + """Return whether the pairwise virtual graph contains a cycle. + + Symmray can contract the direct graded cluster networks, but its current + fermionic contraction path cannot represent an arbitrary mixture of + series-orientation ``P`` projectors and open-orientation ``Q`` projectors + on a cyclic graph. Tree graphs do not need the compatibility route below + and retain the explicit open-edge construction. + """ + vertices = set(tn.tensor_map) + edges = tuple(_pairwise_edges(tn, norm="2norm")) + if not edges: + return False + + parent = {vertex: vertex for vertex in vertices} + + def find(vertex): + while parent[vertex] != vertex: + parent[vertex] = parent[parent[vertex]] + vertex = parent[vertex] + return vertex + + for _, left, right in edges: + root_left = find(left) + root_right = find(right) + if root_left == root_right: + return True + parent[root_left] = root_right + return False + + +def _use_native_fermionic_cluster_open_route(bp, where): + """Whether native graded open observables need the cluster-compatible path.""" + return ( + len(where) > 1 + and _uses_symmray(bp.tn) + and bp.tn.isfermionic() + and _pairwise_graph_has_cycle(bp.tn) + ) + + def _connected_term_from_edges(tn, edges, *, tids=()): """Validate and canonicalize an explicit edge-resolved term.""" records = _edge_records(tn) @@ -601,26 +946,25 @@ def _get_d2_partial_trace_excited( inplace=False, ) + # Keep the ket, including a lazily inserted gate, before the graded bra. + # Symmray's fermionic contraction uses tensor order for parity routing. + for tid in stn.tensor_map: + local |= ket_tn.tensor_map[tid].reindex(kixmaps[tid]) + for tid, tensor in ket_tn.tensor_map.items(): + if tid not in stn.tensor_map: + local |= tensor + + # D2BP owns the graded bra tensors and their virtual dual-index map. + # ``tensor.conj()`` has the same numerical blocks but retains the ket + # virtual labels; that misses the fermionic bra ordering when boundary + # messages are attached. for tid in stn.tensor_map: - tensor = ket_tn.tensor_map[tid] - local |= tensor.reindex(kixmaps[tid]) - # D2BP owns the graded bra tensors and their virtual dual-index map. - # ``tensor.conj()`` has the same numerical blocks but retains the ket - # virtual labels; that misses the fermionic bra ordering when boundary - # messages are attached. bra_reindex = { bp.index_dual_map.get(index, index): new_index for index, new_index in bixmaps[tid].items() } local |= bp.tensor_dual_map[tid].reindex(bra_reindex) - # ``contract=False`` works for both adjacent and separated supports. The - # added gate tensor carries the graded physical routing between the ket and - # bra; original site tensors retain their ids above. - for tid, tensor in ket_tn.tensor_map.items(): - if tid not in stn.tensor_map: - local |= tensor - for index, tid in boundary_inds: data = bp.messages[index, tid] local |= qtn.Tensor( @@ -644,6 +988,7 @@ def _get_d2_partial_trace_excited( ), complement=True, layout="series", + fermionic=False, ) else: vacuum = ar.do( @@ -676,6 +1021,10 @@ def _get_d2_edge_partial_trace_excited( exclude=(), gate=None, gate_inds=(), + projector_layout="series", + gate_as_operator=False, + projector_index_order="bra-ket", + fermionic_q=False, ): """Build a D2 local RDM network with explicit P/Q edge choices. @@ -684,21 +1033,42 @@ def _get_d2_edge_partial_trace_excited( ``exclude`` are traced directly. This is deliberately separate from :func:`_get_d2_partial_trace_excited`, whose non-excluded bonds are all ``Q`` for Quimb's local-region convention. + + ``fermionic_q`` applies the native graded cup/cap phase to open-bond Q + tensors. It is gate-aware because diagonal density operators and + off-diagonal hopping/pairing operators use different physical ordering + conventions. """ import quimb.tensor as qtn stn = bp.tn._select_tids(tids) excited_edges = set(excited_edges) exclude = set(exclude) + gate_inds = tuple(gate_inds) + if projector_index_order not in {"bra-ket", "ket-bra"}: + raise ValueError( + "projector_index_order must be 'bra-ket' or 'ket-bra'" + ) kixmaps = {tid: {} for tid in stn.tensor_map} bixmaps = {tid: {} for tid in stn.tensor_map} projector_inds = {} boundary_inds = [] + gate_index_map = {} for index, region_tids in stn.ind_map.items(): region_tids = tuple(region_tids) if index in bp.output_inds: - if index in partial_trace_map: + if gate_as_operator and index in gate_inds: + (tid,) = region_tids + kix = qtn.rand_uuid() + kixmaps[tid][index] = kix + # ``tensor_network_gate_inds`` represents a gate with its + # original physical labels on the first (bra/output) legs + # and fresh labels on the second (ket/input) legs. Preserve + # that convention so native fermionic gate phases see the + # same graded ordering as Quimb's exact contraction path. + gate_index_map[index] = (index, kix) + elif index in partial_trace_map: (tid,) = region_tids bixmaps[tid][index] = partial_trace_map[index] elif index in exclude: @@ -711,7 +1081,10 @@ def _get_d2_edge_partial_trace_excited( bix = qtn.rand_uuid() kixmaps[tid][index] = kix bixmaps[tid][index] = bix - projector_inds.setdefault(index, {})[tid] = (bix, kix) + if projector_index_order == "bra-ket": + projector_inds.setdefault(index, {})[tid] = (bix, kix) + else: + projector_inds.setdefault(index, {})[tid] = (kix, bix) else: (tid,) = region_tids kix = qtn.rand_uuid() @@ -721,7 +1094,7 @@ def _get_d2_edge_partial_trace_excited( boundary_inds.append((index, tid)) local = qtn.TensorNetwork() - if gate is None: + if gate is None or gate_as_operator: ket_tn = stn else: ket_tn = qtn.tensor_network_gate_inds( @@ -734,16 +1107,41 @@ def _get_d2_edge_partial_trace_excited( inplace=False, ) + # Keep the ket, including a lazily inserted gate, before the graded bra. + # Symmray's fermionic contraction uses tensor order for parity routing. for tid in stn.tensor_map: local |= ket_tn.tensor_map[tid].reindex(kixmaps[tid]) + if gate_as_operator: + try: + gate_inds_local = tuple( + gate_index_map[index][side] + for side in (0, 1) + for index in gate_inds + ) + except KeyError as exc: + raise ValueError( + "gate_inds must be physical output indices in the selected " + "D2BP region" + ) from exc + gate_tensor = gate + if ar.do("ndim", gate_tensor) == 2: + gate_tensor = ar.do( + "reshape", + gate_tensor, + tuple(stn.ind_size(index) for index in gate_inds) * 2, + ) + local |= qtn.Tensor(gate_tensor, inds=gate_inds_local) + else: + for tid, tensor in ket_tn.tensor_map.items(): + if tid not in stn.tensor_map: + local |= tensor + + for tid in stn.tensor_map: bra_reindex = { bp.index_dual_map.get(index, index): new_index for index, new_index in bixmaps[tid].items() } local |= bp.tensor_dual_map[tid].reindex(bra_reindex) - for tid, tensor in ket_tn.tensor_map.items(): - if tid not in stn.tensor_map: - local |= tensor for index, tid in boundary_inds: local |= qtn.Tensor( @@ -766,7 +1164,8 @@ def _get_d2_edge_partial_trace_excited( index, p0, complement=index in excited_edges, - layout="series", + layout=projector_layout, + fermionic=fermionic_q, ) else: p0 = ar.do("einsum", "i,j->ij", ml.reshape(-1), mr.reshape(-1)) @@ -792,6 +1191,88 @@ def _rho_trace(rho): return ar.do("trace", rho) +def _validate_contraction_cost_limits( + max_flops_log10, + max_peak_memory_log2, +): + """Validate optional log-cost limits for explicit term contractions.""" + for name, value in ( + ("max_flops_log10", max_flops_log10), + ("max_peak_memory_log2", max_peak_memory_log2), + ): + if value is not None: + if not isinstance(value, (int, float, np.integer, np.floating)): + raise TypeError(f"{name} must be a real number or None") + if not np.isfinite(value) or value < 0: + raise ValueError(f"{name} must be finite and non-negative") + return ( + None if max_flops_log10 is None else float(max_flops_log10), + None + if max_peak_memory_log2 is None + else float(max_peak_memory_log2), + ) + + +def _contract_cost_record(tree): + """Extract the standard Cotengra log-cost diagnostics from a tree.""" + return { + "flops_log10": float(tree.total_flops(log=10)), + "peak_memory_log2": float(tree.peak_size(log=2)), + } + + +def _contract_with_cost_limits( + network, + *, + optimize, + contract_opts, + max_flops_log10=None, + max_peak_memory_log2=None, +): + """Contract ``network`` or return its cost record when over budget. + + The cost pass builds a Cotengra tree but does not contract any tensor + data. If accepted, the same tree is supplied to the numerical contraction + so path search is not repeated. ``peak_memory_log2`` follows Cotengra's + convention: log2 of the largest concurrently live scalar tensor size. + """ + if max_flops_log10 is None and max_peak_memory_log2 is None: + return ( + True, + network.contract(optimize=optimize, **contract_opts), + None, + ) + + if "get" in contract_opts: + raise TypeError( + "contract_opts['get'] cannot be combined with contraction cost " + "limits; use the loop-series cost diagnostics instead" + ) + tree = network.contract( + get="tree", + optimize=optimize, + **contract_opts, + ) + cost = _contract_cost_record(tree) + accepted = ( + ( + max_flops_log10 is None + or cost["flops_log10"] <= max_flops_log10 + ) + and ( + max_peak_memory_log2 is None + or cost["peak_memory_log2"] <= max_peak_memory_log2 + ) + ) + if not accepted: + return False, None, cost + return ( + True, + network.contract(optimize=tree, **contract_opts), + cost, + ) + + def _term_sites(tn, where): """Normalize a Quimb local-term key to an ordered site tuple.""" has_site = getattr(tn, "has_site", None) @@ -945,6 +1426,8 @@ def _partial_trace_loop_cluster( optimize, contract_opts, info, + max_flops_log10=None, + max_peak_memory_log2=None, ): """Contract native D2BP generalized-loop cluster density matrices.""" if bp.__class__.__name__ != "D2BP": @@ -988,6 +1471,16 @@ def _partial_trace_loop_cluster( from quimb.tensor.belief_propagation import gen_region_counts term_cache = {} if info is None else info.setdefault("cluster_rho_terms", {}) + term_cost_cache = ( + {} + if info is None + else info.setdefault("cluster_rho_term_costs", {}) + ) + skipped_terms = ( + {} + if info is None + else info.setdefault("cluster_rho_skipped_terms", {}) + ) rhos = [] counts = [] for region, count in gen_region_counts(regions, autocomplete=autocomplete): @@ -1000,11 +1493,21 @@ def _partial_trace_loop_cluster( region, partial_trace_map=partial_trace_map, ) - rho_r = cluster.contract( - output_inds=output_inds, + cluster_contract_opts = dict(contract_opts) + cluster_contract_opts.setdefault("output_inds", output_inds) + accepted, rho_r, cost = _contract_with_cost_limits( + cluster, optimize=optimize, - **contract_opts, - ).to_dense(kix, bix) + contract_opts=cluster_contract_opts, + max_flops_log10=max_flops_log10, + max_peak_memory_log2=max_peak_memory_log2, + ) + if not accepted: + skipped_terms[region] = cost + continue + if cost is not None: + term_cost_cache[region] = cost + rho_r = rho_r.to_dense(kix, bix) term_cache[cache_key] = rho_r if normalized == "local": @@ -1222,6 +1725,92 @@ def _edge_series_terms_for_support(bp, tids, gloops, *, cache): return terms, inner_bonds +def _open_edge_series_terms_for_support(bp, tids, gloops, *, cache): + """Parse open rho terms, allowing dangling Q edges at ``tids``.""" + inner_bonds = frozenset(bp.tn._select_tids(tids).inner_inds()) + allowed_tids = frozenset(tids) + + if isinstance(gloops, (int, np.integer)): + cutoff = _validate_nonnegative_degree(gloops) + if cache is None: + terms = _enumerate_open_edge_loops( + bp.tn, + cutoff, + allowed_tids=allowed_tids, + excluded_edges=inner_bonds, + ) + else: + terms = cache.terms_for( + bp.tn, + cutoff, + allowed_tids, + excluded_edges=inner_bonds, + ) + return terms, inner_bonds + + if gloops is None: + max_degree = sum( + index not in inner_bonds + for index, _, _ in _pairwise_edges(bp.tn, norm="2norm") + ) + if cache is None: + terms = _enumerate_open_edge_loops( + bp.tn, + max_degree, + allowed_tids=allowed_tids, + excluded_edges=inner_bonds, + ) + else: + terms = cache.terms_for( + bp.tn, + max_degree, + allowed_tids, + excluded_edges=inner_bonds, + ) + return terms, inner_bonds + + edge_labels = { + index + for index, _, _ in _pairwise_edges(bp.tn, norm="2norm") + if index not in inner_bonds + } + terms = [] + seen = set() + for item in tuple(gloops): + if isinstance(item, LoopSeriesTerm): + edges = item.edges + elif hasattr(item, "edges"): + edges = item.edges + else: + try: + edges = tuple(item) + except TypeError as exc: + raise TypeError( + "explicit open rho terms must be LoopSeriesTerm objects " + "or iterables of virtual-edge labels" + ) from exc + if not edges or not set(edges).issubset(edge_labels): + unknown = set(edges).difference(edge_labels) + raise ValueError( + "explicit open rho terms must use pairwise virtual edges " + f"outside the selected support; unknown edges: {unknown!r}" + ) + term = _open_term_from_edges( + bp.tn, + edges, + allowed_tids=allowed_tids, + excluded_edges=inner_bonds, + ) + if term in seen: + raise ValueError(f"duplicate open rho term: {term.edges!r}") + seen.add(term) + terms.append(term) + + return tuple( + sorted(terms, key=lambda term: (term.degree, tuple(map(repr, term.edges)))), + ), inner_bonds + + def _edge_series_suppression( weights, *, @@ -1344,59 +1933,326 @@ def _partial_trace_edge_loop_series( return rho -def _local_expectation_edge_loop_series( +def _partial_trace_open_loop_series( bp, where, - gate, gloops, *, normalized, - multi_excitation_correct, - tol_correction, - maxiter_correction, optimize, contract_opts, cache, info, + max_flops_log10, + max_peak_memory_log2, ): - """Direct graded scalar counterpart of the explicit edge RDM series.""" + """Contract the explicit open-edge rho loop-series expansion.""" + if bp.__class__.__name__ != "D2BP": + raise ValueError( + "partial_trace_open_loop_series_expand currently requires " + "norm='2norm'" + ) if normalized == "prod": normalized = True - if normalized not in (True, False, "local", "separate"): + if normalized not in (True, False, "separate"): raise ValueError( - "normalized must be one of True, False, 'prod', 'local', or " - "'separate'" + "normalized must be one of True, False, 'prod', or 'separate'" ) + _align_symmray_d2bp_messages(bp) bp.normalize_message_pairs() bp.normalize_tensors() tags = [bp.tn.site_tag(coo) for coo in where] tids = frozenset(bp.tn._get_tids_from_tags(tags, "any")) - terms, inner_bonds = _edge_series_terms_for_support( - bp, tids, gloops, cache=cache - ) - if _uses_symmray(bp.tn) and len(where) > 1 and terms: - raise NotImplementedError( - "fermionic explicit-edge loop corrections for multi-site gates " - "are not supported yet; use gloops=0, a one-site term, or a " - "separate exact/path observable route" - ) + if not tids: + raise ValueError("where must contain at least one site in the network") + kix = [bp.tn.site_ind(coo) for coo in where] - norm_cache = {} if info is None else info.setdefault("edge_series_norm_terms", {}) - norm_terms = {} - gate_terms = {} - for term in terms: - region = frozenset((*tids, *term.tids)) - cache_key = (term.edges, region, tuple(where)) + import quimb.tensor as qtn + + # Keep the physical bra indices stable when ``info`` is reused for a + # sequence of cutoffs. Besides avoiding needless index churn, this is + # required for cached native (e.g. Symmray) rho terms to remain valid. + where_key = tuple(where) + if info is None: + bix = [qtn.rand_uuid() for _ in where] + partial_trace_map = dict(zip(kix, bix)) + else: + partial_trace_maps = info.setdefault("open_rho_partial_trace_maps", {}) + map_key = (where_key, tuple(kix)) try: - norm_e = norm_cache[cache_key] + partial_trace_map = partial_trace_maps[map_key] except KeyError: - norm_e = _get_d2_edge_partial_trace_excited( - bp, - region, - excited_edges=term.edges, - exclude=inner_bonds, + bix = [qtn.rand_uuid() for _ in where] + partial_trace_map = dict(zip(kix, bix)) + partial_trace_maps[map_key] = partial_trace_map + bix = [partial_trace_map[index] for index in kix] + output_inds = (*kix, *bix) + terms, inner_bonds = _open_edge_series_terms_for_support( + bp, + tids, + gloops, + cache=cache, + ) + requested_terms = terms + max_flops_log10, max_peak_memory_log2 = _validate_contraction_cost_limits( + max_flops_log10, + max_peak_memory_log2, + ) + + if _use_native_fermionic_cluster_open_route(bp, where): + # See the scalar counterpart below. This preserves a native + # fermionic rho on cyclic graphs while avoiding the unsupported mixed + # P/Q contraction path in Symmray. + cluster_info = {} + rho = _partial_trace_loop_cluster( + bp, + where, + gloops, + combine="sum", + normalized=normalized, + autocomplete=True, + grow_from="alldangle", + strict_size=False, + optimize=optimize, + contract_opts=contract_opts, + info=cluster_info, + max_flops_log10=max_flops_log10, + max_peak_memory_log2=max_peak_memory_log2, + ) + term_families = { + term.edges: _open_term_family(bp.tn, term) + for term in requested_terms + } + family_counts = Counter(term_families.values()) + family_weights = {family: 0.0 for family in family_counts} + if info is not None: + cluster_region_costs = { + (where_key, region): cost + for region, cost in cluster_info.get( + "cluster_rho_term_costs", {} + ).items() + } + cluster_region_skipped = { + (where_key, region): cost + for region, cost in cluster_info.get( + "cluster_rho_skipped_terms", {} + ).items() + } + info["open_rho_requested_terms"] = requested_terms + info["open_rho_terms_list"] = requested_terms + info["open_rho_term_costs"] = dict( + cluster_info.get("cluster_rho_term_costs", {}) + ) + info["open_rho_skipped_terms"] = dict( + cluster_info.get("cluster_rho_skipped_terms", {}) + ) + info["open_rho_cost_limits"] = { + "max_flops_log10": max_flops_log10, + "max_peak_memory_log2": max_peak_memory_log2, + } + info["open_rho_edge_term_costs"] = {} + info["open_rho_edge_skipped_terms"] = {} + info["open_rho_cluster_region_costs"] = cluster_region_costs + info[ + "open_rho_cluster_region_skipped_terms" + ] = cluster_region_skipped + info["open_rho_weights"] = {} + info["open_rho_term_families"] = term_families + info["open_rho_family_counts"] = dict(family_counts) + info["open_rho_family_weights"] = family_weights + info["open_rho_base_weight"] = _rho_trace(rho) + info["open_rho_support_tids"] = tids + info["open_rho_excluded_edges"] = inner_bonds + info["open_rho_native_route"] = "graded_cluster_compatible" + return rho + + term_cache = {} if info is None else info.setdefault("open_rho_terms", {}) + term_cost_cache = ( + {} if info is None else info.setdefault("open_rho_term_costs", {}) + ) + skipped_terms = ( + {} if info is None else info.setdefault("open_rho_skipped_terms", {}) + ) + rho_terms = {} + accepted_terms = [] + for term in terms: + if term.edges in skipped_terms: + continue + region = frozenset((*tids, *term.tids)) + cache_key = (term.edges, region, where_key) + try: + rho_e = term_cache[cache_key] + except KeyError: + rho_network = _get_d2_edge_partial_trace_excited( + bp, + region, + excited_edges=term.edges, + partial_trace_map=partial_trace_map, + exclude=inner_bonds, + projector_layout="open", + ) + accepted, rho_e, cost = _contract_with_cost_limits( + rho_network, + optimize=optimize, + contract_opts={"output_inds": output_inds, **contract_opts}, + max_flops_log10=max_flops_log10, + max_peak_memory_log2=max_peak_memory_log2, + ) + if not accepted: + skipped_terms[term.edges] = cost + continue + rho_e = rho_e.to_dense(kix, bix) + term_cache[cache_key] = rho_e + term_cost_cache[term.edges] = cost + rho_terms[term.edges] = rho_e + accepted_terms.append(term) + + base_cache = ( + {} + if info is None + else info.setdefault("open_rho_base_terms", {}) + ) + base_key = (where_key, tuple(kix), tids, inner_bonds) + try: + base = base_cache[base_key] + except KeyError: + base_network = _get_d2_edge_partial_trace_excited( + bp, + tids, + partial_trace_map=partial_trace_map, + exclude=inner_bonds, + projector_layout="open", + ) + accepted, base, base_cost = _contract_with_cost_limits( + base_network, + optimize=optimize, + contract_opts={"output_inds": output_inds, **contract_opts}, + max_flops_log10=max_flops_log10, + max_peak_memory_log2=max_peak_memory_log2, + ) + if not accepted: + raise ValueError( + "the unexcited open rho configuration exceeds the " + "contraction cost limits: " + f"{base_cost!r}" + ) + base = base.to_dense(kix, bix) + if info is not None: + base_cache[base_key] = base + + weights = {edges: _rho_trace(rho_e) for edges, rho_e in rho_terms.items()} + terms = tuple(accepted_terms) + term_families = { + term.edges: _open_term_family(bp.tn, term) + for term in terms + } + family_counts = Counter(term_families.values()) + family_weights = { + family: sum( + weight + for edges, weight in weights.items() + if term_families[edges] == family + ) + for family in family_counts + } + + # This is an explicit configuration sum, so attached and disconnected + # path-plus-loop terms are already present. Applying Quimb's scalar + # multi-excitation resummation here would reweight those terms a second + # time. + rho = base + for rho_e in rho_terms.values(): + rho = rho + rho_e + + if normalized in (True, "separate"): + rho = rho / _rho_trace(rho) + elif (bp.sign, bp.exponent) != (1.0, 0.0): + rho = rho * bp.sign * 10**bp.exponent + + if info is not None: + info["open_rho_edge_term_costs"] = { + (where_key, edges): cost + for edges, cost in term_cost_cache.items() + } + info["open_rho_edge_skipped_terms"] = { + (where_key, edges): cost + for edges, cost in skipped_terms.items() + } + info["open_rho_cluster_region_costs"] = {} + info["open_rho_cluster_region_skipped_terms"] = {} + info["open_rho_requested_terms"] = requested_terms + info["open_rho_terms_list"] = terms + info["open_rho_term_costs"] = dict(term_cost_cache) + info["open_rho_skipped_terms"] = dict(skipped_terms) + info["open_rho_cost_limits"] = { + "max_flops_log10": max_flops_log10, + "max_peak_memory_log2": max_peak_memory_log2, + } + info["open_rho_weights"] = weights + info["open_rho_term_families"] = term_families + info["open_rho_family_counts"] = dict(family_counts) + info["open_rho_family_weights"] = family_weights + info["open_rho_base_weight"] = _rho_trace(base) + info["open_rho_support_tids"] = tids + info["open_rho_excluded_edges"] = inner_bonds + return rho + + +def _local_expectation_edge_loop_series( + bp, + where, + gate, + gloops, + *, + normalized, + multi_excitation_correct, + tol_correction, + maxiter_correction, + optimize, + contract_opts, + cache, + info, +): + """Direct graded scalar counterpart of the explicit edge RDM series.""" + if normalized == "prod": + normalized = True + if normalized not in (True, False, "local", "separate"): + raise ValueError( + "normalized must be one of True, False, 'prod', 'local', or " + "'separate'" + ) + _align_symmray_d2bp_messages(bp) + bp.normalize_message_pairs() + bp.normalize_tensors() + + tags = [bp.tn.site_tag(coo) for coo in where] + tids = frozenset(bp.tn._get_tids_from_tags(tags, "any")) + terms, inner_bonds = _edge_series_terms_for_support( + bp, tids, gloops, cache=cache + ) + if _uses_symmray(bp.tn) and len(where) > 1 and terms: + raise NotImplementedError( + "fermionic explicit-edge loop corrections for multi-site gates " + "are not supported yet; use gloops=0, a one-site term, or a " + "separate exact/path observable route" + ) + kix = [bp.tn.site_ind(coo) for coo in where] + norm_cache = {} if info is None else info.setdefault("edge_series_norm_terms", {}) + norm_terms = {} + gate_terms = {} + for term in terms: + region = frozenset((*tids, *term.tids)) + cache_key = (term.edges, region, tuple(where)) + try: + norm_e = norm_cache[cache_key] + except KeyError: + norm_e = _get_d2_edge_partial_trace_excited( + bp, + region, + excited_edges=term.edges, + exclude=inner_bonds, ).contract(optimize=optimize, **contract_opts) norm_cache[cache_key] = norm_e gate_e = _get_d2_edge_partial_trace_excited( @@ -1448,6 +2304,368 @@ def _local_expectation_edge_loop_series( return value, norm +def _gate_needs_fermionic_open_q(gate): + """Return whether an open native gate has local off-diagonal action.""" + dense = _symmray_to_dense(gate) + if dense.ndim == 0: + return False + if dense.ndim == 2: + dimension = int(np.sqrt(dense.shape[0])) + if dimension * dimension != dense.shape[0]: + return False + dense = dense.reshape(dimension, dimension, dimension, dimension) + if dense.ndim != 4: + return False + matrix = dense.reshape( + dense.shape[0] * dense.shape[1], + dense.shape[2] * dense.shape[3], + ) + diagonal = np.diag(np.diag(matrix)) + return not np.allclose(matrix, diagonal, rtol=1e-12, atol=1e-14) + + +def _local_expectation_open_loop_series( + bp, + where, + gate, + gloops, + *, + normalized, + optimize, + contract_opts, + cache, + info, + max_flops_log10, + max_peak_memory_log2, +): + """Contract a gate through the explicit open-edge loop series.""" + if normalized == "prod": + normalized = True + if normalized not in (True, False, "separate"): + raise ValueError( + "normalized must be one of True, False, 'prod', or 'separate'" + ) + + _align_symmray_d2bp_messages(bp) + bp.normalize_message_pairs() + bp.normalize_tensors() + + tags = [bp.tn.site_tag(coo) for coo in where] + tids = frozenset(bp.tn._get_tids_from_tags(tags, "any")) + if not tids: + raise ValueError("where must contain at least one site in the network") + + kix = [bp.tn.site_ind(coo) for coo in where] + terms, inner_bonds = _open_edge_series_terms_for_support( + bp, + tids, + gloops, + cache=cache, + ) + max_flops_log10, max_peak_memory_log2 = _validate_contraction_cost_limits( + max_flops_log10, + max_peak_memory_log2, + ) + where_key = tuple(where) + fermionic_q = _uses_symmray(bp.tn) and _gate_needs_fermionic_open_q(gate) + + if _use_native_fermionic_cluster_open_route(bp, where): + # The explicit open-edge decomposition is exact for dense networks + # and for fermionic trees. On a cyclic native Symmray graph, however, + # a mixed P/Q network can require a non-pairwise fermionic contraction + # that Symmray does not currently support. The direct cluster form is + # algebraically equivalent at this level and keeps the gate inside + # the graded ket/bra contraction, so it is a safe native fallback. + cluster_info = {} + value, normalization = _local_expectation_loop_cluster( + bp, + where, + gate, + gloops, + combine="sum", + normalized=normalized, + autocomplete=True, + grow_from="alldangle", + strict_size=False, + optimize=optimize, + contract_opts=contract_opts, + info=cluster_info, + max_flops_log10=max_flops_log10, + max_peak_memory_log2=max_peak_memory_log2, + ) + term_families = { + term.edges: _open_term_family(bp.tn, term) for term in terms + } + family_counts = Counter(term_families.values()) + if info is not None: + cluster_region_costs = { + (where_key, region): cost + for region, cost in cluster_info.get( + "cluster_scalar_term_costs", {} + ).items() + } + cluster_region_skipped = { + (where_key, region): cost + for region, cost in cluster_info.get( + "cluster_scalar_skipped_terms", {} + ).items() + } + info["open_scalar_requested_terms"] = terms + info["open_scalar_terms"] = terms + info["open_scalar_skipped_terms"] = dict( + cluster_info.get("cluster_scalar_skipped_terms", {}) + ) + info["open_scalar_term_costs"] = dict( + cluster_info.get("cluster_scalar_term_costs", {}) + ) + info["open_scalar_norm_weights"] = {} + info["open_scalar_gate_terms"] = {} + info["open_scalar_term_families"] = term_families + info["open_scalar_family_counts"] = dict(family_counts) + info["open_scalar_family_weights"] = {} + info["open_scalar_base_weight"] = normalization + info["open_scalar_numerator"] = value * normalization + info["open_scalar_denominator"] = normalization + info["open_scalar_excluded_edges"] = inner_bonds + info["open_scalar_native_route"] = ( + "graded_cluster_compatible" + ) + info["open_scalar_fermionic_q_phase"] = fermionic_q + info["open_scalar_cost_limits"] = { + "max_flops_log10": max_flops_log10, + "max_peak_memory_log2": max_peak_memory_log2, + } + info["open_scalar_edge_term_costs"] = {} + info["open_scalar_edge_skipped_terms"] = {} + info["open_scalar_cluster_region_costs"] = cluster_region_costs + info[ + "open_scalar_cluster_region_skipped_terms" + ] = cluster_region_skipped + return value, normalization + + norm_cache = ( + {} if info is None else info.setdefault("open_scalar_norm_terms", {}) + ) + norm_cost_cache = ( + {} + if info is None + else info.setdefault("open_scalar_norm_term_costs", {}) + ) + gate_cache = ( + {} + if info is None + else info.setdefault("open_scalar_gate_term_cache", {}) + ) + gate_cost_cache = ( + {} + if info is None + else info.setdefault("open_scalar_gate_term_costs", {}) + ) + norm_terms = {} + gate_terms = {} + accepted_terms = [] + term_costs = {} if info is None else info.setdefault( + "open_scalar_term_costs", {} + ) + skipped_terms = {} if info is None else info.setdefault( + "open_scalar_skipped_terms", {} + ) + for term in terms: + if (where_key, term.edges) in skipped_terms: + continue + region = frozenset((*tids, *term.tids)) + cache_key = (term.edges, region, where_key, fermionic_q) + try: + norm_e = norm_cache[cache_key] + norm_cost = norm_cost_cache.get((where_key, term.edges)) + except KeyError: + norm_network = _get_d2_edge_partial_trace_excited( + bp, + region, + excited_edges=term.edges, + exclude=inner_bonds, + projector_layout=( + "open" if _uses_symmray(bp.tn) else "series" + ), + fermionic_q=fermionic_q, + ) + accepted, norm_e, norm_cost = _contract_with_cost_limits( + norm_network, + optimize=optimize, + contract_opts=contract_opts, + max_flops_log10=max_flops_log10, + max_peak_memory_log2=max_peak_memory_log2, + ) + if not accepted: + skipped_terms[(where_key, term.edges)] = {"norm": norm_cost} + continue + norm_cache[cache_key] = norm_e + norm_cost_cache[(where_key, term.edges)] = norm_cost + + gate_cache_key = (where_key, term.edges, fermionic_q, id(gate)) + try: + gate_e = gate_cache[gate_cache_key] + gate_cost = gate_cost_cache.get(gate_cache_key) + accepted = True + except KeyError: + gate_network = _get_d2_edge_partial_trace_excited( + bp, + region, + excited_edges=term.edges, + exclude=inner_bonds, + gate=gate, + gate_inds=kix, + projector_layout=( + "open" if _uses_symmray(bp.tn) else "series" + ), + gate_as_operator=True, + fermionic_q=fermionic_q, + ) + accepted, gate_e, gate_cost = _contract_with_cost_limits( + gate_network, + optimize=optimize, + contract_opts=contract_opts, + max_flops_log10=max_flops_log10, + max_peak_memory_log2=max_peak_memory_log2, + ) + if accepted: + gate_cache[gate_cache_key] = gate_e + gate_cost_cache[gate_cache_key] = gate_cost + if not accepted: + skipped_terms[(where_key, term.edges)] = { + "norm": norm_cost, + "gate": gate_cost, + } + norm_cache.pop(cache_key, None) + continue + accepted_terms.append(term) + term_costs[(where_key, term.edges)] = { + "norm": norm_cost, + "gate": gate_cost, + "flops_log10": max( + (norm_cost or {"flops_log10": 0.0})["flops_log10"], + (gate_cost or {"flops_log10": 0.0})["flops_log10"], + ), + "peak_memory_log2": max( + (norm_cost or {"peak_memory_log2": 0.0})[ + "peak_memory_log2" + ], + (gate_cost or {"peak_memory_log2": 0.0})[ + "peak_memory_log2" + ], + ), + } + norm_terms[term.edges] = norm_e + gate_terms[term.edges] = gate_e + + base_key = (where_key, tuple(kix), tids, inner_bonds, fermionic_q) + base_cache = ( + {} if info is None else info.setdefault("open_scalar_base_terms", {}) + ) + try: + base_norm = base_cache[base_key] + except KeyError: + base_network = _get_d2_edge_partial_trace_excited( + bp, + tids, + exclude=inner_bonds, + projector_layout=( + "open" if _uses_symmray(bp.tn) else "series" + ), + fermionic_q=fermionic_q, + ) + accepted, base_norm, base_cost = _contract_with_cost_limits( + base_network, + optimize=optimize, + contract_opts=contract_opts, + max_flops_log10=max_flops_log10, + max_peak_memory_log2=max_peak_memory_log2, + ) + if not accepted: + raise ValueError( + "the unexcited open scalar configuration exceeds the " + "contraction cost limits: " + f"{base_cost!r}" + ) + if info is not None: + base_cache[base_key] = base_norm + + base_gate_network = _get_d2_edge_partial_trace_excited( + bp, + tids, + exclude=inner_bonds, + gate=gate, + gate_inds=kix, + projector_layout=( + "open" if _uses_symmray(bp.tn) else "series" + ), + gate_as_operator=True, + fermionic_q=fermionic_q, + ) + accepted, base_value, base_gate_cost = _contract_with_cost_limits( + base_gate_network, + optimize=optimize, + contract_opts=contract_opts, + max_flops_log10=max_flops_log10, + max_peak_memory_log2=max_peak_memory_log2, + ) + if not accepted: + raise ValueError( + "the unexcited open scalar gate configuration exceeds the " + "contraction cost limits: " + f"{base_gate_cost!r}" + ) + + norm = base_norm + sum(norm_terms.values()) + raw_value = base_value + sum(gate_terms.values()) + value = raw_value + if normalized: + value = value / norm + elif (bp.sign, bp.exponent) != (1.0, 0.0): + value = value * bp.sign * 10**bp.exponent + + if info is not None: + term_families = { + term.edges: _open_term_family(bp.tn, term) + for term in accepted_terms + } + family_counts = Counter(term_families.values()) + family_weights = { + family: sum( + norm_terms[edges] + for edges, term_family in term_families.items() + if term_family == family + ) + for family in family_counts + } + info["open_scalar_requested_terms"] = terms + info["open_scalar_terms"] = tuple(accepted_terms) + info["open_scalar_edge_term_costs"] = dict(term_costs) + info["open_scalar_edge_skipped_terms"] = dict(skipped_terms) + info["open_scalar_cluster_region_costs"] = {} + info["open_scalar_cluster_region_skipped_terms"] = {} + info["open_scalar_cost_limits"] = { + "max_flops_log10": max_flops_log10, + "max_peak_memory_log2": max_peak_memory_log2, + } + info["open_scalar_norm_weights"] = dict(norm_terms) + info["open_scalar_gate_terms"] = dict(gate_terms) + info["open_scalar_term_families"] = term_families + info["open_scalar_family_counts"] = dict(family_counts) + info["open_scalar_family_weights"] = family_weights + info["open_scalar_base_weight"] = base_norm + info["open_scalar_numerator"] = raw_value + info["open_scalar_denominator"] = norm + info["open_scalar_excluded_edges"] = inner_bonds + info["open_scalar_native_route"] = ( + "graded_open_projectors" + if _uses_symmray(bp.tn) + else "dense_open_projectors" + ) + info["open_scalar_fermionic_q_phase"] = fermionic_q + return value, norm + + def _local_expectation_loop_cluster( bp, where, @@ -1462,6 +2680,8 @@ def _local_expectation_loop_cluster( optimize, contract_opts, info, + max_flops_log10=None, + max_peak_memory_log2=None, ): """Contract one gate through the graded D2BP loop-cluster network.""" if combine != "sum": @@ -1496,6 +2716,16 @@ def _local_expectation_loop_cluster( term_cache = ( {} if info is None else info.setdefault("cluster_norm_terms", {}) ) + term_cost_cache = ( + {} + if info is None + else info.setdefault("cluster_scalar_term_costs", {}) + ) + skipped_terms = ( + {} + if info is None + else info.setdefault("cluster_scalar_skipped_terms", {}) + ) norm_terms = [] gate_terms = [] counts = [] @@ -1505,17 +2735,41 @@ def _local_expectation_loop_cluster( try: norm_e = term_cache[cache_key] except KeyError: - norm_e = _get_d2_cluster_norm(bp, region).contract( + norm_network = _get_d2_cluster_norm(bp, region) + accepted, norm_e, norm_cost = _contract_with_cost_limits( + norm_network, optimize=optimize, - **contract_opts, + contract_opts=contract_opts, + max_flops_log10=max_flops_log10, + max_peak_memory_log2=max_peak_memory_log2, ) + if not accepted: + skipped_terms[region] = {"norm": norm_cost} + continue term_cache[cache_key] = norm_e + else: + norm_cost = None gate_e = _get_d2_cluster_norm( bp, region, gate=gate, gate_inds=kix, - ).contract(optimize=optimize, **contract_opts) + ) + accepted, gate_e, gate_cost = _contract_with_cost_limits( + gate_e, + optimize=optimize, + contract_opts=contract_opts, + max_flops_log10=max_flops_log10, + max_peak_memory_log2=max_peak_memory_log2, + ) + if not accepted: + skipped_terms[region] = {"gate": gate_cost} + continue + if max_flops_log10 is not None or max_peak_memory_log2 is not None: + term_cost_cache[region] = { + "norm": norm_cost, + "gate": gate_cost, + } if normalized == "local": gate_e = gate_e / norm_e norm_e = 1.0 @@ -1580,6 +2834,9 @@ def _contract_loop_series( optimize, contract_opts, normalize, + max_flops_log10=None, + max_peak_memory_log2=None, + return_diagnostics=False, ): if normalize: if bp.__class__.__name__ == "D2BP": @@ -1588,10 +2845,24 @@ def _contract_loop_series( bp.normalize_tensors() weights = {} + accepted_terms = [] + contraction_costs = {} + skipped_terms = [] for term in terms: - weights[term.edges] = _get_edge_excited(bp, term).contract( - optimize=optimize, **contract_opts + accepted, weight, cost = _contract_with_cost_limits( + _get_edge_excited(bp, term), + optimize=optimize, + contract_opts=contract_opts, + max_flops_log10=max_flops_log10, + max_peak_memory_log2=max_peak_memory_log2, ) + if not accepted: + skipped_terms.append(term) + continue + weights[term.edges] = weight + accepted_terms.append(term) + if cost is not None: + contraction_costs[term.edges] = cost estimate, correction, suppression = _process_weights( weights, mantissa=bp.sign, @@ -1601,7 +2872,10 @@ def _contract_loop_series( maxiter_correction=maxiter_correction, strip_exponent=strip_exponent, ) - return estimate, weights, correction, suppression + result = estimate, weights, correction, suppression + if return_diagnostics: + return result + (tuple(accepted_terms), contraction_costs, tuple(skipped_terms)) + return result def _build_bp( @@ -1773,7 +3047,9 @@ def loop_series_expand( multi_excitation_correct: bool = True, tol_correction: float = 1e-12, maxiter_correction: int = 100, - optimize: str = "auto-hq", + optimize: Any = "auto-hq", + max_flops_log10: float | None = None, + max_peak_memory_log2: float | None = None, strip_exponent: bool = False, progbar: bool = False, contract_opts: dict[str, Any] | None = None, @@ -1795,6 +3071,10 @@ def loop_series_expand( bond gauges. The loop-series formal cancellation assumes a fixed point; set ``require_fixed_point=False`` only for an explicitly exploratory boundary approximation. + + Optional ``max_flops_log10`` and ``max_peak_memory_log2`` limits are + applied to each explicit Q-edge contraction using Cotengra's tree + diagnostics. Skipped terms are exposed on the returned result. """ if run_bp and ( not isinstance(max_iterations, (int, np.integer)) or max_iterations < 1 @@ -1833,7 +3113,7 @@ def loop_series_expand( cache = cache or LoopSeriesCache() terms = _parse_gloops(bp.tn, gloops, cache=cache) - estimate, weights, correction, suppression = _contract_loop_series( + result = _contract_loop_series( bp, terms, multi_excitation_correct=multi_excitation_correct, @@ -1843,12 +3123,24 @@ def loop_series_expand( optimize=optimize, contract_opts=contract_opts, normalize=True, + max_flops_log10=max_flops_log10, + max_peak_memory_log2=max_peak_memory_log2, + return_diagnostics=True, ) + ( + estimate, + weights, + correction, + suppression, + accepted_terms, + contraction_costs, + skipped_terms, + ) = result return LoopSeriesResult( estimate=estimate, gloops=gloops, norm=str(norm).lower(), - terms=terms, + terms=accepted_terms, loop_weights=weights, free_energy_correction=correction, suppression_factors=suppression, @@ -1857,6 +3149,13 @@ def loop_series_expand( bp_iterations=info.get("iterations"), bp_max_mdiff=info.get("max_mdiff"), bp=bp, + requested_terms=terms, + contraction_costs=contraction_costs, + skipped_terms=skipped_terms, + cost_limits={ + "max_flops_log10": max_flops_log10, + "max_peak_memory_log2": max_peak_memory_log2, + }, _cache=cache, _contract_defaults={ "tol_correction": tol_correction, @@ -1887,7 +3186,7 @@ def partial_trace_loop_series_expand( grow_from: str = "alldangle", strict_size: bool = False, multi_excitation_correct: bool = True, - optimize: str = "auto-hq", + optimize: Any = "auto-hq", info: dict[str, Any] | None = None, contract_opts: dict[str, Any] | None = None, **bp_opts, @@ -2080,6 +3379,316 @@ def partial_trace_edge_loop_series_expand( ) +def partial_trace_open_loop_series_sweep( + tn, + supports, + cutoffs, + *, + messages=None, + gauges=None, + run_bp: bool = True, + bp_runner: str = "plain", + relay_opts: dict[str, Any] | None = None, + max_iterations: int = 1000, + tol: float = 5e-6, + tol_abs: float | None = None, + tol_rolling_diff: float | None = 0.0, + diis: bool | dict[str, Any] = False, + damping: float = 0.0, + update: str = "sequential", + require_fixed_point: bool = True, + normalized: bool | str = True, + cache: OpenLoopSeriesCache | None = None, + optimize: Any = "auto-hq", + max_flops_log10: float | None = None, + max_peak_memory_log2: float | None = None, + contract_opts: dict[str, Any] | None = None, + **bp_opts, +) -> OpenLoopSeriesSweepResult: + """Sweep open-rho cutoffs and supports after one D2BP construction. + + Parameters + ---------- + tn : TensorNetwork + A PEPS-like network for native D2BP. + supports : iterable of sequences + The physical sites to retain for each rho. List and tuple site + coordinates are normalized to hashable tuples. + cutoffs : iterable of int or int + Maximum numbers of excited Q edges to evaluate, in order. + messages, gauges, run_bp, ... + BP controls matching :func:`partial_trace_open_loop_series_expand`. + + Returns + ------- + OpenLoopSeriesSweepResult + Rhos, per-cutoff family diagnostics, support-local contraction caches, + and the shared D2BP object. + + Notes + ----- + ``supports`` can contain one-site, two-site, or larger retained regions. + Virtual bonds internal to each support are contracted exactly by the + underlying open-rho expansion. For native fermionic PEPS, the returned + rhos are diagnostics; evaluate operators through the graded scalar APIs. + """ + if isinstance(cutoffs, (int, np.integer)): + cutoffs = (int(cutoffs),) + else: + cutoffs = tuple(cutoffs) + if not cutoffs: + raise ValueError("cutoffs must contain at least one non-negative integer") + if any( + not isinstance(cutoff, (int, np.integer)) or cutoff < 0 + for cutoff in cutoffs + ): + raise ValueError("cutoffs must contain only non-negative integers") + cutoffs = tuple(dict.fromkeys(int(cutoff) for cutoff in cutoffs)) + + supports = tuple( + tuple( + tuple(site) if isinstance(site, (list, tuple)) else site + for site in support + ) + for support in supports + ) + if not supports: + raise ValueError("supports must contain at least one retained site set") + if any(not support for support in supports): + raise ValueError("each support must contain at least one site") + + contract_opts = {} if contract_opts is None else dict(contract_opts) + cache = cache or OpenLoopSeriesCache() + bp, bp_info = _build_bp( + tn, + norm="2norm", + messages=messages, + gauges=gauges, + run_bp=run_bp, + bp_runner=bp_runner, + relay_opts=relay_opts, + max_iterations=max_iterations, + tol=tol, + tol_abs=tol_abs, + tol_rolling_diff=tol_rolling_diff, + diis=diis, + damping=damping, + update=update, + optimize=optimize, + bp_opts=bp_opts, + progbar=False, + ) + if require_fixed_point and run_bp and not bp_info.get("converged", False): + raise RuntimeError( + "partial_trace_open_loop_series_sweep requires converged BP " + "messages; pass require_fixed_point=False for an exploratory sweep" + ) + + rhos = {} + diagnostics = {} + infos = {} + for support in supports: + support_key = tuple(support) + support_info = {} + support_rhos = {} + support_diagnostics = {} + for cutoff in cutoffs: + rho = _partial_trace_open_loop_series( + bp, + support, + cutoff, + normalized=normalized, + optimize=optimize, + max_flops_log10=max_flops_log10, + max_peak_memory_log2=max_peak_memory_log2, + contract_opts=contract_opts, + cache=cache, + info=support_info, + ) + support_rhos[cutoff] = rho + support_diagnostics[cutoff] = { + "term_count": len(support_info["open_rho_terms_list"]), + "requested_term_count": len( + support_info["open_rho_requested_terms"] + ), + "skipped_term_count": len( + support_info["open_rho_skipped_terms"] + ), + "family_counts": dict( + support_info["open_rho_family_counts"] + ), + "family_weights": dict( + support_info["open_rho_family_weights"] + ), + "base_weight": support_info["open_rho_base_weight"], + "trace": _rho_trace(rho), + "term_costs": dict(support_info["open_rho_term_costs"]), + "edge_term_costs": dict( + support_info["open_rho_edge_term_costs"] + ), + "edge_skipped_terms": dict( + support_info["open_rho_edge_skipped_terms"] + ), + "cluster_region_costs": dict( + support_info["open_rho_cluster_region_costs"] + ), + "cluster_region_skipped_terms": dict( + support_info[ + "open_rho_cluster_region_skipped_terms" + ] + ), + "cost_limits": dict(support_info["open_rho_cost_limits"]), + } + rhos[support_key] = support_rhos + diagnostics[support_key] = support_diagnostics + infos[support_key] = support_info + + return OpenLoopSeriesSweepResult( + rhos=rhos, + diagnostics=diagnostics, + infos=infos, + bp=bp, + cache=cache, + bp_converged=bp_info.get("converged"), + bp_iterations=bp_info.get("iterations"), + bp_max_mdiff=bp_info.get("max_mdiff"), + ) + + +def partial_trace_open_loop_series_expand( + tn, + where, + gloops=None, + *, + messages=None, + gauges=None, + run_bp: bool = True, + bp_runner: str = "plain", + relay_opts: dict[str, Any] | None = None, + max_iterations: int = 1000, + tol: float = 5e-6, + tol_abs: float | None = None, + tol_rolling_diff: float | None = 0.0, + diis: bool | dict[str, Any] = False, + damping: float = 0.0, + update: str = "sequential", + require_fixed_point: bool = True, + normalized: bool | str = True, + cache: OpenLoopSeriesCache | None = None, + optimize: Any = "auto-hq", + max_flops_log10: float | None = None, + max_peak_memory_log2: float | None = None, + info: dict[str, Any] | None = None, + contract_opts: dict[str, Any] | None = None, + **bp_opts, +): + """Compute a long-range local rho from an explicit open-edge series. + + The integer ``gloops`` cutoff counts excited ``Q`` virtual edges. A + retained edge subset may have degree one only at one of the selected + physical rho sites; every other touched tensor must have either zero or at + least two excited edges. This keeps the open excitation paths connecting + separated sites, closed loops, and attached or disconnected combinations + of a path with closed loops. The selected rho sites' internal virtual bonds are + contracted exactly and are not expanded. + + This is intentionally separate from + :func:`partial_trace_loop_series_expand`, whose integer cutoff follows + Quimb's local tensor-region convention. It is also separate from + :func:`partial_trace_edge_loop_series_expand`, which retains only closed + connected generalized-loop terms and applies the scalar + multi-excitation resummation. Here the configuration sum is explicit, so + all retained terms have unit coefficient and are normalized only after + summation. + + ``messages`` can be supplied from a previously converged D2BP run with + ``run_bp=False``. This is the intended route for measuring many + long-range rho supports after one BP solve. ``gloops=None`` enumerates up + to the number of eligible pairwise virtual bonds and can be expensive; + use an integer cutoff for practical calculations. + + Parameters + ---------- + tn : TensorNetwork + A PEPS-like network for the native D2BP calculation. + where : sequence + The physical sites to retain in the reduced density matrix. + gloops : int or iterable, optional + Maximum number of excited Q edges, or explicit virtual-edge subsets. + normalized : bool or {"prod", "separate"}, optional + Whether to normalize the final explicit configuration sum. + optimize : str or path optimizer, optional + Quimb contraction optimizer. A reusable optimizer from + :func:`pepsy.build_contraction` can be passed here and is forwarded + unchanged to every term contraction. + max_flops_log10, max_peak_memory_log2 : float, optional + Optional Cotengra tree-cost limits. Terms over either limit are + skipped and recorded in ``info``; the unexcited base configuration + must still fit both limits. + cache : OpenLoopSeriesCache, optional + Reusable open-term geometry cache for the same topology and support. + info : dict, optional + Receives the contracted terms and their trace weights. Reuse the same + dictionary for a cutoff sweep to avoid recontracting earlier terms. + The current cutoff's ``open_rho_family_counts`` and + ``open_rho_family_weights`` separate open paths, closed loops, and + path-plus-loop configurations. Cost metadata is exposed in the stable + ``open_rho_edge_*`` and ``open_rho_cluster_region_*`` fields; the + older ``open_rho_term_costs`` fields are route-specific aliases. + """ + contract_opts = {} if contract_opts is None else dict(contract_opts) + where = tuple(where) + if not where: + raise ValueError("where must contain at least one site") + if normalized not in (True, False, "prod", "separate"): + raise ValueError( + "normalized must be one of True, False, 'prod', or 'separate'" + ) + if run_bp and ( + not isinstance(max_iterations, (int, np.integer)) or max_iterations < 1 + ): + raise ValueError("max_iterations must be a positive integer when run_bp=True") + + bp, bp_info = _build_bp( + tn, + norm="2norm", + messages=messages, + gauges=gauges, + run_bp=run_bp, + bp_runner=bp_runner, + relay_opts=relay_opts, + max_iterations=max_iterations, + tol=tol, + tol_abs=tol_abs, + tol_rolling_diff=tol_rolling_diff, + diis=diis, + damping=damping, + update=update, + optimize=optimize, + bp_opts=bp_opts, + progbar=False, + ) + if require_fixed_point and run_bp and not bp_info.get("converged", False): + raise RuntimeError( + "partial_trace_open_loop_series_expand requires converged BP " + "messages; pass require_fixed_point=False for an exploratory " + "estimate" + ) + + return _partial_trace_open_loop_series( + bp, + where, + gloops, + normalized=normalized, + optimize=optimize, + max_flops_log10=max_flops_log10, + max_peak_memory_log2=max_peak_memory_log2, + contract_opts=contract_opts, + cache=cache or OpenLoopSeriesCache(), + info=info, + ) + + def partial_trace_loop_cluster_expand( tn, where, @@ -2103,7 +3712,7 @@ def partial_trace_loop_cluster_expand( autocomplete: bool = True, grow_from: str = "alldangle", strict_size: bool = False, - optimize: str = "auto-hq", + optimize: Any = "auto-hq", info: dict[str, Any] | None = None, contract_opts: dict[str, Any] | None = None, **bp_opts, @@ -2499,3 +4108,174 @@ def compute_local_expectation_edge_loop_series( if return_all: return expecs return functools.reduce(operator.add, expecs.values()) + + +def compute_local_expectation_open_loop_series( + tn, + terms, + gloops=None, + *, + messages=None, + gauges=None, + run_bp: bool = True, + bp_runner: str = "plain", + relay_opts: dict[str, Any] | None = None, + max_iterations: int = 1000, + tol: float = 5e-6, + tol_abs: float | None = None, + tol_rolling_diff: float | None = 0.0, + diis: bool | dict[str, Any] = False, + damping: float = 0.0, + update: str = "sequential", + require_fixed_point: bool = True, + normalized: bool | str = True, + cache: OpenLoopSeriesCache | None = None, + optimize: Any = "auto-hq", + max_flops_log10: float | None = None, + max_peak_memory_log2: float | None = None, + info: dict[str, Any] | None = None, + return_all: bool = False, + contract_opts: dict[str, Any] | None = None, + **bp_opts, +): + """Compute fermion-safe expectations from open-edge loop terms. + + The terms mapping has the usual site-or-sites to gate form. Unlike the + open rho API, this function evaluates the observable as a scalar numerator + and identity denominator. Dense networks use direct gate insertion over + the explicit open-path, closed-loop, and path-plus-loop configurations; + native fermionic networks use native graded open-bond projectors and keep + the gate in the ket/bra contraction, so the graded gate ordering is + preserved without materializing a diagnostic rho. + + Integer gloops counts excited virtual edges. For native fermionic PEPS, + this is the preferred route for even two-site operators such as hopping, + pairing, and density terms. + + ``optimize`` is forwarded unchanged to Quimb's + ``TensorNetwork.contract`` calls. Callers can pass the reusable Cotengra + optimizer returned by :func:`pepsy.build_contraction` to reuse contraction + path searches across the loop terms. + + ``max_flops_log10`` and ``max_peak_memory_log2`` optionally filter each + explicit configuration, or each cyclic native cluster region, using + Cotengra's tree diagnostics. A contraction is performed only when both + limits pass; skipped terms are reported in + ``info["open_scalar_edge_skipped_terms"]`` or + ``info["open_scalar_cluster_region_skipped_terms"]`` depending on the + native route. The older ``open_scalar_skipped_terms`` field is a + route-specific compatibility alias. + """ + if not hasattr(terms, "items"): + raise TypeError("terms must be a mapping from sites to operators") + if not terms: + raise ValueError("terms must contain at least one operator") + if normalized == "prod": + normalized = True + if normalized not in (True, False, "separate"): + raise ValueError( + "normalized must be one of True, False, 'prod', or 'separate'" + ) + contract_opts = {} if contract_opts is None else dict(contract_opts) + bp, bp_info = _build_bp( + tn, + norm="2norm", + messages=messages, + gauges=gauges, + run_bp=run_bp, + bp_runner=bp_runner, + relay_opts=relay_opts, + max_iterations=max_iterations, + tol=tol, + tol_abs=tol_abs, + tol_rolling_diff=tol_rolling_diff, + diis=diis, + damping=damping, + update=update, + optimize=optimize, + bp_opts=bp_opts, + progbar=False, + ) + if require_fixed_point and run_bp and not bp_info.get("converged", False): + raise RuntimeError( + "compute_local_expectation_open_loop_series requires converged " + "BP messages; pass require_fixed_point=False for an exploratory " + "estimate" + ) + + cache = cache or OpenLoopSeriesCache() + term_info = ( + {} + if info is None + else info.setdefault("open_scalar_normalization_by_term", {}) + ) + support_info = ( + {} + if info is None + else info.setdefault("open_scalar_supports", {}) + ) + expecs = {} + for where, gate in terms.items(): + sites = _term_sites(bp.tn, where) + value, normalization = _local_expectation_open_loop_series( + bp, + sites, + gate, + gloops, + normalized=normalized, + optimize=optimize, + max_flops_log10=max_flops_log10, + max_peak_memory_log2=max_peak_memory_log2, + contract_opts=contract_opts, + cache=cache, + info=info, + ) + term_info[where] = normalization + expecs[where] = value + if info is not None: + support_key = tuple(sites) + support_edge_costs = { + key: value + for key, value in info["open_scalar_edge_term_costs"].items() + if key[0] == support_key + } + support_edge_skipped = { + key: value + for key, value in info[ + "open_scalar_edge_skipped_terms" + ].items() + if key[0] == support_key + } + support_cluster_costs = { + key: value + for key, value in info[ + "open_scalar_cluster_region_costs" + ].items() + if key[0] == support_key + } + support_cluster_skipped = { + key: value + for key, value in info[ + "open_scalar_cluster_region_skipped_terms" + ].items() + if key[0] == support_key + } + support_info[support_key] = { + "terms": tuple(info["open_scalar_terms"]), + "requested_terms": tuple( + info["open_scalar_requested_terms"] + ), + "skipped_terms": dict( + info["open_scalar_skipped_terms"] + ), + "term_costs": dict(info["open_scalar_term_costs"]), + "edge_skipped_terms": support_edge_skipped, + "edge_term_costs": support_edge_costs, + "cluster_region_skipped_terms": support_cluster_skipped, + "cluster_region_costs": support_cluster_costs, + "family_counts": dict(info["open_scalar_family_counts"]), + "family_weights": dict(info["open_scalar_family_weights"]), + } + if return_all: + return expecs + return functools.reduce(operator.add, expecs.values()) diff --git a/tests/test_bp_open_series.py b/tests/test_bp_open_series.py new file mode 100644 index 0000000..4e90f88 --- /dev/null +++ b/tests/test_bp_open_series.py @@ -0,0 +1,464 @@ +"""Tests for the explicit open-edge BP rho loop series.""" + +from collections import Counter +from itertools import combinations + +import numpy as np +import quimb.tensor as qtn + +from pepsy.bp import ( + OpenLoopSeriesCache, + OpenLoopSeriesSweepResult, + compute_local_expectation_open_loop_series, + partial_trace_open_loop_series_expand, + partial_trace_open_loop_series_sweep, + two_norm_bp, +) +from pepsy.bp.series import _open_term_family + + +def _edge_degrees(tn, edges): + degrees = {} + for index in edges: + left, right = tn.ind_map[index] + degrees[left] = degrees.get(left, 0) + 1 + degrees[right] = degrees.get(right, 0) + 1 + return degrees + + +def _notebook_terms(tn, max_degree, allowed_tids, excluded_edges): + """Mirror ``quf.combine_elements`` without importing the notebook code.""" + edges = tuple( + index + for index, tids in tn.ind_map.items() + if len(tids) == 2 and index not in excluded_edges + ) + expected = set() + for degree in range(1, max_degree + 1): + for selected in combinations(edges, degree): + degrees = _edge_degrees(tn, selected) + dangling = { + tid for tid, value in degrees.items() if value == 1 + } + if not dangling or dangling <= allowed_tids: + expected.add(frozenset(selected)) + return expected + + +def test_open_rho_series_keeps_the_long_range_path_and_is_exact_on_a_tree(): + state = qtn.PEPS.rand( + 1, + 4, + bond_dim=2, + phys_dim=2, + seed=1908, + dtype="complex128", + ) + where = ((0, 0), (0, 3)) + exact = state.partial_trace( + where, + max_bond=64, + optimize="auto-hq", + flatten=True, + normalized=True, + ) + info = {} + rho = partial_trace_open_loop_series_expand( + state, + where, + gloops=3, + max_iterations=200, + tol=1e-10, + diis=False, + info=info, + ) + + terms = info["open_rho_terms_list"] + assert len(terms) == 1 + assert terms[0].degree == 3 + assert np.max(np.abs(rho - exact)) < 1e-10 + np.testing.assert_allclose(np.trace(rho), 1.0, atol=1e-12) + + +def test_open_scalar_series_inserts_a_two_site_gate_and_normalizes_after_sum(): + state = qtn.PEPS.rand( + 1, + 4, + bond_dim=2, + phys_dim=2, + seed=1930, + dtype="complex128", + ) + where = ((0, 0), (0, 3)) + gate = np.diag([1.0, 2.0, 3.0, 4.0]) + exact = state.compute_local_expectation_exact( + {where: gate}, normalized=True, optimize="auto-hq" + ) + info = {} + value = compute_local_expectation_open_loop_series( + state, + {where: gate}, + gloops=3, + max_iterations=200, + tol=1e-10, + diis=False, + info=info, + ) + + np.testing.assert_allclose(value, exact, rtol=1e-10, atol=1e-12) + np.testing.assert_allclose(info["open_scalar_denominator"], 1.0, atol=1e-12) + np.testing.assert_allclose(info["open_scalar_numerator"], value, atol=1e-12) + + +def test_open_scalar_series_reports_and_applies_contraction_cost_limits(): + """Cost limits screen terms using Cotengra tree diagnostics.""" + state = qtn.PEPS.rand( + 1, + 4, + bond_dim=2, + phys_dim=2, + seed=1932, + dtype="complex128", + ) + where = ((0, 0), (0, 3)) + gate = np.diag([1.0, 2.0, 3.0, 4.0]) + info = {} + value = compute_local_expectation_open_loop_series( + state, + {where: gate}, + gloops=3, + max_iterations=200, + tol=1e-10, + diis=False, + info=info, + max_flops_log10=2.0, + max_peak_memory_log2=30.0, + ) + + assert info["open_scalar_cost_limits"] == { + "max_flops_log10": 2.0, + "max_peak_memory_log2": 30.0, + } + assert info["open_scalar_edge_term_costs"] == info[ + "open_scalar_term_costs" + ] + assert info["open_scalar_edge_skipped_terms"] == info[ + "open_scalar_skipped_terms" + ] + assert not info["open_scalar_cluster_region_costs"] + assert not info["open_scalar_terms"] + assert info["open_scalar_skipped_terms"] + cost = next(iter(info["open_scalar_skipped_terms"].values()))["norm"] + assert set(cost) == {"flops_log10", "peak_memory_log2"} + assert cost["flops_log10"] > 2.0 + assert np.isfinite(value) + + +def test_open_rho_series_allows_only_rho_site_dangling_vertices(): + state = qtn.PEPS.rand( + 2, + 3, + bond_dim=2, + phys_dim=2, + seed=1909, + dtype="complex128", + ) + where = ((0, 0), (0, 2)) + tags = [state.site_tag(coo) for coo in where] + allowed_tids = frozenset(state._get_tids_from_tags(tags, "any")) + excluded_edges = frozenset(state._select_tids(allowed_tids).inner_inds()) + + terms = OpenLoopSeriesCache().terms_for( + state, + 4, + allowed_tids, + excluded_edges=excluded_edges, + ) + assert any( + {tid for tid, degree in _edge_degrees(state, term.edges).items() if degree == 1} + == allowed_tids.intersection(_edge_degrees(state, term.edges)) + and len(term.edges) == 2 + for term in terms + ) + assert any( + not any(degree == 1 for degree in _edge_degrees(state, term.edges).values()) + and len(term.edges) == 4 + for term in terms + ) + for term in terms: + dangling = { + tid for tid, degree in _edge_degrees(state, term.edges).items() if degree == 1 + } + assert dangling <= allowed_tids + + +def test_open_rho_terms_match_notebook_filter_and_classify_disconnected_loops(): + state = qtn.PEPS.rand( + 3, + 3, + bond_dim=2, + phys_dim=2, + seed=1913, + dtype="complex128", + ) + where = ((0, 0), (0, 2)) + tags = [state.site_tag(coo) for coo in where] + allowed_tids = frozenset(state._get_tids_from_tags(tags, "any")) + excluded_edges = frozenset(state._select_tids(allowed_tids).inner_inds()) + terms = OpenLoopSeriesCache().terms_for( + state, + 6, + allowed_tids, + excluded_edges=excluded_edges, + ) + + assert {frozenset(term.edges) for term in terms} == _notebook_terms( + state, + 6, + allowed_tids, + excluded_edges, + ) + families = Counter(_open_term_family(state, term) for term in terms) + assert families["open_path"] + assert families["closed_loop"] + assert families["path_plus_loop"] + + +def test_four_by_four_long_range_cutoffs_have_expected_geometry(): + """The first 4x4 terms are paths, plaquettes, and attached loops.""" + state = qtn.PEPS.rand( + 4, + 4, + bond_dim=2, + phys_dim=2, + seed=1915, + dtype="complex128", + ) + where = ((1, 1), (3, 3)) + allowed_tids = frozenset( + state._get_tids_from_tags([state.site_tag(coo) for coo in where], "any") + ) + excluded_edges = frozenset(state._select_tids(allowed_tids).inner_inds()) + cache = OpenLoopSeriesCache() + + geometries = {} + for cutoff in (4, 5, 6): + terms = cache.terms_for( + state, + cutoff, + allowed_tids, + excluded_edges=excluded_edges, + ) + geometries[cutoff] = Counter( + (len(term.edges), _open_term_family(state, term)) + for term in terms + ) + + assert geometries[4] == Counter( + {(4, "open_path"): 6, (4, "closed_loop"): 9} + ) + assert geometries[5] == geometries[4] + Counter( + {(5, "path_plus_loop"): 6} + ) + assert geometries[6] == geometries[5] + Counter( + { + (6, "open_path"): 14, + (6, "closed_loop"): 12, + (6, "path_plus_loop"): 14, + } + ) + + +def test_open_rho_series_reports_term_families_and_reuses_one_bp_run(): + state = qtn.PEPS.rand( + 3, + 3, + bond_dim=2, + phys_dim=2, + seed=1914, + dtype="complex128", + ) + where = ((0, 0), (0, 2)) + bp = two_norm_bp( + state, + max_iterations=200, + tol=1e-10, + diis=False, + ) + info = {} + cache = OpenLoopSeriesCache() + rho = partial_trace_open_loop_series_expand( + state, + where, + gloops=4, + messages=bp.messages, + run_bp=False, + cache=cache, + info=info, + ) + + assert info["open_rho_family_counts"]["open_path"] + assert info["open_rho_family_counts"]["closed_loop"] + assert set(info["open_rho_term_families"].values()) == { + "open_path", + "closed_loop", + } + assert info["open_rho_edge_term_costs"] + assert not info["open_rho_cluster_region_costs"] + np.testing.assert_allclose(np.trace(rho), 1.0, atol=1e-12) + + other_info = {} + other_rho = partial_trace_open_loop_series_expand( + state, + ((0, 0), (2, 2)), + gloops=4, + messages=bp.messages, + run_bp=False, + cache=cache, + info=other_info, + ) + assert other_info["open_rho_terms_list"] + np.testing.assert_allclose(np.trace(other_rho), 1.0, atol=1e-12) + assert len(cache.terms_by_key) == 2 + + +def test_open_rho_series_reuses_one_d2bp_message_set(): + state = qtn.PEPS.rand( + 1, + 4, + bond_dim=2, + phys_dim=2, + seed=1910, + dtype="complex128", + ) + where = ((0, 0), (0, 3)) + bp = two_norm_bp( + state, + max_iterations=200, + tol=1e-10, + diis=False, + ) + fresh = partial_trace_open_loop_series_expand( + state, + where, + gloops=3, + max_iterations=200, + tol=1e-10, + diis=False, + ) + reused = partial_trace_open_loop_series_expand( + state, + where, + gloops=3, + messages=bp.messages, + run_bp=False, + ) + np.testing.assert_allclose(reused, fresh, rtol=1e-10, atol=1e-12) + + +def test_open_rho_series_incrementally_reuses_contracted_terms(): + state = qtn.PEPS.rand( + 2, + 3, + bond_dim=2, + phys_dim=2, + seed=1911, + dtype="complex128", + ) + where = ((0, 0), (0, 2)) + bp = two_norm_bp( + state, + max_iterations=200, + tol=1e-10, + diis=False, + ) + cache = OpenLoopSeriesCache() + info = {} + partial_trace_open_loop_series_expand( + state, + where, + gloops=2, + messages=bp.messages, + run_bp=False, + cache=cache, + info=info, + ) + first_terms = dict(info["open_rho_terms"]) + first_base = info["open_rho_base_terms"] + + partial_trace_open_loop_series_expand( + state, + where, + gloops=4, + messages=bp.messages, + run_bp=False, + cache=cache, + info=info, + ) + + assert len(info["open_rho_terms"]) > len(first_terms) + assert info["open_rho_base_terms"] is first_base + for key, rho_term in first_terms.items(): + assert info["open_rho_terms"][key] is rho_term + + +def test_open_rho_series_sweep_reuses_bp_across_supports_and_cutoffs(): + state = qtn.PEPS.rand( + 1, + 4, + bond_dim=2, + phys_dim=2, + seed=1916, + dtype="complex128", + ) + supports = (((0, 0), (0, 3)), ((0, 0), (0, 1), (0, 3))) + result = partial_trace_open_loop_series_sweep( + state, + supports, + (0, 2, 3), + max_iterations=200, + tol=1e-10, + diis=False, + ) + + assert isinstance(result, OpenLoopSeriesSweepResult) + assert result.bp_converged + assert result.bp_iterations is not None + for support in supports: + assert tuple(support) in result.rhos + for cutoff in (0, 2, 3): + rho = result.get_rho(support, cutoff) + np.testing.assert_allclose(np.trace(rho), 1.0, atol=1e-12) + assert result.diagnostics[tuple(support)][cutoff]["term_count"] >= 0 + + +def test_open_rho_series_is_exact_for_a_tree_with_a_multi_site_support(): + state = qtn.PEPS.rand( + 1, + 4, + bond_dim=2, + phys_dim=2, + seed=1917, + dtype="complex128", + ) + where = ((0, 0), (0, 1), (0, 3)) + exact = state.partial_trace( + where, + max_bond=64, + optimize="auto-hq", + flatten=True, + normalized=True, + ) + info = {} + rho = partial_trace_open_loop_series_expand( + state, + where, + gloops=2, + max_iterations=200, + tol=1e-10, + diis=False, + info=info, + ) + + np.testing.assert_allclose(rho, exact, rtol=1e-10, atol=1e-12) + assert info["open_rho_excluded_edges"] diff --git a/tests/test_bp_symmray.py b/tests/test_bp_symmray.py index 6cd5ff4..793dad3 100644 --- a/tests/test_bp_symmray.py +++ b/tests/test_bp_symmray.py @@ -10,9 +10,11 @@ ctg = pytest.importorskip("cotengra") from pepsy.bp import ( # noqa: E402 + OpenLoopSeriesCache, compute_boundary_expectation, compute_bp_path_expectation, compute_local_expectation_edge_loop_series, + compute_local_expectation_open_loop_series, compute_local_expectation_loop_cluster, compute_local_expectation_loop_series, compute_path_cluster_expectation, @@ -21,6 +23,8 @@ loop_series_expand, one_norm_bp, partial_trace_edge_loop_series_expand, + partial_trace_open_loop_series_expand, + partial_trace_open_loop_series_sweep, partial_trace_loop_series_expand, partial_trace_loop_cluster_expand, partitioned_expand, @@ -33,6 +37,7 @@ SymPEPS, ps_to_peps, site_charge_alternating, + site_charge_from_map, ) @@ -500,6 +505,19 @@ def test_spinful_eta_pair_measurement_survives_su_and_bp_gauges(symmetry): assert exact_auto == pytest.approx(jw_value, rel=1e-10, abs=1e-10) assert exact_greedy == pytest.approx(exact_auto, rel=1e-10, abs=1e-10) + open_info = {} + open_value = compute_local_expectation_open_loop_series( + state.tn, + terms, + gloops=3, + max_iterations=200, + tol=1e-10, + diis=False, + info=open_info, + ) + assert open_value == pytest.approx(exact_auto, rel=1e-10, abs=1e-10) + assert open_info["open_scalar_fermionic_q_phase"] is True + su = gauge_all( state.tn, start="su", @@ -566,6 +584,123 @@ def test_spinful_eta_pair_measurement_survives_su_and_bp_gauges(symmetry): assert bp_helper == pytest.approx(exact_auto, rel=1e-10, abs=1e-10) +@pytest.mark.parametrize("symmetry", ("U1", "U1U1")) +def test_fermionic_long_range_doublon_product_survives_cluster_gauges(symmetry): + """A nontrivial long-range doublon correlator survives native clusters.""" + fermion = Fermion(spinful=True, symmetry=symmetry) + doublon = 2 if symmetry == "U1" else (1, 1) + empty = 0 if symmetry == "U1" else (0, 0) + occupations = { + (0, 0): doublon, + (1, 1): doublon, + (0, 1): empty, + (1, 0): empty, + } + peps = ps_to_peps( + 2, + 2, + fermion=fermion, + occupations=occupations, + dtype="complex128", + ) + state = SymPEPS( + peps=peps, + symmetry=symmetry, + edges=tuple(qtn.edges_2d_square(2, 2)), + fermionic=True, + phys_sectors=fermion.physical_sectors, + site_charge=occupations, + site_ind_id="k{},{}", + ) + + gate = fermion.hopping_gate(0.31, t=(1.0, 0.0)) + state.apply_gates( + ( + (gate, ((0, 0), (1, 0))), + (gate, ((1, 1), (0, 1))), + ), + method="direct", + contract="split", + max_bond=4, + cutoff=0.0, + ) + + operator = fermion.operator_term( + [(1.0, ((0, "double"), (1, "double")))], + sites=(0, 1), + ) + where = ((0, 0), (1, 1)) + terms = {where: operator} + norm_before = complex(state.tn.norm()) + + exact = state.tn.compute_local_expectation_exact( + terms, + normalized=True, + optimize="auto-hq", + ) + + su = gauge_all( + state.tn, + start="su", + target="su", + norm="2norm", + su_options={"max_iterations": 8, "tol": 0.0}, + ) + cluster = compute_path_cluster_expectation( + su.core, + terms, + gauges=su.gauges, + max_distance=1, + fillin=True, + max_bond=None, + normalized=True, + optimize="auto-hq", + ) + + bp_helper = compute_bp_path_expectation( + state.tn, + terms, + max_distance=1, + fillin=True, + max_bond=None, + normalized=True, + optimize="auto-hq", + bp_options={ + "run_opts": { + "max_iterations": 150, + "tol": 1e-10, + "diis": False, + } + }, + conversion_options={"smudge": 1e-12}, + ) + + with pytest.warns(UserWarning, match="not a compressed one"): + compressed = compute_path_cluster_expectation( + su.core, + terms, + gauges=su.gauges, + max_distance=1, + fillin=True, + max_bond=8, + normalized=True, + optimize="auto-hq", + ) + + expected_type = ( + "U1FermionicArray" if symmetry == "U1" else "U1U1FermionicArray" + ) + assert exact.real > 0.1 + assert cluster == pytest.approx(exact, rel=1e-10, abs=1e-10) + assert bp_helper == pytest.approx(exact, rel=1e-10, abs=1e-10) + assert compressed == pytest.approx(exact, rel=1e-10, abs=1e-10) + assert complex(state.tn.norm()) == pytest.approx(norm_before) + assert all( + type(tensor.data).__name__ == expected_type + for tensor in state.tn.tensors + ) + + def test_fermionic_long_range_boundary_expectation_matches_exact(): """The boundary route preserves a distant native density operator.""" state = SymPEPS.random( @@ -1267,6 +1402,114 @@ def test_fermionic_partial_trace_loop_series_keeps_native_rho(): ) +def test_fermionic_open_rho_is_native_diagnostic_and_scalar_path_is_graded(): + """Long-range open rho stays native while observables use graded gates.""" + state = SymPEPS.random( + 1, + 4, + symmetry="U1", + bond_dim=3, + phys_dim=2, + fermionic=True, + seed=1912, + dtype="complex128", + ) + where = ((0, 0), (0, 3)) + bp = two_norm_bp( + state.tn, + max_iterations=200, + tol=1e-10, + diis=False, + ) + info = {} + rho = partial_trace_open_loop_series_expand( + state.tn, + where, + gloops=3, + messages=bp.messages, + run_bp=False, + cache=OpenLoopSeriesCache(), + info=info, + ) + + assert type(rho).__name__ == "U1FermionicArray" + assert info["open_rho_family_counts"] == {"open_path": 1} + np.testing.assert_allclose( + np.trace(rho.to_dense()), + 1.0, + rtol=1e-10, + atol=1e-12, + ) + + fermion = Fermion(spinful=False, symmetry="U1") + local_where = ((0, 1),) + local_gate = fermion.chemical_potential_operator() + exact = state.tn.compute_local_expectation_exact( + {local_where: local_gate}, + normalized=True, + optimize="auto-hq", + ) + graded = compute_local_expectation_loop_series( + state.tn, + {local_where: local_gate}, + gloops=0, + messages=bp.messages, + run_bp=False, + max_iterations=200, + tol=1e-10, + diis=False, + ) + np.testing.assert_allclose(graded, exact, rtol=1e-10, atol=1e-12) + + +@pytest.mark.parametrize( + "symmetry,phys_dim,site_charge", + ( + ("U1", 2, None), + ("U1U1", 4, site_charge_alternating((1, 0), (0, 1))), + ("Z2", 2, None), + ), +) +def test_fermionic_open_rho_sweep_supports_native_symmetries( + symmetry, + phys_dim, + site_charge, +): + """One-BP open-rho sweeps preserve native U1, U1U1, and Z2 arrays.""" + state = SymPEPS.random( + 1, + 4, + symmetry=symmetry, + bond_dim=2, + phys_dim=phys_dim, + fermionic=True, + site_charge=site_charge, + seed=1920 + len(symmetry), + dtype="complex128", + ) + supports = (((0, 0), (0, 3)), ((0, 0), (0, 1), (0, 3))) + result = partial_trace_open_loop_series_sweep( + state.tn, + supports, + (0, 2, 3), + max_iterations=200, + tol=1e-10, + diis=False, + ) + + assert result.bp_converged + for support in supports: + for cutoff in (0, 2, 3): + rho = result.get_rho(support, cutoff) + assert type(rho).__module__.startswith("symmray") + np.testing.assert_allclose( + np.trace(rho.to_dense()), + 1.0, + rtol=1e-10, + atol=1e-12, + ) + + def test_fermionic_local_expectation_loop_series_aligns_charge_support(): """Graded gate insertion is exact on a fermionic D2BP tree.""" state = SymPEPS.random( @@ -1368,6 +1611,401 @@ def test_explicit_edge_loop_series_uses_fermion_safe_gate_path(): np.testing.assert_allclose(np.trace(rho.to_dense()), 1.0) +def test_open_scalar_series_uses_native_fermion_projectors_for_long_range_hopping(): + """Long-range native hopping stays graded while using one BP solve.""" + state = SymPEPS.random( + 1, + 4, + symmetry="U1", + bond_dim=3, + phys_dim=2, + fermionic=True, + seed=1931, + dtype="complex128", + ) + where = ((0, 0), (0, 3)) + gate = Fermion(spinful=False, symmetry="U1").hopping_operator() + exact = state.tn.compute_local_expectation_exact( + {where: gate}, normalized=True, optimize="auto-hq" + ) + bp = two_norm_bp( + state.tn, + max_iterations=200, + tol=1e-10, + diis=False, + ) + info = {} + value = compute_local_expectation_open_loop_series( + state.tn, + {where: gate}, + gloops=3, + messages=bp.messages, + run_bp=False, + max_iterations=200, + tol=1e-10, + diis=False, + info=info, + ) + + np.testing.assert_allclose(value, exact, rtol=1e-10, atol=1e-12) + assert info["open_scalar_native_route"] == "graded_open_projectors" + assert info["open_scalar_fermionic_q_phase"] is True + assert info["open_scalar_family_counts"] == {"open_path": 1} + + reverse_where = where[::-1] + reverse_gate = gate.transpose((1, 0, 3, 2)) + reverse_exact = state.tn.compute_local_expectation_exact( + {reverse_where: reverse_gate}, + normalized=True, + optimize="auto-hq", + ) + reverse_value = compute_local_expectation_open_loop_series( + state.tn, + {reverse_where: reverse_gate}, + gloops=3, + messages=bp.messages, + run_bp=False, + max_iterations=200, + tol=1e-10, + diis=False, + ) + np.testing.assert_allclose( + reverse_value, + reverse_exact, + rtol=1e-10, + atol=1e-12, + ) + + +@pytest.mark.parametrize( + "symmetry,phys_dim,site_charge,gate_name,phase", + ( + ("U1", 4, None, "density", False), + ("U1U1", 4, site_charge_alternating((1, 0), (0, 1)), "hopping", True), + ("Z2", 2, None, "density", False), + ), +) +def test_open_scalar_series_native_gate_phase_matches_symmetry( + symmetry, + phys_dim, + site_charge, + gate_name, + phase, +): + """Native open projectors handle diagonal and off-diagonal gates.""" + state = SymPEPS.random( + 1, + 4, + symmetry=symmetry, + bond_dim=3, + phys_dim=phys_dim, + fermionic=True, + site_charge=site_charge, + seed=1960 + len(symmetry), + dtype="complex128", + ) + fermion = Fermion( + spinful=phys_dim == 4, + symmetry=symmetry, + ) + gate = getattr(fermion, f"{gate_name}_operator")() + where = ((0, 0), (0, 3)) + exact = state.tn.compute_local_expectation_exact( + {where: gate}, + normalized=True, + optimize="auto-hq", + ) + info = {} + value = compute_local_expectation_open_loop_series( + state.tn, + {where: gate}, + gloops=3, + max_iterations=200, + tol=1e-10, + diis=False, + info=info, + ) + + np.testing.assert_allclose(value, exact, rtol=1e-10, atol=1e-12) + assert info["open_scalar_fermionic_q_phase"] is phase + + +def test_open_scalar_series_cyclic_native_u1_matches_exact_d3(): + """Cyclic native U1 open corrections remain exact at full cutoff.""" + site_charge = site_charge_from_map( + {(0, 0): 2, (1, 2): 2}, + default=0, + ) + state = SymPEPS.random( + 2, + 3, + symmetry="U1", + bond_dim=3, + phys_dim=4, + fermionic=True, + site_charge=site_charge, + seed=2053, + dtype="complex128", + ) + fermion = Fermion(spinful=True, symmetry="U1") + where = ((0, 0), (1, 2)) + gate = fermion.operator_term( + [(1.0, ((0, "double"), (1, "double")))], + sites=(0, 1), + ) + exact = state.tn.compute_local_expectation_exact( + {where: gate}, + normalized=True, + optimize="auto-hq", + ) + bp = two_norm_bp( + state.tn, + max_iterations=200, + tol=1e-9, + diis=False, + ) + info = {} + value = compute_local_expectation_open_loop_series( + state.tn, + {where: gate}, + gloops=7, + messages=bp.messages, + run_bp=False, + max_iterations=200, + tol=1e-9, + diis=False, + info=info, + ) + np.testing.assert_allclose(value, exact, rtol=1e-10, atol=1e-12) + assert info["open_scalar_native_route"] == "graded_cluster_compatible" + + rho_info = {} + rho = partial_trace_open_loop_series_expand( + state.tn, + where, + gloops=7, + messages=bp.messages, + run_bp=False, + max_iterations=200, + tol=1e-9, + diis=False, + info=rho_info, + ) + assert type(rho).__name__ == "U1FermionicArray" + np.testing.assert_allclose(np.trace(rho.to_dense()), 1.0, atol=1e-12) + assert rho_info["open_rho_native_route"] == "graded_cluster_compatible" + + +@pytest.mark.parametrize( + "symmetry,site_charge,seed", + ( + ( + "U1", + site_charge_from_map({(0, 0): 2, (1, 2): 2}, default=0), + 2083, + ), + ( + "U1", + site_charge_from_map({(0, 0): 2, (1, 2): 2}, default=0), + 2085, + ), + ( + "U1U1", + site_charge_alternating((1, 0), (0, 1)), + 2084, + ), + ( + "U1U1", + site_charge_alternating((1, 0), (0, 1)), + 2086, + ), + ), +) +def test_cyclic_native_open_series_gate_probes_and_rho_match_cluster( + symmetry, + site_charge, + seed, +): + """Cyclic native rho is cluster-consistent and gate probes are exact.""" + state = SymPEPS.random( + 2, + 3, + symmetry=symmetry, + bond_dim=3, + phys_dim=4, + fermionic=True, + site_charge=site_charge, + seed=seed, + dtype="complex128", + ) + where = ((0, 0), (1, 2)) + fermion = Fermion(spinful=True, symmetry=symmetry) + bp = two_norm_bp( + state.tn, + max_iterations=200, + tol=1e-9, + diis=False, + ) + + open_rho = partial_trace_open_loop_series_expand( + state.tn, + where, + gloops=7, + messages=bp.messages, + run_bp=False, + max_iterations=200, + tol=1e-9, + diis=False, + ) + cluster_rho = partial_trace_loop_cluster_expand( + state.tn, + where, + gloops=7, + messages=bp.messages, + run_bp=False, + max_iterations=200, + tol=1e-9, + diis=False, + ) + assert type(open_rho).__name__ == f"{symmetry}FermionicArray" + assert type(cluster_rho).__name__ == f"{symmetry}FermionicArray" + np.testing.assert_allclose( + open_rho.to_dense(), + cluster_rho.to_dense(), + rtol=1e-10, + atol=1e-12, + ) + np.testing.assert_allclose( + np.trace(open_rho.to_dense()), + 1.0, + rtol=1e-10, + atol=1e-12, + ) + + for gate in ( + fermion.density_operator(), + fermion.hopping_operator(), + fermion.eta_pair_operator(), + ): + exact = state.tn.compute_local_expectation_exact( + {where: gate}, + normalized=True, + optimize="auto-hq", + ) + value = compute_local_expectation_open_loop_series( + state.tn, + {where: gate}, + gloops=7, + messages=bp.messages, + run_bp=False, + max_iterations=200, + tol=1e-9, + diis=False, + ) + np.testing.assert_allclose(value, exact, rtol=1e-10, atol=1e-12) + + reverse_where = where[::-1] + reverse_gate = gate.transpose((1, 0, 3, 2)) + reverse_exact = state.tn.compute_local_expectation_exact( + {reverse_where: reverse_gate}, + normalized=True, + optimize="auto-hq", + ) + reverse_value = compute_local_expectation_open_loop_series( + state.tn, + {reverse_where: reverse_gate}, + gloops=7, + messages=bp.messages, + run_bp=False, + max_iterations=200, + tol=1e-9, + diis=False, + ) + np.testing.assert_allclose( + reverse_value, + reverse_exact, + rtol=1e-10, + atol=1e-12, + ) + + +def test_cyclic_native_open_series_honors_contraction_cost_limits(): + """Safe cyclic fallback still records and applies Cotengra budgets.""" + site_charge = site_charge_from_map( + {(0, 0): 2, (1, 2): 2}, + default=0, + ) + state = SymPEPS.random( + 2, + 3, + symmetry="U1", + bond_dim=3, + phys_dim=4, + fermionic=True, + site_charge=site_charge, + seed=2093, + dtype="complex128", + ) + where = ((0, 0), (1, 2)) + gate = Fermion(spinful=True, symmetry="U1").operator_term( + [(1.0, ((0, "double"), (1, "double")))], + sites=(0, 1), + ) + exact = state.tn.compute_local_expectation_exact( + {where: gate}, + normalized=True, + optimize="auto-hq", + ) + info = {} + value = compute_local_expectation_open_loop_series( + state.tn, + {where: gate}, + gloops=7, + max_iterations=200, + tol=1e-9, + diis=False, + max_flops_log10=11.0, + max_peak_memory_log2=30.0, + info=info, + ) + np.testing.assert_allclose(value, exact, rtol=1e-10, atol=1e-12) + assert info["open_scalar_native_route"] == "graded_cluster_compatible" + assert not info["open_scalar_edge_term_costs"] + assert not info["open_scalar_edge_skipped_terms"] + assert info["open_scalar_cluster_region_costs"] + assert info["open_scalar_cluster_region_skipped_terms"] == {} + assert info["open_scalar_term_costs"] + for costs in info["open_scalar_term_costs"].values(): + for cost in costs.values(): + if cost is not None: + assert cost["flops_log10"] <= 11.0 + assert cost["peak_memory_log2"] <= 30.0 + + rho_info = {} + rho = partial_trace_open_loop_series_expand( + state.tn, + where, + gloops=7, + max_iterations=200, + tol=1e-9, + diis=False, + max_flops_log10=11.0, + max_peak_memory_log2=30.0, + info=rho_info, + ) + assert type(rho).__name__ == "U1FermionicArray" + assert rho_info["open_rho_native_route"] == "graded_cluster_compatible" + assert not rho_info["open_rho_edge_term_costs"] + assert not rho_info["open_rho_edge_skipped_terms"] + assert rho_info["open_rho_cluster_region_costs"] + assert rho_info["open_rho_cluster_region_skipped_terms"] == {} + assert rho_info["open_rho_term_costs"] + np.testing.assert_allclose(np.trace(rho.to_dense()), 1.0, atol=1e-12) + for cost in rho_info["open_rho_term_costs"].values(): + assert cost["flops_log10"] <= 11.0 + assert cost["peak_memory_log2"] <= 30.0 + + def test_explicit_edge_loop_series_preserves_dense_edge_degree_terms(): """Edge-degree terms are distinct from the local-region cutoff API.""" state = qtn.PEPS.rand(2, 2, bond_dim=2, seed=1906, dtype="complex128") From e3630a198a37c3ad1cb7c4ef0277dd10c231bf5a Mon Sep 17 00:00:00 2001 From: rezaquant Date: Fri, 31 Jul 2026 08:24:01 -0700 Subject: [PATCH 37/70] tree: support ternary virtual binary roots --- docs/api/optimizers/tree.md | 26 +++-- src/pepsy/optimizers/tree/layout.py | 144 ++++++++++++++++++++++--- src/pepsy/optimizers/tree/optimizer.py | 56 +++++++++- src/pepsy/optimizers/tree/ttn.py | 28 ++++- src/pepsy/tensors/constructors.py | 13 ++- tests/test_optimize_tree.py | 54 ++++++++++ 6 files changed, 289 insertions(+), 32 deletions(-) diff --git a/docs/api/optimizers/tree.md b/docs/api/optimizers/tree.md index 7820fac..5331635 100644 --- a/docs/api/optimizers/tree.md +++ b/docs/api/optimizers/tree.md @@ -47,6 +47,17 @@ structured sub-MPOs may include it in their support. `TreeLayoutFinder` keeps the site fixed at the root while its path, Steiner, congestion, greedy, and Nevergrad objectives permute only the remaining leaf sites. +For the conventional binary TTN with a three-leg top tensor, pass +`max_arity=2, top_arity=3` to `TreePlan.from_order`, `TreeLayoutFinder`, or +`TreeTensorNetwork.from_order`. The structural root then has three **virtual** +child bonds and no parent bond; every non-root internal tensor has two child +bonds and one parent bond. Thus the root is still in the rank-three binary +class, rather than being a genuinely wider tensor. `top_arity=3` is not +combined with `root_qubit`, because adding a physical root leg would make a +rank-four tensor. `TreePlan.is_binary()` accepts this ternary-root convention, +while `TreePlan.is_strictly_binary()` requests two children at every internal +node. + Gates are absorbed into the tree: - **single-qubit gates** are contracted into their site tensor with no bond @@ -305,7 +316,8 @@ Because the geometry (`plan`) and naming live in `_EXTRA_PROPS`, they survive `.copy()` and every Quimb view, exactly like `site_ind_id` does for an MPS. Build one with `TreeTensorNetwork.from_plan(plan)` (product `|0...0>`), `TreeTensorNetwork.from_order(order, structure=...)` (build the plan and the -product state in one step), or `TreeTensorNetwork.rand(plan, D=..., seed=...)` +product state in one step; its `top_arity=3` option exposes the ternary virtual +root), or `TreeTensorNetwork.rand(plan, D=..., seed=...)` (a random state, canonicalised around the root by default). `TreeOptimizer` builds and evolves its state on this class, delegating all node/qubit naming and geometry queries to it. @@ -355,8 +367,8 @@ graded product tree is normalized by an exact graded norm contraction, so its represented norm is one rather than an arbitrary constructor scalar. `pepsy.hrs_to_ttn(..., chi=...)` creates the corresponding random symmetric tree with the requested charge-sector bond dimension and accepts the same -`root_qubit=` option. These constructors keep the Symmray arrays native; they -do not materialize dense tensor data. +`root_qubit=`, `max_arity=`, and `top_arity=` options. These constructors keep +the Symmray arrays native; they do not materialize dense tensor data. `pepsy.TreeSampler(state)` samples every registered physical site, including the optional root site. Its cached canonical arrays use parent, physical, then @@ -468,12 +480,14 @@ any arity, controlled by two knobs on `TreeLayoutFinder` / `TreePlan.from_order` A caller may bypass the finder entirely by passing an explicit `TreePlan` via `TreeOptimizer(..., tree=plan)`. `TreePlan` is exported from both `pepsy` and `pepsy.optimizers.tree`. Build one with -`TreePlan.from_order(order, weights=..., structure=..., max_arity=...)`, or -- for +`TreePlan.from_order(order, weights=..., structure=..., max_arity=..., top_arity=...)`, or -- for a fully hand-specified arbitrary-arity tree -- with `TreePlan.from_children(children, qubit_of_leaf)`, which validates that the children map and leaf assignment describe a single rooted tree covering qubits -`0..n-1` exactly once. `TreePlan.max_arity()` and `TreePlan.is_binary()` report -the shape. +`0..n-1` exactly once. Set `top_arity=3` with `max_arity=2` for the +three-virtual-bond root convention described above. `TreePlan.max_arity()` and +`TreePlan.is_binary()` report the shape; `TreePlan.is_strictly_binary()` is the +strict two-child-at-every-internal-node predicate. For an automatic arity choice, call `finder.recommend_arities((2, 3, 4))`. This is also what the finder and diff --git a/src/pepsy/optimizers/tree/layout.py b/src/pepsy/optimizers/tree/layout.py index 7b8b7fa..4e4a30c 100644 --- a/src/pepsy/optimizers/tree/layout.py +++ b/src/pepsy/optimizers/tree/layout.py @@ -486,7 +486,7 @@ def __init__( @classmethod def from_order(cls, order, *, weights=None, structure="quality", max_arity=2, community_frac=0.35, star_frac=0.75, - dense_max=512, root_qubit=None): + dense_max=512, root_qubit=None, top_arity=None): """Build a rooted tree by recursive partition of ``order``. Parameters @@ -525,6 +525,14 @@ def from_order(cls, order, *, weights=None, structure="quality", Maximum subsystem size for dense spectral reordering. root_qubit : int, optional Qubit label carried by the top tensor rather than a leaf. + top_arity : int, optional + Number of virtual child bonds on the structural root. Set + ``top_arity=3`` with ``max_arity=2`` for the conventional binary + TTN with a ternary top tensor: the root has three virtual legs and + every non-root internal tensor has two child legs plus one parent + leg. This keeps every tensor rank at most three. It is incompatible + with ``root_qubit`` when greater than two because that would make a + rank-four root tensor. """ order = list(order) if not order and root_qubit is None: @@ -538,6 +546,24 @@ def from_order(cls, order, *, weights=None, structure="quality", root_qubit = int(root_qubit) except (TypeError, ValueError) as exc: raise ValueError("root_qubit must be an integer or None.") from exc + if top_arity is not None: + try: + top_arity = int(top_arity) + except (TypeError, ValueError) as exc: + raise ValueError( + "top_arity must be an integer >= 2 or None." + ) from exc + if top_arity < 2: + raise ValueError("top_arity must be >= 2 or None.") + if top_arity > len(order): + raise ValueError( + "top_arity cannot exceed the number of non-root qubits." + ) + if root_qubit is not None and top_arity != 2: + raise ValueError( + "top_arity > 2 cannot be combined with root_qubit: " + "the root would have a rank-four tensor." + ) all_qubits = order + ([] if root_qubit is None else [root_qubit]) if sorted(all_qubits) != list(range(len(all_qubits))): raise ValueError( @@ -585,14 +611,15 @@ def make_internal(child_ids): parent[c] = nid return nid - def kary_split(qs): - """Split ``qs`` into up to ``max_arity`` contiguous balanced parts. + def kary_split(qs, arity=None): + """Split ``qs`` into up to ``arity`` contiguous balanced parts. Cut points use ``floor(i * L / k)`` so the two-way case reproduces the previous ``mid = len(qs) // 2`` bisection exactly. """ length = len(qs) - k = length if max_arity is None else max_arity + k = max_arity if arity is None else arity + k = length if k is None else k k = min(k, length) if k <= 1: return [qs] @@ -654,18 +681,19 @@ def is_near_clique(qs): total = m * (m - 1) // 2 return total > 0 and strong / total >= float(star_frac) - def split(qs): + def split(qs, arity=None): """Return the child qubit-groups for the internal node over ``qs``.""" + arity_limit = max_arity if arity is None else arity groups = None - if structure == "adaptive": + if structure == "adaptive" and arity is None: comps = communities(qs) if comps is not None and len(comps) >= 2: - if max_arity is None or len(comps) <= max_arity: + if arity_limit is None or len(comps) <= arity_limit: groups = comps # else: too many communities for the arity cap; fall back to # a spectral k-ary split (deeper recursion still resolves # communities inside each part). - elif (max_arity is None or len(qs) <= max_arity) \ + elif (arity_limit is None or len(qs) <= arity_limit) \ and is_near_clique(qs): # A densely coupled block is flattest as a star of leaves. groups = [[q] for q in qs] @@ -677,14 +705,15 @@ def split(qs): ) if spectral: qs2 = spectral - groups = kary_split(qs2) + groups = kary_split(qs2, arity_limit) return groups - def build(qs): + def build(qs, *, is_root=False): qs = list(qs) if len(qs) == 1: return make_leaf(qs[0]) - groups = split(qs) + root_limit = top_arity if is_root else None + groups = split(qs, root_limit) if len(groups) < 2: # Degenerate split (e.g. all mass in one part): force a split so # recursion always makes progress. @@ -694,7 +723,7 @@ def build(qs): return make_internal(child_ids) if order: - root = build(order) + root = build(order, is_root=True) if root_qubit is not None and root in qubit_of_leaf: # With one non-root qubit, ``build`` returns that physical # leaf itself. The top qubit needs its own tensor, so insert a @@ -984,10 +1013,61 @@ def max_arity(self): """Return the largest number of children over all internal nodes.""" return max((len(ch) for ch in self.children.values()), default=0) - def is_binary(self): + @property + def top_arity(self): + """Return the number of virtual child bonds on the structural root.""" + return len(self.children.get(self.root, ())) + + def virtual_degree(self, nid): + """Return the number of virtual tree bonds incident on ``nid``.""" + if nid not in self.children: + raise ValueError(f"node {nid!r} is not present in the tree") + return len(self.children[nid]) + int(nid in self.parent) + + def max_virtual_degree(self): + """Return the largest number of virtual bonds on any tensor.""" + return max( + (self.virtual_degree(nid) for nid in self.children), + default=0, + ) + + def max_tensor_rank(self): + """Return the largest number of virtual/physical legs on a node.""" + return max( + ( + self.virtual_degree(nid) + + int(nid in self.qubit_of_node) + for nid in self.children + ), + default=0, + ) + + def is_strictly_binary(self): """Return ``True`` when every internal node has exactly two children.""" return all(len(ch) in (0, 2) for ch in self.children.values()) + def is_binary(self, *, allow_ternary_root=True): + """Return whether the tree is binary below an optional ternary root. + + A conventional binary TTN has two child bonds entering every + non-root internal tensor and one parent bond leaving it. Its top tensor + has no parent, so it may carry three child bonds without increasing + the maximum tensor rank. Pass ``allow_ternary_root=False`` to request + the older strictly-binary predicate. + """ + if not allow_ternary_root: + return self.is_strictly_binary() + for nid, children in self.children.items(): + if not children: + continue + allowed = (2, 3) if nid == self.root else (2,) + if len(children) not in allowed: + return False + # A ternary virtual root is binary only when it has no additional + # physical leg. This also keeps explicit hand-built rank-four roots + # out of the binary predicate. + return self.max_tensor_rank() <= 3 + def max_bond_cut(self): """Return the largest qubit bipartition induced by any tree bond. @@ -1155,7 +1235,8 @@ def __repr__(self): return ( f"TreePlan(n={self.n}, root={self.root}, " f"internal_nodes={n_internal}, " - f"max_arity={self.max_arity()}{root_site})" + f"max_arity={self.max_arity()}, top_arity={self.top_arity}" + f"{root_site})" ) @@ -1188,6 +1269,12 @@ class TreeLayoutFinder: wider trees). An iterable of candidate arities makes :meth:`run` *search* them and keep the objective-best plan; this is the default ``(2, 3, 4)``. Pass a scalar to opt back into a single fixed tree. + top_arity : int, optional + Override the structural root's number of virtual child bonds. With + ``max_arity=2, top_arity=3`` the finder builds the conventional binary + TTN whose top tensor has three virtual legs while all non-root internal + tensors remain two-in/one-out. This keeps the maximum tensor rank at + three. It cannot be combined with ``root_qubit`` when greater than two. chi : int, optional Bond-dimension budget used to bias the default arity search toward plans that stay exact at ``chi`` (see :meth:`recommend_arities`). ``None`` @@ -1254,7 +1341,8 @@ class TreeLayoutFinder: """ def __init__(self, gates=None, n=None, *, supports=None, structure="quality", - max_arity=(2, 3, 4), community_frac=0.35, star_frac=0.75, + max_arity=(2, 3, 4), top_arity=None, + community_frac=0.35, star_frac=0.75, dense_max=512, objective="path", weight_mode="count", chi=None, max_operator_qubits=8, hybrid_weights=None, refine=None, refine_budget=None, topology_refine=None, topology_budget=None, @@ -1330,6 +1418,25 @@ def __init__(self, gates=None, n=None, *, supports=None, structure="quality", self.leaf_qubits = tuple( q for q in range(self.n) if q != self.root_qubit ) + if top_arity is not None: + try: + top_arity = int(top_arity) + except (TypeError, ValueError) as exc: + raise ValueError( + "top_arity must be an integer >= 2 or None." + ) from exc + if top_arity < 2: + raise ValueError("top_arity must be >= 2 or None.") + if top_arity > len(self.leaf_qubits): + raise ValueError( + "top_arity cannot exceed the number of non-root qubits." + ) + if root_qubit is not None and top_arity != 2: + raise ValueError( + "top_arity > 2 cannot be combined with root_qubit: " + "the root would have a rank-four tensor." + ) + self.top_arity = top_arity self.supports = tuple(normalized_supports) self.structure = structure self.max_arity, self.arity_candidates = _normalize_arity_candidates( @@ -1720,7 +1827,7 @@ def _tensor_cost_key(self, plan): for node, children in plan.children.items(): if not children: continue - virtual_degree = len(children) + (1 if node in plan.parent else 0) + virtual_degree = plan.virtual_degree(node) physical_legs = 1 if node in plan.qubit_of_node else 0 degrees.append(virtual_degree) log_sizes.append(virtual_degree * log_chi + physical_legs) @@ -2095,6 +2202,7 @@ def _build_plan(self, weights, *, structure=None, star_frac=self.star_frac, dense_max=self.dense_max, root_qubit=self.root_qubit, + top_arity=self.top_arity, ) self._plan_cache[key] = (weights, plan) return plan @@ -2976,6 +3084,7 @@ def _balanced_plan(self): self.leaf_qubits, structure="balanced", root_qubit=self.root_qubit, + top_arity=self.top_arity, ) return self._balanced_plan_cache @@ -3059,8 +3168,11 @@ def report(self, plan=None, *, include_edge_loads=True): ), "root": plan.root, "root_qubit": plan.root_qubit, + "top_arity": plan.top_arity, "is_binary": plan.is_binary(), + "is_strictly_binary": plan.is_strictly_binary(), "max_arity": plan.max_arity(), + "max_tensor_rank": plan.max_tensor_rank(), "arity_histogram": arity_histogram, "score": float(weighted_sum), "max_path": int(max(dists)) if dists else 0, diff --git a/src/pepsy/optimizers/tree/optimizer.py b/src/pepsy/optimizers/tree/optimizer.py index 4cff216..40da1dd 100644 --- a/src/pepsy/optimizers/tree/optimizer.py +++ b/src/pepsy/optimizers/tree/optimizer.py @@ -326,6 +326,11 @@ class TreeOptimizer: optimizer's ``chi`` (structures that stay exact at ``chi`` are preferred). Pass ``max_arity=2`` to force a fixed binary tree. Ignored when an explicit ``tree`` is supplied. + top_arity : int, optional + Number of virtual child bonds on the structural root when the layout + is built automatically. Set ``top_arity=3`` with ``max_arity=2`` for + the conventional binary TTN with a three-leg top tensor. This keeps + every tensor rank at most three. layout_objective : {"path", "congestion", "compression", "hypergraph", "hybrid"} Objective used when building an automatic tree. ``"path"`` is the backward-compatible interaction-path heuristic; ``"congestion"`` @@ -442,7 +447,8 @@ def __init__(self, gates=None, n=None, *, chi=64, cutoff=_DEFAULT_CUTOFF, cutoff_mode=_DEFAULT_CUTOFF_MODE, mode="auto", two_site_mode=None, - structure="quality", max_arity=(2, 3, 4), community_frac=0.35, + structure="quality", max_arity=(2, 3, 4), top_arity=None, + community_frac=0.35, star_frac=0.75, layout_objective="path", layout_weight_mode="count", layout_time_decay=None, layout_time_window=None, layout=None, tree=None, @@ -459,6 +465,9 @@ def __init__(self, gates=None, n=None, *, chi=64, # iterator and silently degrade to an interaction-free layout. if hasattr(gates, "__next__"): gates = list(gates) + layout_top_arity = ( + layout.top_arity if isinstance(layout, TreeLayoutFinder) else None + ) if layout is not None: if tree is not None: raise ValueError("pass either layout= or tree=, not both.") @@ -474,6 +483,8 @@ def __init__(self, gates=None, n=None, *, chi=64, "layout must be a TreeLayoutFinder or TreePlan; " "pass an entangled TreeTensorNetwork as state= or tn=." ) + if top_arity is None and layout_top_arity is not None: + top_arity = layout_top_arity if state is not None: if tn is not None: raise ValueError("pass either state= or tn=, not both.") @@ -490,6 +501,20 @@ def __init__(self, gates=None, n=None, *, chi=64, raise ValueError( "root_qubit must be an integer or None." ) from exc + if top_arity is not None: + try: + top_arity = int(top_arity) + except (TypeError, ValueError) as exc: + raise ValueError( + "top_arity must be an integer >= 2 or None." + ) from exc + if top_arity < 2: + raise ValueError("top_arity must be >= 2 or None.") + if root_qubit is not None and top_arity != 2: + raise ValueError( + "top_arity > 2 cannot be combined with root_qubit: " + "the root would have a rank-four tensor." + ) self.G, self.where, self.event_types = self._normalize_gate_queue(gates) self.layout_finder = layout if isinstance(layout, TreeLayoutFinder) else None @@ -590,6 +615,7 @@ def __init__(self, gates=None, n=None, *, chi=64, # iterable of candidate arities to search; forward it to the finder, # which normalizes and (for a candidate set) searches it chi-aware. self.max_arity = max_arity + self.top_arity = top_arity self.community_frac = float(community_frac) self.star_frac = float(star_frac) self.layout_objective = str(layout_objective) @@ -632,6 +658,7 @@ def __init__(self, gates=None, n=None, *, chi=64, gates=self._layout_gate_stream(), n=self.n, structure=structure, max_arity=self.max_arity, community_frac=self.community_frac, star_frac=self.star_frac, + top_arity=self.top_arity, objective=self.layout_objective, weight_mode=self.layout_weight_mode, time_decay=self.layout_time_decay, @@ -651,6 +678,14 @@ def __init__(self, gates=None, n=None, *, chi=64, raise ValueError( "root_qubit does not match the supplied tree/layout plan." ) + if self.top_arity is not None and tree.top_arity != int(self.top_arity): + raise ValueError( + "top_arity does not match the supplied tree/layout plan." + ) + if self.top_arity is None and tree.top_arity >= 2: + # Preserve the root convention when an explicit plan is handed in, + # so later candidate searches and plots keep the same topology. + self.top_arity = tree.top_arity self.plan = tree if product_state_source is not None: @@ -1435,8 +1470,11 @@ def layout_report(self): "n_qubits": self.n, "root": self.plan.root, "root_qubit": self.plan.root_qubit, + "top_arity": self.plan.top_arity, "is_binary": self.plan.is_binary(), + "is_strictly_binary": self.plan.is_strictly_binary(), "max_arity": self.plan.max_arity(), + "max_tensor_rank": self.plan.max_tensor_rank(), } return self.layout_finder.report(self.plan) @@ -1491,6 +1529,7 @@ def select_layout_for_compression( chi=self.chi, max_operator_qubits=self.max_operator_qubits, root_qubit=self.plan.root_qubit, + top_arity=self.top_arity, ) candidates = finder.candidate_plans( chi=self.chi, @@ -1534,6 +1573,7 @@ def select_layout_for_compression( mode=self.mode, structure=self.structure, max_arity=self.max_arity, + top_arity=self.top_arity, community_frac=self.community_frac, star_frac=self.star_frac, tree=plan, @@ -1632,6 +1672,7 @@ def plot_layout(self, plan=None, *, layout_kwargs=None, **plot_kwargs): chi=self.chi, max_operator_qubits=self.max_operator_qubits, root_qubit=self.plan.root_qubit, + top_arity=self.top_arity, ) if plan is None: if layout_kwargs: @@ -1658,6 +1699,7 @@ def plot_rubberband(self, plan=None, *, layout_kwargs=None, **plot_kwargs): chi=self.chi, max_operator_qubits=self.max_operator_qubits, root_qubit=self.plan.root_qubit, + top_arity=self.top_arity, ) if plan is None: if layout_kwargs: @@ -1684,6 +1726,7 @@ def plot_tent(self, plan=None, *, layout_kwargs=None, **plot_kwargs): chi=self.chi, max_operator_qubits=self.max_operator_qubits, root_qubit=self.plan.root_qubit, + top_arity=self.top_arity, ) if plan is None: if layout_kwargs: @@ -4576,6 +4619,7 @@ def copy(self): mode=self.mode, structure=self.structure, max_arity=self.max_arity, + top_arity=self.top_arity, community_frac=self.community_frac, star_frac=self.star_frac, tree=self.plan, @@ -4648,13 +4692,14 @@ def find_tree_layout(cls, gates, n=None, *, structure="quality", star_frac=0.75, layout_objective="path", layout_weight_mode="count", layout_time_decay=None, layout_time_window=None, - root_qubit=None, + root_qubit=None, top_arity=None, max_operator_qubits=_DEFAULT_MAX_OPERATOR_QUBITS): """Return the :class:`TreePlan` a :class:`TreeLayoutFinder` would use.""" return TreeLayoutFinder( gates=gates, n=n, structure=structure, max_arity=max_arity, community_frac=community_frac, star_frac=star_frac, objective=layout_objective, + top_arity=top_arity, weight_mode=layout_weight_mode, time_decay=layout_time_decay, time_window=layout_time_window, @@ -4666,7 +4711,8 @@ def find_tree_layout(cls, gates, n=None, *, structure="quality", def convergence_sweep(cls, gates, n=None, chi_values=(2, 4, 8, 16, 32), *, ops=None, structure="quality", max_arity=(2, 3, 4), community_frac=0.35, star_frac=0.75, tree=None, - root_qubit=None, dense_cap=1 << 14): + root_qubit=None, top_arity=None, + dense_cap=1 << 14): """Replay ``gates`` at several ``chi`` and report convergence. The tree structure is built once and reused for every ``chi`` so the @@ -4702,8 +4748,8 @@ def convergence_sweep(cls, gates, n=None, chi_values=(2, 4, 8, 16, 32), *, chi_values = sorted(int(c) for c in chi_values) if tree is None: probe = cls(gates, n=n, structure=structure, max_arity=max_arity, - community_frac=community_frac, star_frac=star_frac, - root_qubit=root_qubit, run=False) + top_arity=top_arity, community_frac=community_frac, + star_frac=star_frac, root_qubit=root_qubit, run=False) tree = probe.plan n = probe.n elif n is None: diff --git a/src/pepsy/optimizers/tree/ttn.py b/src/pepsy/optimizers/tree/ttn.py index 4ddce86..ef2c34a 100644 --- a/src/pepsy/optimizers/tree/ttn.py +++ b/src/pepsy/optimizers/tree/ttn.py @@ -19,6 +19,9 @@ ``site_tag_id.format(q)`` (default ``"I{}"``) and the physical index ``site_ind_id.format(q)`` (default ``"k{}"``) for qubit ``q``; these are structural leaves by default; +* a plan may use ``top_arity=3`` for the conventional binary TTN with three + virtual bonds entering the top tensor. Other internal nodes then have two + child bonds plus one parent bond, so every tensor remains rank three; * a plan may designate one additional ``root_qubit`` carried by the top tensor. A binary root then has exactly two child bonds plus this physical leg. Other internal nodes remain ancillary bond carriers. This class supplies the @@ -313,6 +316,25 @@ def plan(self): """The :class:`TreePlan` describing the tree structure.""" return self._plan + @property + def top_arity(self): + """Number of virtual child bonds entering the structural root.""" + return self._plan.top_arity + + @property + def max_virtual_degree(self): + """Largest number of virtual bonds incident on any live tensor.""" + return self._plan.max_virtual_degree() + + @property + def max_tensor_rank(self): + """Largest virtual/physical leg count in the live tree.""" + return self._plan.max_tensor_rank() + + def is_binary(self, *, allow_ternary_root=True): + """Whether the TTN is binary below an optional ternary top tensor.""" + return self._plan.is_binary(allow_ternary_root=allow_ternary_root) + @property def node_tag_id(self): """Format string for the structural node tag (e.g. ``"N{}"``).""" @@ -1918,18 +1940,22 @@ def zero_charge(value): def from_order(cls, order, *, weights=None, structure="quality", max_arity=2, community_frac=0.35, star_frac=0.75, dtype=complex, site_tag_id="I{}", site_ind_id="k{}", - node_tag_id="N{}", root_qubit=None): + node_tag_id="N{}", root_qubit=None, top_arity=None): """Build a product state on a tree partitioned from ``order``. Convenience wrapper that first builds a :class:`TreePlan` with :meth:`TreePlan.from_order` and then :meth:`from_plan`. ``max_arity`` and ``structure`` control the tree shape (see :meth:`TreePlan.from_order`); the defaults reproduce the binary tree. + Set ``top_arity=3`` with ``max_arity=2`` for a binary TTN with a + three-virtual-leg top tensor; all lower internal tensors remain rank + three (two child bonds plus one parent bond). """ plan = TreePlan.from_order( order, weights=weights, structure=structure, max_arity=max_arity, community_frac=community_frac, star_frac=star_frac, root_qubit=root_qubit, + top_arity=top_arity, ) return cls.from_plan( plan, diff --git a/src/pepsy/tensors/constructors.py b/src/pepsy/tensors/constructors.py index e148cac..489b68f 100644 --- a/src/pepsy/tensors/constructors.py +++ b/src/pepsy/tensors/constructors.py @@ -619,6 +619,7 @@ def ps_to_ttn( root_qubit=None, structure="balanced", max_arity=2, + top_arity=None, community_frac=0.35, star_frac=0.75, chi: int = 1, @@ -657,7 +658,7 @@ def ps_to_ttn( root_qubit : int, optional Qubit carried by the top tensor rather than a structural leaf. When an explicit ``tree`` is supplied, this must match its root site. - structure, max_arity, community_frac, star_frac + structure, max_arity, top_arity, community_frac, star_frac Forwarded to :meth:`TreePlan.from_order`. chi : int, optional If greater than one, expand every virtual bond to at least ``chi``. @@ -718,6 +719,7 @@ def ps_to_ttn( order, structure=structure, max_arity=max_arity, + top_arity=top_arity, community_frac=community_frac, star_frac=star_frac, root_qubit=root_qubit, @@ -845,6 +847,7 @@ def hrs_to_ttn( root_qubit=None, structure="balanced", max_arity=2, + top_arity=None, community_frac=0.35, star_frac=0.75, seed=None, @@ -862,9 +865,10 @@ def hrs_to_ttn( With ``fermion=`` the physical sites receive the model's charge sectors, while virtual-only internal nodes are neutral and every virtual tree edge is a conjugate pair of Symmray charge-sector indices. ``root_qubit`` places - one physical site on the top tensor. ``chi`` is the requested total - virtual-bond dimension. All block-sparse and fermionic operations are - delegated to Symmray/Quimb. + one physical site on the top tensor. ``top_arity=3`` with ``max_arity=2`` + gives the conventional three-virtual-leg binary root. ``chi`` is the + requested total virtual-bond dimension. All block-sparse and fermionic + operations are delegated to Symmray/Quimb. """ from ..optimizers.tree import TreePlan, TreeTensorNetwork @@ -895,6 +899,7 @@ def hrs_to_ttn( order, structure=structure, max_arity=max_arity, + top_arity=top_arity, community_frac=community_frac, star_frac=star_frac, root_qubit=root_qubit, diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index c0f15fa..3321517 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -2992,6 +2992,60 @@ def test_ttn_from_plan_is_product_state(): assert np.linalg.norm(sv[1:]) < 1e-12 +def test_binary_tree_supports_a_three_virtual_leg_top_tensor(): + """A ternary virtual root keeps every tensor in the binary rank class.""" + plan = TreePlan.from_order( + range(9), structure="balanced", max_arity=2, top_arity=3, + ) + assert plan.top_arity == 3 + assert plan.is_binary() + assert not plan.is_strictly_binary() + assert all( + len(children) in (0, 2) + for node, children in plan.children.items() + if node != plan.root + ) + + ttn = TreeTensorNetwork.from_plan(plan) + assert len(ttn.node_tensor(plan.root).inds) == 3 + assert ttn.max_virtual_degree == 3 + assert ttn.max_tensor_rank == 3 + assert ttn.validate(check_canonical=True) is ttn + + ordered = TreeTensorNetwork.from_order( + range(9), max_arity=2, top_arity=3, + ) + assert ordered.top_arity == 3 + assert len(ordered.node_tensor(ordered.plan.root).inds) == 3 + + finder = TreeLayoutFinder([], n=9, max_arity=2, top_arity=3) + found = finder.run() + report = finder.report(found) + assert found.top_arity == 3 + assert found.is_binary() + assert report["top_arity"] == 3 + assert report["max_tensor_rank"] == 3 + + automatic = TreeOptimizer( + [], n=9, max_arity=2, top_arity=3, run=False, + ) + assert automatic.plan.top_arity == 3 + assert automatic.tn.max_tensor_rank == 3 + + layout = TreeLayoutFinder([], n=9, max_arity=2, top_arity=3) + from_layout = TreeOptimizer([], layout=layout, run=False) + assert from_layout.top_arity == 3 + assert from_layout.plan.top_arity == 3 + + product = pepsy.ps_to_ttn(9, max_arity=2, top_arity=3) + assert product.top_arity == 3 + assert product.max_tensor_rank == 3 + + random = pepsy.hrs_to_ttn(9, max_arity=2, top_arity=3, seed=11) + assert random.top_arity == 3 + assert random.max_tensor_rank == 3 + + def test_ps_to_ttn_matches_product_state_constructor_api(): """The high-level TTN constructor mirrors ``ps_to_mps`` amplitudes.""" theta = 0.31 From 6ad21cb4dda7abd4e4847f4953c4d63143ca01bf Mon Sep 17 00:00:00 2001 From: rezaquant Date: Fri, 31 Jul 2026 09:01:44 -0700 Subject: [PATCH 38/70] tree: optimize layouts across all scales --- docs/api/optimizers/tree.md | 11 + src/pepsy/optimizers/tree/layout.py | 668 +++++++++++++++++++++++-- src/pepsy/optimizers/tree/optimizer.py | 4 +- tests/test_optimize_tree.py | 65 +++ 4 files changed, 713 insertions(+), 35 deletions(-) diff --git a/docs/api/optimizers/tree.md b/docs/api/optimizers/tree.md index 5331635..02eb779 100644 --- a/docs/api/optimizers/tree.md +++ b/docs/api/optimizers/tree.md @@ -421,6 +421,15 @@ NNI topology moves using the full hyperedge score. Pass set explicit budgets for a larger search. Dense operators wider than `max_operator_qubits` still use the documented conservative rank bound. +For a whole-tree optimization, use `objective="full_tree"` (also accepted as +`"tree"` or `"cotengra"`). This evaluates dynamic operator-Schmidt demand, +working tensor width, estimated work/write volume, and route length across +every hierarchical scale, not only the root cut. It enables bounded subtree +reconfiguration and simulated annealing by default; override these with +`topology_refine="subtree"`, `topology_budget=`, `search="anneal"`, and +`search_budget=`. The result is still a cheap layout proxy rather than a real +TTN replay, so the state-aware pilot remains the final accuracy check. + Use `order="quality"` with `finder.run()` (or set it on the finder) for the MPS-style higher-quality offline search. It enables bounded greedy leaf refinement, bounded binary-tree nearest-neighbor-interchange (NNI) topology @@ -581,6 +590,8 @@ the accuracy/bond-dimension cost. For offline quality searches, add `search="nevergrad"`. Nevergrad starts from the spectral/greedy plan, proposes leaf orders, and keeps its result only when it improves the same chi-aware objective. It never acts on a live optimizer. +For `objective="full_tree"`, use `search="anneal"` to explore subtree +replacements at multiple scales without the optional Nevergrad dependency. Install the optional dependency with `pip install pepsy[layout]`: ```python diff --git a/src/pepsy/optimizers/tree/layout.py b/src/pepsy/optimizers/tree/layout.py index 4e4a30c..198fb38 100644 --- a/src/pepsy/optimizers/tree/layout.py +++ b/src/pepsy/optimizers/tree/layout.py @@ -146,10 +146,13 @@ def _normalize_topology_refinement(refine): "joint": "nni", "joint_greedy": "nni", "greedy_topology": "nni", + "reconfigure": "subtree", + "subtree_reconfigure": "subtree", + "all_scales": "subtree", } name = aliases.get(name, name) - if name != "nni": - raise ValueError("topology_refine must be None or 'nni'.") + if name not in {"nni", "subtree"}: + raise ValueError("topology_refine must be None, 'nni', or 'subtree'.") return name @@ -158,10 +161,16 @@ def _normalize_layout_search(search): if search is None or search is False: return None name = str(search).replace("-", "_").strip().lower() - aliases = {"ng": "nevergrad", "never_grad": "nevergrad"} + aliases = { + "ng": "nevergrad", + "never_grad": "nevergrad", + "simulated_annealing": "anneal", + "subtree_anneal": "anneal", + "annealing": "anneal", + } name = aliases.get(name, name) - if name != "nevergrad": - raise ValueError("search must be None or 'nevergrad'.") + if name not in {"nevergrad", "anneal"}: + raise ValueError("search must be None, 'nevergrad', or 'anneal'.") return name @@ -270,15 +279,20 @@ def _normalize_layout_objective(objective): "hypergraph_load": "hypergraph", "per_edge": "hypergraph", "per_edge_load": "hypergraph", + "full": "full_tree", + "tree": "full_tree", + "cotengra": "full_tree", + "all_scales": "full_tree", } name = aliases.get(name, name) if name not in { - "path", "congestion", "hybrid", "compression", "hypergraph" + "path", "congestion", "hybrid", "compression", "hypergraph", + "full_tree", }: raise ValueError( f"Unknown tree layout objective {objective!r}. " "Expected 'path', 'congestion', 'compression', 'hypergraph', " - "or 'hybrid'." + "'full_tree', or 'hybrid'." ) return name @@ -1288,7 +1302,7 @@ class TreeLayoutFinder: (see :meth:`TreePlan.from_order`). dense_max : int Maximum subsystem size for dense spectral reordering. - objective : {"path", "congestion", "compression", "hypergraph", "hybrid"} + objective : {"path", "congestion", "compression", "hypergraph", "full_tree", "hybrid"} Layout objective. `"path"` preserves the co-occurrence/path-length heuristic; `"congestion"` selects among layout candidates using the predicted operator-Schmidt load on tree edges. `"hybrid"` combines @@ -1298,6 +1312,10 @@ class TreeLayoutFinder: objective. `"hypergraph"` is the direct multi-site mode: it ranks plans from the full support hyperedges and per-edge Schmidt loads, then applies bounded leaf and binary-topology refinement by default. + `"full_tree"` evaluates dynamic bond pressure, tensor width, estimated + work, write volume, and route length across every tree scale. It is + the high-quality, Cotengra-inspired mode and is opt-in because its + bounded subtree search is more expensive. order : {None, "quality"}, optional Optional high-quality offline mode. `"quality"` enables bounded greedy refinement and opportunistic Nevergrad refinement; omitted @@ -1309,22 +1327,27 @@ class TreeLayoutFinder: Optional fixed-plan local search used by :meth:`run` and recommendation methods. `"greedy"` tries adjacent leaf-label swaps before simulation; it never changes a live :class:`TreeOptimizer` tree. - topology_refine : {None, "nni"} - Optional joint topology refinement for binary candidates. `"nni"` - tries bounded nearest-neighbor interchange moves on internal edges, - retaining only objective-improving trees. It never changes a live + topology_refine : {None, "nni", "subtree"} + Optional joint topology refinement. `"nni"` tries bounded + nearest-neighbor interchange moves on binary internal edges; + `"subtree"` reconfigures descendant subtrees at all scales. Both + retain only accepted candidates and never change a live :class:`TreeOptimizer` tree. refine_budget : int, optional Maximum greedy swap proposals per candidate plan. Defaults to at most 64 proposals when refinement is enabled. topology_budget : int, optional - Maximum NNI proposals per candidate plan. Defaults to at most 64 - proposals when topology refinement is enabled. - search : {None, "nevergrad"} + Maximum topology proposals per candidate plan. Defaults to at most 64 + proposals when topology refinement is enabled. For ``"subtree"``, + proposals are sampled across the available descendant scales. + search : {None, "nevergrad", "anneal"} Optional offline derivative-free refinement. It is never run unless - requested and requires the optional ``nevergrad`` package. + requested. `"nevergrad"` refines leaf order and requires the optional + package; `"anneal"` performs bounded simulated annealing over subtree + reconfigurations and has no additional dependency. search_budget : int - Number of Nevergrad objective evaluations per candidate plan. + Number of offline search evaluations per candidate plan. For + ``search="anneal"``, this is the number of subtree proposals. seed : int Reproducible seed used by the optional Nevergrad stage. nevergrad_optimizer : str @@ -1491,6 +1514,7 @@ def __init__(self, gates=None, n=None, *, supports=None, structure="quality", self._plan_cache = {} self._edge_load_cache = {} self._rank_diagnostics_cache = {} + self._full_tree_profile_cache = {} self._schmidt_rank_cache = {} self._similarity_cache = {} self._congestion_weights_cache = None @@ -1615,6 +1639,8 @@ def _resolve_search_settings( topology_refine = self.topology_refine if self.objective == "hypergraph" and topology_refine is None: topology_refine = "nni" + elif self.objective == "full_tree" and topology_refine is None: + topology_refine = "subtree" else: topology_refine = _normalize_topology_refinement(topology_refine) if topology_budget is _DEFAULT_SEARCH_OPTION: @@ -1628,6 +1654,8 @@ def _resolve_search_settings( if search is _DEFAULT_SEARCH_OPTION: search = self.search + if self.objective == "full_tree" and search is None: + search = "anneal" else: search = _normalize_layout_search(search) if search_budget is _DEFAULT_SEARCH_OPTION: @@ -1746,6 +1774,104 @@ def _nni_edges(plan): if len(plan.children.get(child, ())) == 2 ) + @staticmethod + def _subtree_nodes(plan, root): + """Return all node ids in the rooted subtree at ``root``.""" + nodes = set() + stack = [root] + while stack: + node = stack.pop() + if node in nodes: + continue + nodes.add(node) + stack.extend(plan.children[node]) + return nodes + + def _plan_with_subtree_reconfiguration(self, plan, subtree_root, rng): + """Rebuild one descendant subtree while preserving its attachment.""" + old_nodes = self._subtree_nodes(plan, subtree_root) + subtree_qubits = sorted( + plan.qubit_of_leaf[node] + for node in old_nodes + if node in plan.qubit_of_leaf + ) + if len(subtree_qubits) < 4: + return None + + global_weights = self._similarity_weights( + self._congestion_pair_weights() + if self.objective == "full_tree" else None + ) + local_index = {q: i for i, q in enumerate(subtree_qubits)} + local_weights = { + (local_index[qa], local_index[qb]): weight + for (qa, qb), weight in global_weights.items() + if qa in local_index and qb in local_index + } + local_order = list(range(len(subtree_qubits))) + rng.shuffle(local_order) + local_structure = ( + "adaptive" if self.structure == "adaptive" else "balanced" + ) + local_root_qubit = ( + len(subtree_qubits) + if subtree_root == plan.root and plan.root_qubit is not None + else None + ) + lower_arities = [ + len(children) + for node, children in plan.children.items() + if node != plan.root and children + ] + local_max_arity = max(lower_arities, default=2) + local_top_arity = ( + plan.top_arity if subtree_root == plan.root else None + ) + local_plan = TreePlan.from_order( + local_order, + weights=local_weights, + structure=local_structure, + max_arity=local_max_arity, + community_frac=self.community_frac, + star_frac=self.star_frac, + dense_max=self.dense_max, + root_qubit=local_root_qubit, + top_arity=local_top_arity, + ) + + next_node = max(plan.children, default=-1) + 1 + local_to_global = {local_plan.root: subtree_root} + for local_node in local_plan.children: + if local_node != local_plan.root: + local_to_global[local_node] = next_node + next_node += 1 + + children = { + node: tuple(child_ids) + for node, child_ids in plan.children.items() + if node not in old_nodes + } + qubit_of_leaf = { + node: qubit + for node, qubit in plan.qubit_of_leaf.items() + if node not in old_nodes + } + for local_node, local_children in local_plan.children.items(): + global_node = local_to_global[local_node] + children[global_node] = tuple( + local_to_global[child] for child in local_children + ) + for local_node, local_qubit in local_plan.qubit_of_leaf.items(): + qubit_of_leaf[local_to_global[local_node]] = ( + subtree_qubits[local_qubit] + ) + return TreePlan.from_children( + children, + qubit_of_leaf, + root=plan.root, + root_qubit=plan.root_qubit, + ) + def _path_score_and_max(self, plan): """Return the weighted interaction path sum and longest active path.""" score = 0.0 @@ -1850,6 +1976,18 @@ def _objective_key(self, plan): return self._path_score_and_max(plan) if self.objective == "congestion": return self._congestion_key(plan) + if self.objective == "full_tree": + profile = self.full_tree_profile(plan) + return ( + profile["peak_tensor_log2"], + profile["peak_work_log2"], + profile["log_total_write"], + profile["log_total_work"], + profile["peak_edge_demand_log2"], + profile["total_edge_demand_log2"], + profile["total_route_length"], + self.score(plan), + ) if self.objective in {"compression", "hypergraph"}: loads = self.edge_loads(plan) values = tuple(loads.values()) @@ -1880,6 +2018,16 @@ def _selection_loss(self, plan, chi): key = self._objective_key(plan) if self.objective == "path": value = key[0] + elif self.objective == "full_tree": + value = ( + key[0] + + 0.50 * key[1] + + 0.10 * key[2] + + 0.10 * key[3] + + 0.01 * key[4] + + 0.001 * key[5] + + 1.0e-6 * key[6] + ) elif self.objective in {"congestion", "compression", "hypergraph"}: value = key[0] + 1.0e-6 * key[1] + 1.0e-12 * key[2] else: @@ -1896,6 +2044,9 @@ def _discard_plan_cache(self, plan): cached = self._rank_diagnostics_cache.get(id(plan)) if cached is not None and cached[0] is plan: del self._rank_diagnostics_cache[id(plan)] + cached = self._full_tree_profile_cache.get(id(plan)) + if cached is not None and cached[0] is plan: + del self._full_tree_profile_cache[id(plan)] def _refine_plan_greedy(self, plan, *, chi, budget, progbar=False): """Greedily improve a fixed topology through adjacent leaf swaps.""" @@ -2037,6 +2188,198 @@ def _refine_plan_topology(self, plan, *, chi, budget, progbar=False): "final_key": current_key, } + def _refine_plan_subtree( + self, plan, *, chi, budget, seed, progbar=False + ): + """Greedily accept subtree replacements across all tree scales.""" + initial_key = self._selection_key(plan, chi) + if budget < 1: + return plan, { + "method": "subtree", + "search": "greedy", + "evaluations": 0, + "accepted_moves": 0, + "initial_key": initial_key, + "final_key": initial_key, + "scales_visited": (), + } + rng = np.random.default_rng(seed) + current = plan + current_key = initial_key + evaluations = 0 + accepted_moves = 0 + visited_scales = set() + progress = None + if progbar: + from tqdm import tqdm # pylint: disable=import-outside-toplevel + + progress = tqdm( + total=budget, + desc="tree layout subtree", + leave=False, + ) + while evaluations < budget: + current_candidates = [] + for node, children in current.children.items(): + if not children: + continue + subtree_size = sum( + 1 + for descendant in self._subtree_nodes(current, node) + if descendant in current.qubit_of_leaf + ) + if subtree_size >= 4: + current_candidates.append((node, subtree_size)) + if not current_candidates: + break + evaluations += 1 + if progress is not None: + progress.update() + nodes = np.asarray( + [node for node, _size in current_candidates], dtype=int + ) + weights = np.asarray( + [np.log2(size) for _node, size in current_candidates] + ) + weights /= weights.sum() + node = int(rng.choice(nodes, p=weights)) + visited_scales.add(_tree_node_scales(current)[node]) + candidate = self._plan_with_subtree_reconfiguration( + current, node, rng + ) + if candidate is None: + continue + candidate_key = self._selection_key(candidate, chi) + if candidate_key < current_key: + self._discard_plan_cache(current) + current = candidate + current_key = candidate_key + accepted_moves += 1 + else: + self._discard_plan_cache(candidate) + if progress is not None: + progress.close() + return current, { + "method": "subtree", + "search": "greedy", + "evaluations": evaluations, + "accepted_moves": accepted_moves, + "initial_key": initial_key, + "final_key": current_key, + "scales_visited": tuple(sorted(visited_scales)), + } + + def _anneal_plan_subtree( + self, plan, *, chi, budget, seed, progbar=False + ): + """Anneal subtree replacements across all available tree scales.""" + initial_key = self._selection_key(plan, chi) + candidates = [] + for node, children in plan.children.items(): + if not children: + continue + subtree_size = sum( + 1 + for descendant in self._subtree_nodes(plan, node) + if descendant in plan.qubit_of_leaf + ) + if subtree_size >= 4: + candidates.append((node, subtree_size)) + if budget < 1 or not candidates: + return plan, { + "method": "subtree", + "search": "anneal", + "evaluations": 0, + "accepted_moves": 0, + "initial_key": initial_key, + "final_key": initial_key, + "scales_visited": (), + } + + rng = np.random.default_rng(seed) + current = plan + current_key = initial_key + current_loss = self._selection_loss(current, chi) + best = current + best_key = current_key + evaluations = 0 + accepted_moves = 0 + visited_scales = set() + initial_temperature = max(1.0, abs(current_loss) * 0.05) + progress = None + if progbar: + from tqdm import tqdm # pylint: disable=import-outside-toplevel + + progress = tqdm( + total=budget, + desc="tree layout subtree anneal", + leave=False, + ) + + while evaluations < budget: + evaluations += 1 + if progress is not None: + progress.update() + current_candidates = [] + for current_node, current_children in current.children.items(): + if not current_children: + continue + current_size = sum( + 1 + for descendant in self._subtree_nodes(current, current_node) + if descendant in current.qubit_of_leaf + ) + if current_size >= 4: + current_candidates.append((current_node, current_size)) + if not current_candidates: + break + nodes = np.asarray( + [node for node, _size in current_candidates], dtype=int + ) + weights = np.asarray( + [np.log2(size) for _node, size in current_candidates] + ) + weights /= weights.sum() + node = int(rng.choice(nodes, p=weights)) + visited_scales.add(_tree_node_scales(current)[node]) + candidate = self._plan_with_subtree_reconfiguration( + current, node, rng + ) + if candidate is None: + continue + candidate_loss = self._selection_loss(candidate, chi) + delta = candidate_loss - current_loss + fraction = evaluations / max(1, budget) + temperature = initial_temperature * max(1.0e-3, 1.0 - fraction) + accept = delta <= 0.0 or rng.random() < np.exp( + -min(700.0, delta / temperature) + ) + if accept: + current = candidate + current_loss = candidate_loss + current_key = self._selection_key(current, chi) + accepted_moves += 1 + if current_key < best_key: + self._discard_plan_cache(best) + best = current + best_key = current_key + else: + self._discard_plan_cache(candidate) + + if progress is not None: + progress.close() + if best is not current: + self._discard_plan_cache(current) + return best, { + "method": "subtree", + "search": "anneal", + "evaluations": evaluations, + "accepted_moves": accepted_moves, + "initial_key": initial_key, + "final_key": best_key, + "scales_visited": tuple(sorted(visited_scales)), + } + def _refine_plan_nevergrad( self, plan, *, chi, budget, seed, optimizer_name, progbar=False ): @@ -2160,6 +2503,14 @@ def _improve_plan(self, plan, *, chi, settings, progbar=False): budget=settings["topology_budget"], progbar=progbar, ) + elif settings["topology_refine"] == "subtree": + plan, info["topology_refinement"] = self._refine_plan_subtree( + plan, + chi=chi, + budget=settings["topology_budget"], + seed=settings["seed"], + progbar=progbar, + ) if settings["refine"] == "greedy": plan, info["refinement"] = self._refine_plan_greedy( plan, @@ -2167,7 +2518,15 @@ def _improve_plan(self, plan, *, chi, settings, progbar=False): budget=settings["refine_budget"], progbar=progbar, ) - if settings["search"] == "nevergrad": + if settings["search"] == "anneal": + plan, info["search"] = self._anneal_plan_subtree( + plan, + chi=chi, + budget=settings["search_budget"], + seed=settings["seed"] + 1, + progbar=progbar, + ) + elif settings["search"] == "nevergrad": plan, info["search"] = self._refine_plan_nevergrad( plan, chi=chi, @@ -2417,23 +2776,25 @@ def recommend_layered( refine : {None, "greedy"}, optional Override the finder refinement setting. `"greedy"` performs a bounded adjacent leaf-swap search on each candidate tree. - topology_refine : {None, "nni"}, optional - Override the optional binary-tree topology refinement. `"nni"` - performs bounded nearest-neighbor interchange proposals. It is a - no-op for the non-binary layered structure. + topology_refine : {None, "nni", "subtree"}, optional + Override the optional topology refinement. ``"nni"`` performs + bounded nearest-neighbor interchange proposals and is a no-op for + the non-binary layered structure; ``"subtree"`` rebuilds selected + descendant subtrees and is useful for all-scale objectives. refine_budget : int, optional Maximum greedy proposals per candidate. When omitted, an enabled greedy search uses at most ``min(n - 1, 64)`` proposals. topology_budget : int, optional - Maximum NNI proposals per candidate. When omitted, an enabled NNI + Maximum topology proposals per candidate. When omitted, an enabled search uses at most 64 proposals. - search : {None, "nevergrad"}, optional + search : {None, "nevergrad", "anneal"}, optional Override the finder offline search setting. Nevergrad optimizes - only the returned fixed plan; it never mutates a live TTN. + only the returned fixed plan; annealing explores subtree + replacements. Neither mutates a live TTN. search_budget, seed, nevergrad_optimizer - Optional Nevergrad configuration for each candidate plan. + Optional offline-search configuration for each candidate plan. progbar : bool, optional - Display greedy and Nevergrad search progress for each candidate. + Display local-search progress for each candidate. Returns ------- @@ -2472,6 +2833,18 @@ def recommend_layered( else: order = [int(q) for q in order] + # ``recommend_layered`` compares a fixed block-family. Keep the + # all-scale objective, but do not silently replace that family with + # arbitrary subtree topologies unless the caller explicitly asks for + # those search stages. + if self.objective == "full_tree": + if topology_refine is _DEFAULT_SEARCH_OPTION: + settings["topology_refine"] = None + settings["topology_budget"] = None + if search is _DEFAULT_SEARCH_OPTION: + settings["search"] = None + settings["search_budget"] = self.search_budget + candidates = [] for bs in options: plan = TreePlan.build_layered( @@ -2504,6 +2877,7 @@ def recommend_layered( "max_path": report["max_path"], "max_edge_load": report["max_edge_load"], "peak_bond_growth": report["peak_bond_growth"], + "full_tree_profile": report["full_tree"], "max_virtual_degree": report["max_virtual_degree"], "total_virtual_degree": report["total_virtual_degree"], "estimated_max_tensor_log2": report[ @@ -2562,17 +2936,23 @@ def run( ``chi`` and the fixed-plan ``refine`` / ``search`` controls can be overridden for this call. Pass ``progbar=True`` to display greedy and - Nevergrad search progress. Omitted values inherit the corresponding + offline search progress. Omitted values inherit the corresponding finder settings, so the original zero-argument behavior is unchanged. The explicit ``objective="hypergraph"`` mode is the one exception: when no refinement controls are supplied, it enables bounded greedy and binary-NNI stages so the full support hyperedges directly influence the returned layout. + ``objective="full_tree"`` enables bounded all-scale subtree + reconfiguration and annealing by default when no search controls are + explicitly supplied. It evaluates every tree scale using dynamic + bond-pressure and tensor-work proxies. + ``order="quality"`` is a convenience mode matching the MPS layout - API: it enables bounded greedy refinement and opportunistic Nevergrad - refinement when the optional dependency is installed. If Nevergrad is - unavailable, quality mode falls back to greedy refinement. Pass + API: it enables bounded greedy refinement and opportunistic offline + refinement. If Nevergrad is unavailable, quality mode falls back to + greedy refinement. For ``objective="full_tree"``, quality mode uses + the dependency-free annealing stage. Pass ``search=None`` or ``refine=None`` explicitly to disable either stage. """ if order is _DEFAULT_ORDER: @@ -2583,9 +2963,15 @@ def run( if refine is _DEFAULT_SEARCH_OPTION: refine = "greedy" if topology_refine is _DEFAULT_SEARCH_OPTION: - topology_refine = "nni" + topology_refine = ( + "subtree" if self.objective == "full_tree" else "nni" + ) if search is _DEFAULT_SEARCH_OPTION: - search = "nevergrad" if _nevergrad_available() else None + search = ( + "anneal" + if self.objective == "full_tree" + else ("nevergrad" if _nevergrad_available() else None) + ) if chi is _DEFAULT_CHI: chi = self.chi else: @@ -2774,7 +3160,7 @@ def recommend_arities( :meth:`recommend_layered`. They are applied to each arity candidate before selecting one final immutable plan. progbar : bool, optional - Display greedy and Nevergrad search progress for each candidate. + Display local-search progress for each candidate. """ if chi is _DEFAULT_CHI: chi = self.chi @@ -2844,6 +3230,7 @@ def recommend_arities( "max_path": report["max_path"], "max_edge_load": report["max_edge_load"], "peak_bond_growth": report["peak_bond_growth"], + "full_tree_profile": report["full_tree"], "estimated_max_tensor_log2": report[ "estimated_max_tensor_log2" ], @@ -3021,6 +3408,214 @@ def edge_loads(self, plan=None): ) return dict(loads) + @staticmethod + def _log2_add(total, value): + """Add two positive quantities represented by their log2 values.""" + if value == -np.inf: + return float(total) + if total == -np.inf: + return float(value) + return float(np.logaddexp2(total, value)) + + def _support_span(self, plan, support): + """Return the mask, nodes, and rooted edges spanned by ``support``.""" + support = tuple(dict.fromkeys(support)) + support_mask = 0 + for site in support: + support_mask |= 1 << site + site_nodes = [plan.node_of_qubit[site] for site in support] + if len(site_nodes) == 2: + path = plan.node_path(site_nodes[0], site_nodes[1]) + span_nodes = set(path) + crossed_edges = [ + (u, v) if plan.parent.get(v) == u else (v, u) + for u, v in zip(path, path[1:]) + ] + else: + span_nodes = set() + anchor = site_nodes[0] + for site_node in site_nodes: + span_nodes.update(plan.node_path(anchor, site_node)) + crossed_edges = [ + (parent, node) + for node in span_nodes + if (parent := plan.parent.get(node)) in span_nodes + ] + return support_mask, span_nodes, crossed_edges + + def full_tree_profile(self, plan=None): + """Return a dynamic, all-scale cost profile for ``plan``. + + The profile is a cheap layout proxy, not a replacement for replaying + the circuit. It accumulates uncapped operator-Schmidt demand on every + tree edge, tracks the capped working bond pressure at the configured + ``chi``, estimates tensor widths and write/work volume for every + touched node, and groups those quantities by hierarchical tree scale. + """ + if plan is None: + plan = self.run() + cache_key = id(plan) + cached = self._full_tree_profile_cache.get(cache_key) + if cached is not None and cached[0] is plan: + return cached[1] + + below = plan.subtree_qubit_masks() + node_scales = _tree_node_scales(plan) + edges = tuple( + (parent, child) + for parent, children in plan.children.items() + for child in children + ) + demand_log = {edge: 0.0 for edge in edges} + bond_log = {edge: 0.0 for edge in edges} + log_chi = ( + float(np.log2(self.chi)) if self.chi is not None else float("inf") + ) + scales = {} + for scale in sorted(set(node_scales.values())): + scales[scale] = { + "node_count": 0, + "edge_count": 0, + "peak_tensor_log2": 0.0, + "log_total_tensor_size": -np.inf, + "peak_edge_demand_log2": 0.0, + "total_edge_demand_log2": 0.0, + } + for node, scale in node_scales.items(): + scales[scale]["node_count"] += 1 + for parent, child in edges: + scales[node_scales[child]]["edge_count"] += 1 + + def node_log_size(node): + incident = [] + if node in plan.parent: + incident.append((plan.parent[node], node)) + incident.extend((node, child) for child in plan.children[node]) + return float( + int(node in plan.qubit_of_node) + + sum(bond_log[edge] for edge in incident) + ) + + peak_tensor_log2 = 0.0 + peak_work_log2 = 0.0 + log_total_write = -np.inf + log_total_work = -np.inf + total_route_length = 0 + event_count = 0 + exact_events = 0 + bounded_events = 0 + bound_reasons = {} + + for payload, support, event_type, temporal_factor in zip( + self.payloads, + self.supports, + self.event_types, + self.temporal_factors, + ): + support = tuple(dict.fromkeys(support)) + if ( + temporal_factor <= 0.0 + or len(support) < 2 + or str(event_type).lower() in { + "measure", "reset", "measure_reset", "cap" + } + ): + continue + support_mask, span_nodes, crossed_edges = self._support_span( + plan, support + ) + if not crossed_edges: + continue + event_count += 1 + total_route_length += len(crossed_edges) + for edge in crossed_edges: + _parent, child = edge + left_mask = support_mask & below[child] + if not left_mask or left_mask == support_mask: + continue + left = tuple( + site for site in support if left_mask & (1 << site) + ) + info = self._schmidt_rank_info(payload, support, left) + delta = float(temporal_factor) * float( + np.log2(max(1, int(info["rank"]))) + ) + demand_log[edge] += delta + bond_log[edge] = min(log_chi, bond_log[edge] + delta) + if info["exact"]: + exact_events += 1 + else: + bounded_events += 1 + reason = info["reason"] + bound_reasons[reason] = bound_reasons.get(reason, 0) + 1 + + tensor_logs = [node_log_size(node) for node in span_nodes] + event_write_log = max(tensor_logs, default=0.0) + event_work_logs = [ + log_size + 2.0 * len(support) + for log_size in tensor_logs + ] + event_work_log = max(event_work_logs, default=0.0) + peak_tensor_log2 = max(peak_tensor_log2, event_write_log) + peak_work_log2 = max(peak_work_log2, event_work_log) + log_total_write = self._log2_add( + log_total_write, event_write_log + ) + log_total_work = self._log2_add( + log_total_work, event_work_log + ) + for node, log_size in zip(span_nodes, tensor_logs): + scale = scales[node_scales[node]] + scale["peak_tensor_log2"] = max( + scale["peak_tensor_log2"], log_size + ) + scale["log_total_tensor_size"] = self._log2_add( + scale["log_total_tensor_size"], log_size + ) + for edge in crossed_edges: + scale = scales[node_scales[edge[1]]] + scale["peak_edge_demand_log2"] = max( + scale["peak_edge_demand_log2"], demand_log[edge] + ) + + for node in plan.children: + log_size = node_log_size(node) + scale = scales[node_scales[node]] + scale["peak_tensor_log2"] = max( + scale["peak_tensor_log2"], log_size + ) + scale["log_total_tensor_size"] = self._log2_add( + scale["log_total_tensor_size"], log_size + ) + for edge, demand in demand_log.items(): + scale = scales[node_scales[edge[1]]] + scale["peak_edge_demand_log2"] = max( + scale["peak_edge_demand_log2"], demand + ) + scale["total_edge_demand_log2"] += demand + + profile = { + "event_count": event_count, + "peak_tensor_log2": float(peak_tensor_log2), + "peak_work_log2": float(peak_work_log2), + "log_total_write": float( + 0.0 if log_total_write == -np.inf else log_total_write + ), + "log_total_work": float( + 0.0 if log_total_work == -np.inf else log_total_work + ), + "peak_edge_demand_log2": float(max(demand_log.values(), default=0.0)), + "total_edge_demand_log2": float(sum(demand_log.values())), + "peak_bond_log2": float(max(bond_log.values(), default=0.0)), + "total_route_length": int(total_route_length), + "exact_events": int(exact_events), + "bounded_events": int(bounded_events), + "bound_reasons": bound_reasons, + "scales": scales, + } + self._full_tree_profile_cache[cache_key] = (plan, profile) + return profile + def _congestion_key(self, plan): """Return the lexicographic key used by the load-aware objective.""" loads = self.edge_loads(plan) @@ -3133,6 +3728,10 @@ def report(self, plan=None, *, include_edge_loads=True): hybrid_cost = None if self.objective == "hybrid" and loads is not None: hybrid_cost = self._hybrid_key(plan)[0] + full_tree = ( + self.full_tree_profile(plan) + if self.objective == "full_tree" else None + ) arity_histogram = {} for node, children in plan.children.items(): if children: @@ -3166,6 +3765,7 @@ def report(self, plan=None, *, include_edge_loads=True): } if self.objective == "hypergraph" and loads is not None else None ), + "full_tree": full_tree, "root": plan.root, "root_qubit": plan.root_qubit, "top_arity": plan.top_arity, diff --git a/src/pepsy/optimizers/tree/optimizer.py b/src/pepsy/optimizers/tree/optimizer.py index 40da1dd..86ee922 100644 --- a/src/pepsy/optimizers/tree/optimizer.py +++ b/src/pepsy/optimizers/tree/optimizer.py @@ -331,7 +331,7 @@ class TreeOptimizer: is built automatically. Set ``top_arity=3`` with ``max_arity=2`` for the conventional binary TTN with a three-leg top tensor. This keeps every tensor rank at most three. - layout_objective : {"path", "congestion", "compression", "hypergraph", "hybrid"} + layout_objective : {"path", "congestion", "compression", "hypergraph", "full_tree", "hybrid"} Objective used when building an automatic tree. ``"path"`` is the backward-compatible interaction-path heuristic; ``"congestion"`` selects a candidate using predicted operator-Schmidt edge load; @@ -339,6 +339,8 @@ class TreeOptimizer: estimated local tensor cost at ``chi``; ``"hypergraph"`` directly scores every original multi-qubit support across every crossed tree edge and enables bounded leaf/NNI refinement by default; + ``"full_tree"`` evaluates dynamic all-scale tensor width, work, write, + bond pressure, and route costs; ``"hybrid"`` combines normalized path, peak-load, and total-load costs. Pass a configured :class:`TreeLayoutFinder` through ``layout=`` to customize its hybrid weights or enable pre-simulation refinement. diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index 3321517..2612ed8 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -1084,6 +1084,71 @@ def test_tree_candidate_plans_include_quality_for_state_aware_pilots(): ) +def test_full_tree_profile_reports_dynamic_cost_at_all_scales(): + """Whole-tree mode exposes width, work, demand, and scale diagnostics.""" + gates = [ + (pepsy.cnot(), (0, 1)), + (pepsy.cnot(), (1, 2)), + (pepsy.cnot(), (3, 4)), + (pepsy.cnot(), (2, 4)), + ] + finder = TreeLayoutFinder( + gates, n=5, max_arity=2, objective="full_tree", chi=4, + ) + plan = TreePlan.from_order(range(5), structure="balanced", max_arity=2) + profile = finder.full_tree_profile(plan) + + assert profile["event_count"] == len(gates) + assert profile["peak_tensor_log2"] >= 1.0 + assert profile["peak_work_log2"] >= profile["peak_tensor_log2"] + assert profile["total_route_length"] > 0 + assert profile["scales"] + assert all( + { + "node_count", + "edge_count", + "peak_tensor_log2", + "peak_edge_demand_log2", + } <= set(scale_info) + for scale_info in profile["scales"].values() + ) + + report = finder.report(plan) + assert report["objective"] == "full_tree" + assert report["full_tree"] == profile + assert len(report["objective_key"]) == 8 + + +def test_full_tree_anneals_subtrees_without_changing_binary_contract(): + """All-scale subtree search preserves the requested binary tree shape.""" + finder = TreeLayoutFinder( + [(pepsy.cnot(), (0, 3)), (pepsy.cnot(), (3, 1)), + (pepsy.cnot(), (2, 5)), (pepsy.cnot(), (4, 5))], + n=6, + max_arity=2, + top_arity=3, + objective="full_tree", + chi=4, + seed=7, + ) + recommendation = finder.recommend_arities( + (2,), + topology_budget=3, + search_budget=4, + refine=None, + ) + candidate = recommendation["candidates"][0] + planning = candidate["planning"] + plan = recommendation["plan"] + + assert plan.top_arity == 3 + assert plan.is_binary() + assert planning["topology_refinement"]["method"] == "subtree" + assert planning["search"]["method"] == "subtree" + assert planning["search"]["search"] == "anneal" + assert candidate["full_tree_profile"]["scales"] + + def test_tree_edge_loads_match_full_edge_reference(): """Steiner-only edge scanning preserves the full congestion calculation.""" rng = np.random.default_rng(109) From b4ffaf65e97546bb8a486298eabeb340faef6935 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Fri, 31 Jul 2026 11:19:49 -0600 Subject: [PATCH 39/70] Add native fermionic PEPO and BP workflows --- docs/api/bp.md | 142 +- docs/api/operators/gates.md | 40 + docs/api/operators/hamiltonians.md | 58 +- docs/api/tensors/symmetric.md | 64 +- .../references/belief_propagation.md | 4 +- src/pepsy/bp/__init__.py | 10 + src/pepsy/bp/series.py | 3020 ++++++++++++++++- src/pepsy/operators/gates.py | 231 +- src/pepsy/operators/hamiltonians.py | 44 +- src/pepsy/tensors/symmetric.py | 359 +- tests/test_bp_open_series.py | 305 ++ tests/test_bp_symmray.py | 115 + tests/test_gate.py | 80 + tests/test_ham.py | 77 + tests/test_symmetric_tensors.py | 232 ++ 15 files changed, 4597 insertions(+), 184 deletions(-) diff --git a/docs/api/bp.md b/docs/api/bp.md index 058c406..8fb4e56 100644 --- a/docs/api/bp.md +++ b/docs/api/bp.md @@ -7,6 +7,42 @@ the top-level package. > API details are maintained as handwritten Markdown in this page. +## Native fermionic PEPO norm workflow + +Native fermionic PEPOs can be evolved with `gate_simple`, including its +simple-update gauge dictionary, and their Frobenius norm can be measured with +`PEPO.norm()`. The D2BP loop-cluster route accepts the evolved PEPO directly: + +```python +from pepsy import Fermion, OneDMap +from pepsy.operators.gates import gate_simple +from pepsy.bp import loop_cluster_expand + +pepo = fermion.to_pepo( + terms, + Lx=2, + Ly=2, + mapper=OneDMap(2, 2, mode="snake-row-major"), + fermionic=True, +) +gauges = {} +pepo.gauge_all_simple_(gauges=gauges, progbar=False) +pepo = gate_simple( + pepo, + fermion.hopping_gate(0.001, t=1.0).H, + where=((0, 0), (0, 1)), + gauges=gauges, + inplace=False, +) +norm = pepo.norm() +correction = loop_cluster_expand(pepo, gloops=2, norm="2norm") +``` + +This path preserves native `U1`, `U1U1`, and `Z2` tensors. Simple-update +gauges are used by `gate_simple`; D2BP recomputes its own norm messages. The +`norm="1norm"` gauge route requires a closed scalar tensor network, so it is +not the direct route for an open PEPO with ket/bra physical indices. + ## Long-range PEPS expectations Use `compute_boundary_expectation` for batched one- and two-site operators, @@ -121,8 +157,8 @@ is completed. For separated sites, use `partial_trace_open_loop_series_expand` when the explicit configuration family should include Q paths between the retained -sites. Its integer `gloops` is a maximum number of excited virtual edges. A -configuration is retained when degree-one Q vertices occur only at the +sites. Use `edge_cutoff` to set the maximum number of excited virtual edges. +A configuration is retained when degree-one Q vertices occur only at the selected rho sites, so the sum contains open paths, closed loops, and path-plus-loop combinations: @@ -137,19 +173,97 @@ bp = two_norm_bp(peps.tn, max_iterations=1000, tol=1e-10) rho = partial_trace_open_loop_series_expand( peps.tn, where=((0, 0), (0, 7)), - gloops=8, + edge_cutoff=8, messages=bp.messages, run_bp=False, ) ``` +The edge geometry is generated lazily, with shortest support-connecting paths +yielded first. Set `max_terms`, `max_enumeration_time`, or +`max_enumeration_memory` to fail before an uncontrolled geometry expansion; +limits raise `OpenLoopEnumerationLimitError` rather than returning a partial +sum. `gloops` remains a compatibility alias, but new code should use the +route-specific names. + +For very distant supports, set `corridor_width` to use the bounded corridor +route. It retains a small weighted-shortest-path beam, inflates those paths +by the requested graph width, and adds connected loop decorations only near +sampled corridor segments: + +```python +rho = partial_trace_open_loop_series_expand( + peps.tn, + where=((0, 0), (999, 999)), + corridor_width=4, + max_path_candidates=8, + loop_decoration_size=6, + corridor_segment_length=32, + max_loop_clusters_per_segment=8, + corridor_max_bond=128, + max_corridor_edges=100_000, +) +``` + +This is an explicitly controlled approximation: disconnected products of +far-separated loop clusters are omitted. Increase the corridor width, +candidate count, decoration size, or boundary bond dimension and compare the +incremental correction using the `open_rho_corridor` and +`open_scalar_corridor` diagnostics. `path_edge_weights` can supply positive +edge costs for ranking routes; unspecified edges have unit cost. + +For a measurement workflow that must inspect the geometry before doing any +numerical contractions, use `diagnose_open_loop_series`. It accepts native +operators made by `Fermion`, records the selected route, paths, loop terms, +Cotengra FLOP estimates, and peak-memory estimates, and can be reused during +measurement: + +```python +from pepsy.bp import ( + OpenLoopObservableTerm, + OpenLoopSeriesDiagnosticCache, + compute_local_expectation_open_loop_series, + diagnose_open_loop_series, +) + +term = OpenLoopObservableTerm( + ((0, 0), (999, 999)), + fermion.hopping_operator(), +) +diagnostic = diagnose_open_loop_series( + peps.tn, + term, + mode="auto", + edge_cutoff=2_000, + max_terms=10_000, + diagnostic_cache=OpenLoopSeriesDiagnosticCache(), +) +value = compute_local_expectation_open_loop_series( + peps.tn, + term, + mode="auto", + diagnostic=diagnostic, +) +``` + +The diagnostic phase builds contraction trees but does not contract tensor +values. `mode="auto"` selects the graded cluster-compatible route first for +cyclic native fermionic supports, and selects a corridor when the support +distance exceeds `auto_corridor_distance`. Reusing the same +`OpenLoopSeriesDiagnosticCache` lets later measurements reuse the geometry and +cost report; distinct operator values with the same support and shape do not +share numerical results. If no `diagnostic` is supplied, scalar measurement +with `mode="auto"` performs this diagnostic pass internally before starting +the numerical contractions. + This path performs an explicit configuration sum and normalizes only after the sum; it does not apply the scalar disconnected-loop resummation used by `partial_trace_edge_loop_series_expand`. For a cutoff sweep, reuse both the converged messages and the two caches. The -same `info` dictionary keeps already-contracted rho terms, while the -`OpenLoopSeriesCache` keeps the eligible edge configurations: +same `info` dictionary keeps already-contracted rho terms, regional +contraction paths, and physical output labels, while the `OpenLoopSeriesCache` +keeps the eligible edge configurations: ```python cache = OpenLoopSeriesCache() @@ -158,7 +272,7 @@ for cutoff in (2, 4, 6, 8): rho = partial_trace_open_loop_series_expand( peps.tn, where=((0, 0), (0, 7)), - gloops=cutoff, + edge_cutoff=cutoff, messages=bp.messages, run_bp=False, cache=cache, @@ -176,6 +290,22 @@ On cyclic native fermionic graphs, those scalar and rho APIs use the equivalent graded loop-cluster contraction internally because Symmray cannot currently contract arbitrary mixed open ``P/Q`` configurations; the returned rho remains native and the gate is still inserted in the ket/bra network. +For that route, pass `cluster_size`, not `edge_cutoff`: + +```python +rho = partial_trace_open_loop_series_expand( + peps.tn, + where=((0, 0), (0, 7)), + cluster_size=8, + messages=bp.messages, + run_bp=False, +) +``` + +The cyclic native route is selected before explicit edge discovery, so its +cluster regions do not pay the open-edge enumeration cost. Inspect +`info["open_rho_cluster_region_costs"]` or +`info["open_scalar_cluster_region_costs"]` for its contraction decisions. The same one-BP/many-support workflow is runnable in the downstream example `../pepsy_examples/symmetric_tensors/peps/bp_open_rho_series.py`; the long-range native doublon comparison is in diff --git a/docs/api/operators/gates.md b/docs/api/operators/gates.md index b449040..5c56634 100644 --- a/docs/api/operators/gates.md +++ b/docs/api/operators/gates.md @@ -1,4 +1,44 @@ # `pepsy.operators.gates` +The gate-to-operator builders accept native Symmray fermionic gates directly; +they do not convert them through dense arrays. Charge-neutral gates work by +default, such as `Fermion.hopping_gate(...)`: + +```python +fermion = pepsy.Fermion(spinful=True, symmetry="U1U1") +gate = fermion.hopping_gate(0.01, t=1.0) + +mpo = pepsy.build_mpo_from_gates(gate, where=(0, 1), max_bond=16) +pepo = pepsy.build_pepo_from_gates( + gate, + where=((0, 0), (0, 1)), + mapper=pepsy.OneDMap(2, 2, mode="snake-row-major"), + max_bond=16, +) +``` + +The resulting tensors remain `U1U1FermionicArray` (or the corresponding +`U1`/`Z2` native type). `build_pepo_from_gates` uses `OneDMap` to choose the +MPO ordering before embedding it on the 2D lattice. A definite-charge native +operator can also be used explicitly: + +```python +charged_mpo = pepsy.build_mpo_from_gates( + charged_gate, + where=(0, 1), + allow_charged=True, +) +charged_pepo = pepsy.build_pepo_from_gates( + charged_gate, + where=((0, 0), (0, 1)), + mapper=pepsy.OneDMap(2, 2, mode="snake-row-major"), + allow_charged=True, +) +``` + +`allow_charged=True` means the returned operator carries the accumulated +charge of the sequential gate product. It is opt-in because charged gates +change the symmetry sector; ordinary charge-preserving evolution should leave +it disabled. > API details are maintained as handwritten Markdown in this page. diff --git a/docs/api/operators/hamiltonians.md b/docs/api/operators/hamiltonians.md index e9a9d42..8cc3bc2 100644 --- a/docs/api/operators/hamiltonians.md +++ b/docs/api/operators/hamiltonians.md @@ -22,10 +22,60 @@ the Jordan-Wigner-compatible MPO convention is wanted. `Fermion.to_mpo(...)` and `SymHamiltonian.to_mpo(..., fermionic=True)` return native graded `FermionicArray` MPO tensors. Explicit mappings can contain -arbitrary neutral multi-site terms; non-contiguous supports are represented by -charged virtual channels. `ham_tn.build_mpo(..., fermionic=True)` selects the -same native path. Pass `to_backend=...` to map the stored Symmray blocks to a -selected array backend. +arbitrary homogeneous-charge multi-site terms; non-contiguous supports are +represented by charged virtual channels, and the open boundary carries the +operator charge when it is nonzero. `ham_tn.build_mpo(..., fermionic=True)` +selects the same native path. Pass `to_backend=...` to map the stored Symmray +blocks to a selected array backend. + +For a mixed-charge operator, request an explicit charge-sector decomposition: + +```python +sectors = fermion.to_mpo( + mixed_terms, + L=4, + fermionic=True, + charge_sectors=True, +) +# sectors[charge] is one homogeneous native MPO. +``` + +The same `charge_sectors=True` option is available on +`SymHamiltonian.to_mpo`, `Fermion.to_pepo`, `SymHamiltonian.to_pepo`, +`ham_tn.build_mpo`, and `ham_tn.build_pepo`; those methods return +`{charge: MPO}` or `{charge: PEPO}`. This keeps each block-sparse tensor +network within one charge sector while preserving the exact sum decomposition. + +The corresponding 2D entry points all use the same `OneDMap` ordering: + +```python +mapper = pepsy.OneDMap(3, 2, mode="snake-row-major") + +pepo = hamiltonian.to_pepo( + Lx=3, + Ly=2, + mapper=mapper, + fermionic=True, +) +pepo = fermion.build_pepo( + {(left, right): native_term}, + Lx=3, + Ly=2, + mapper=mapper, + fermionic=True, +) +pepo = builder.build_pepo( + {(left, right): native_term}, + fermion=fermion, + mapper=mapper, + fermionic=True, +) +``` + +Use a coordinate-keyed mapping for native terms, with one-site support written +as `((x, y),)`. PEPO embedding currently requires `snake` or +`snake-row-major` ordering; transverse lattice bonds are rank one unless +periodic PEPO bonds are requested with `cycle_peps=True`. Native MPO assembly, replay, and exact energy measurement are supported. The native energy path applies the MPO sitewise as a factorized graded MPO-MPS diff --git a/docs/api/tensors/symmetric.md b/docs/api/tensors/symmetric.md index f55fb34..15a1788 100644 --- a/docs/api/tensors/symmetric.md +++ b/docs/api/tensors/symmetric.md @@ -520,7 +520,7 @@ Current support is: - spinful Fermi-Hubbard ``model="fermi_hubbard_u1u1"`` with ``symmetry="U1U1"``, hopping, onsite interaction, nearest-neighbor density interaction, and chemical-potential terms. -- native graded MPO conversion for arbitrary neutral one- or multi-site +- native graded MPO conversion for arbitrary homogeneous-charge one- or multi-site ``FermionicArray`` terms, including non-contiguous support; Spinful total-particle-number ``model="fermi_hubbard"`` with ``symmetry="U1"`` @@ -593,8 +593,66 @@ Native fermionic gate streams from the same ``Fermion`` model can be passed to charge blocks during replay and compression. ``fermionic=False`` remains the explicit Jordan-Wigner compatibility choice for ``SymHamiltonian.to_mpo``. -Pass ``to_backend=`` to ``Fermion.to_mpo`` or ``Fermion.build_mpo`` when the -returned Symmray blocks should use a selected array backend. Native MPO +For a coordinate-keyed native operator on a 2D lattice, ``Fermion.to_pepo`` +provides the corresponding PEPO embedding: + +```python +left, right = (0, 1), (2, 2) +fermion = py.Fermion(spinful=True, symmetry="U1U1") +term = fermion.operator_term( + [(1.0, ((left, "create_up"), (right, "annihilate_up")))], + sites=(left, right), + add_hc=True, +) +pepo = fermion.to_pepo( + {(left, right): term}, + Lx=3, + Ly=3, + fermionic=True, + max_bond=16, +) +assert all(type(tensor.data).__name__.endswith("FermionicArray") for tensor in pepo) +``` + +``to_pepo`` preserves native Symmray grading and supports the spinful +``U1U1``, ``U1``, and ``Z2`` paths as well as spinless ``U1`` and ``Z2``. +Its current PEPO representation uses the selected snake-style MPO ordering +for fermionic channels; the added transverse lattice bonds are dimension one +unless periodic PEPO bonds are requested with ``cyclic=True``. Terms should +be homogeneous in charge for one MPO/PEPO. Native ``fermionic=True`` +construction supports both neutral and nonzero charges by carrying the +operator charge at the open MPO/PEPO boundary. For a mixed-charge collection, +pass ``charge_sectors=True`` to receive ``{charge: PEPO}`` (or ``{charge: +MPO}`` from ``to_mpo``). For odd-parity terms, pass ``label=`` to +``operator_term`` so the native dummy-mode phase metadata is retained. The +Jordan--Wigner compatibility path remains neutral-only. For a one-site +coordinate key, use ``((x, y),)`` rather than ``(x, y)``. + +The same native PEPO route is available from an existing Hamiltonian or as a +model-facing shorthand: + +```python +pepo = hamiltonian.to_pepo( + Lx=3, + Ly=3, + mapper=py.OneDMap(3, 3, mode="snake-row-major"), + fermionic=True, +) +pepo = fermion.build_pepo( + {(left, right): hopping}, + Lx=3, + Ly=3, + mapper=py.OneDMap(3, 3, mode="snake-row-major"), + fermionic=True, +) +``` + +For Hamiltonian-builder workflows, ``ham_tn.build_pepo(..., fermion=fermion, +fermionic=True, mapper=...)`` is equivalent. + +Pass ``to_backend=`` to ``Fermion.to_mpo``, ``Fermion.to_pepo``, or +``Fermion.build_mpo`` when the returned Symmray blocks should use a selected +array backend. Native MPO assembly, replay, and exact energy measurement are supported. Native MPO energy applies the operator sitewise as a factorized graded MPO-MPS contraction, so its cost is controlled by the MPS and MPO bond dimensions. diff --git a/docs/development/references/belief_propagation.md b/docs/development/references/belief_propagation.md index 0db3165..3178d58 100644 --- a/docs/development/references/belief_propagation.md +++ b/docs/development/references/belief_propagation.md @@ -23,8 +23,8 @@ side · **[roots]** foundational / prior art. | `loop_series_expand`, `LoopSeriesTerm`, `LoopSeriesCache` | edge-resolved `P + Q` loop series for D1BP and D2BP; retains excited-bond degree and distinct embeddings/chord subsets | Evenbly et al. 2409.03108 | | `partial_trace_loop_series_expand`, `compute_local_expectation_loop_series` | D2BP local reduced-density-matrix and scalar `P + Q` loop series; keeps physical output legs open and uses native Symmray virtual projectors | Evenbly et al. 2409.03108; quimb local loop-series API | | `partial_trace_edge_loop_series_expand`, `compute_local_expectation_edge_loop_series` | D2BP local RDM and graded scalar observable expansion over canonical explicit Q-edge terms; does not reinterpret Quimb's local-region cutoff | Evenbly et al. 2409.03108; Pepsy API | -| `partial_trace_open_loop_series_expand`, `partial_trace_open_loop_series_sweep` | Explicit D2BP rho configuration sum over open Q paths, closed loops, and attached or disconnected path-plus-loop terms; degree-one vertices are allowed only on retained rho sites; the sweep wrapper reuses one BP solve and cache across supports and cutoffs; cyclic native fermionic graphs use the equivalent native graded cluster contraction when mixed open P/Q contractions are unsupported | Evenbly et al. 2409.03108; Pepsy API | -| `compute_local_expectation_open_loop_series` | Scalar companion for long-range gates: reuses the open-loop family bookkeeping, inserts native gates through the graded open-bond projector route (or the equivalent graded cluster route on cyclic native fermionic graphs), normalizes numerator/denominator after the configuration estimate, and optionally filters explicit terms or cyclic cluster contractions by Cotengra log10-FLOP/log2-peak-size limits | Evenbly et al. 2409.03108; Pepsy API | +| `partial_trace_open_loop_series_expand`, `partial_trace_open_loop_series_sweep` | Lazy/path-first D2BP rho configuration sum over open Q paths, closed loops, and attached or disconnected path-plus-loop terms; `edge_cutoff` counts excited Q edges, bounded discovery uses `max_terms` / enumeration time / approximate geometry memory limits, and regional contraction paths are reused; `corridor_width` activates bounded weighted-shortest-path discovery, sampled connected loop decorations, and optional compressed boundary contraction; cyclic native fermionic graphs select the graded cluster route before edge discovery and use `cluster_size` | Evenbly et al. 2409.03108; Pepsy API | +| `compute_local_expectation_open_loop_series`, `diagnose_open_loop_series` | Scalar companion for long-range gates plus a non-contracting diagnostic pass: accepts native Fermion operators, reuses open-loop geometry, selects exact/corridor/graded-cluster routes, reports path/loop families and Cotengra FLOP/peak-memory estimates, and can reuse cached diagnostics before inserting the gate through the graded open-bond projector route | Evenbly et al. 2409.03108; Pepsy API | | `partial_trace_loop_cluster_expand`, `compute_local_expectation_loop_cluster` | D2BP local reduced-density-matrix and scalar generalized-loop cluster expansion; combines BP-closed regions with inclusion--exclusion counts | Gray et al. 2510.05647; quimb local cluster API | | `loop_expand` | explicit selector between the edge loop series and region loop-cluster expansion; preserves each method's cutoff and result metadata | Pepsy API | | `partitioned_expand`, `pne_expand`, `PNEExpansionResult` | linear and combinatorial partitioned network expansions for D1BP/D2BP, with optional residue, explicit projectors, open outputs, and fixed recursive schedules | Evenbly, Gray & Chan 2512.10910 | diff --git a/src/pepsy/bp/__init__.py b/src/pepsy/bp/__init__.py index c867656..5532d93 100644 --- a/src/pepsy/bp/__init__.py +++ b/src/pepsy/bp/__init__.py @@ -45,6 +45,10 @@ select_bp_candidate, ) from .series import ( + OpenLoopEnumerationLimitError, + OpenLoopObservableTerm, + OpenLoopSeriesDiagnostic, + OpenLoopSeriesDiagnosticCache, OpenLoopSeriesCache, OpenLoopSeriesSweepResult, LoopSeriesCache, @@ -57,6 +61,7 @@ partial_trace_edge_loop_series_expand, partial_trace_open_loop_series_expand, partial_trace_open_loop_series_sweep, + diagnose_open_loop_series, partial_trace_loop_cluster_expand, partial_trace_loop_series_expand, loop_series_expand, @@ -141,11 +146,16 @@ "LoopClusterTerm", "LoopSeriesCache", "OpenLoopSeriesCache", + "OpenLoopEnumerationLimitError", + "OpenLoopObservableTerm", + "OpenLoopSeriesDiagnostic", + "OpenLoopSeriesDiagnosticCache", "OpenLoopSeriesSweepResult", "LoopSeriesResult", "LoopSeriesTerm", "compute_local_expectation_edge_loop_series", "compute_local_expectation_open_loop_series", + "diagnose_open_loop_series", "compute_local_expectation_loop_cluster", "compute_local_expectation_loop_series", "partial_trace_edge_loop_series_expand", diff --git a/src/pepsy/bp/series.py b/src/pepsy/bp/series.py index 0fee4c2..d0c43ae 100644 --- a/src/pepsy/bp/series.py +++ b/src/pepsy/bp/series.py @@ -24,8 +24,11 @@ from collections import Counter, deque from dataclasses import dataclass, field import functools +import heapq from itertools import combinations import operator +import sys +import time from typing import Any, ClassVar import autoray as ar @@ -49,6 +52,10 @@ ) __all__ = [ + "OpenLoopEnumerationLimitError", + "OpenLoopObservableTerm", + "OpenLoopSeriesDiagnostic", + "OpenLoopSeriesDiagnosticCache", "OpenLoopSeriesCache", "OpenLoopSeriesSweepResult", "LoopSeriesCache", @@ -56,6 +63,7 @@ "LoopSeriesTerm", "compute_local_expectation_edge_loop_series", "compute_local_expectation_open_loop_series", + "diagnose_open_loop_series", "compute_local_expectation_loop_cluster", "partial_trace_loop_cluster_expand", "partial_trace_edge_loop_series_expand", @@ -66,6 +74,115 @@ ] +class OpenLoopEnumerationLimitError(RuntimeError): + """Raised when bounded open-series term discovery reaches a limit. + + A partial open-series sum is not returned: silently dropping terms would + turn a mathematically defined expansion into an uncontrolled truncation. + ``reason`` is one of the explicit enumeration limits, including + ``"max_corridor_edges"`` for corridor discovery. + """ + + def __init__(self, reason: str, limit: float, observed: float): + self.reason = reason + self.limit = limit + self.observed = observed + super().__init__( + f"open loop-series enumeration exceeded {reason}={limit!r} " + f"(observed {observed!r}); increase the limit or lower the " + "edge cutoff" + ) + + +@dataclass(frozen=True) +class OpenLoopObservableTerm: + """An observable term with its physical support made explicit. + + This is a small convenience wrapper for callers building terms from + :class:`pepsy.tensors.Fermion`. The operator should be the native dense + or Symmray array returned by methods such as ``fermion.observable(...)`` + or ``fermion.operator_term(...)``. A bare ``Fermion`` helper is not an + observable: it does not identify either the support or the local + operator, and is therefore rejected by the measurement APIs. + """ + + where: Any + operator: Any + label: Any = None + + +@dataclass +class OpenLoopSeriesDiagnostic: + """Geometry and contraction-cost report for open-series measurement. + + ``supports`` is keyed by the physical support tuple. Each value contains + the selected route, discovered terms, corridor/path diagnostics, and + per-term cost records. This object is intentionally numerical-data free: + it can be cached and passed back to + :func:`compute_local_expectation_open_loop_series` to reuse the discovered + geometry without re-enumerating it. + """ + + supports: dict[tuple[Any, ...], dict[str, Any]] + total_flops_log10: float | None = None + peak_memory_log2: float | None = None + cache_hits: int = 0 + + @property + def routes(self) -> dict[tuple[Any, ...], str]: + """Return the selected route for every support.""" + return { + support: record.get("route", "unknown") + for support, record in self.supports.items() + } + + def for_support(self, where): + """Return the cached report for one support.""" + support = tuple(where) + try: + return self.supports[support] + except KeyError as exc: + raise KeyError(f"no diagnostic was built for support {support!r}") from exc + + +@dataclass +class OpenLoopSeriesDiagnosticCache: + """Cache geometry and cost diagnostics for one TN topology. + + The cache deliberately does not retain observable values. Consequently + it is safe to reuse for different native Fermion operators with the same + physical rank and support, while keeping BP-message and gate data out of + the cache's ownership. + """ + + diagnostics_by_key: dict[Any, OpenLoopSeriesDiagnostic] = field( + default_factory=dict + ) + _topology_signature: Any = field(default=None, init=False, repr=False) + + def _check_topology(self, tn) -> None: + signature = LoopSeriesCache._signature(tn) + if self._topology_signature is None: + self._topology_signature = signature + elif self._topology_signature != signature: + raise ValueError( + "OpenLoopSeriesDiagnosticCache belongs to a different " + "tensor-network topology or tensor-id layout; create a fresh " + "cache" + ) + + def get(self, tn, key): + """Return a diagnostic for ``key`` or ``None`` when absent.""" + self._check_topology(tn) + return self.diagnostics_by_key.get(key) + + def put(self, tn, key, diagnostic): + """Store and return ``diagnostic`` after checking the topology.""" + self._check_topology(tn) + self.diagnostics_by_key[key] = diagnostic + return diagnostic + + @dataclass(frozen=True) class LoopSeriesTerm: """One connected loop-series excitation. @@ -100,6 +217,127 @@ def weight(self) -> int: return self.degree +@dataclass(frozen=True) +class _OpenEnumerationLimits: + """Validated limits for lazy open-series geometry discovery.""" + + max_terms: int | None = None + max_enumeration_time: float | None = None + max_enumeration_memory: int | None = None + + @classmethod + def validate( + cls, + *, + max_terms=None, + max_enumeration_time=None, + max_enumeration_memory=None, + ): + if max_terms is not None: + if not isinstance(max_terms, (int, np.integer)) or max_terms < 0: + raise ValueError("max_terms must be a non-negative integer or None") + max_terms = int(max_terms) + if max_enumeration_time is not None: + if ( + not isinstance( + max_enumeration_time, + (int, float, np.integer, np.floating), + ) + or not np.isfinite(max_enumeration_time) + or max_enumeration_time <= 0 + ): + raise ValueError( + "max_enumeration_time must be a finite positive number " + "or None" + ) + max_enumeration_time = float(max_enumeration_time) + if max_enumeration_memory is not None: + if ( + not isinstance(max_enumeration_memory, (int, np.integer)) + or max_enumeration_memory <= 0 + ): + raise ValueError( + "max_enumeration_memory must be a positive byte count " + "or None" + ) + max_enumeration_memory = int(max_enumeration_memory) + return cls( + max_terms=max_terms, + max_enumeration_time=max_enumeration_time, + max_enumeration_memory=max_enumeration_memory, + ) + + +class _OpenEnumerationGuard: + """Check lazy enumeration limits without changing term semantics.""" + + def __init__(self, limits: _OpenEnumerationLimits): + self.limits = limits + self.started = time.perf_counter() + self.emitted = 0 + self.estimated_memory = 0 + + @staticmethod + def _term_memory(term: LoopSeriesTerm) -> int: + # This is deliberately conservative bookkeeping for Python-side + # geometry, not a claim about tensor contraction memory. + return ( + sys.getsizeof(term) + + sys.getsizeof(term.edges) + + sys.getsizeof(term.tids) + + sum(sys.getsizeof(edge) for edge in term.edges) + + sum(sys.getsizeof(tid) for tid in term.tids) + ) + + def check(self): + elapsed = time.perf_counter() - self.started + limit = self.limits.max_enumeration_time + if limit is not None and elapsed >= limit: + raise OpenLoopEnumerationLimitError( + "max_enumeration_time", limit, elapsed + ) + + def accept(self, term: LoopSeriesTerm): + self.check() + if ( + self.limits.max_terms is not None + and self.emitted >= self.limits.max_terms + ): + raise OpenLoopEnumerationLimitError( + "max_terms", self.limits.max_terms, self.emitted + 1 + ) + term_memory = self._term_memory(term) + if ( + self.limits.max_enumeration_memory is not None + and self.estimated_memory + term_memory + > self.limits.max_enumeration_memory + ): + raise OpenLoopEnumerationLimitError( + "max_enumeration_memory", + self.limits.max_enumeration_memory, + self.estimated_memory + term_memory, + ) + self.emitted += 1 + self.estimated_memory += term_memory + + def diagnostics(self): + return { + "terms": self.emitted, + "elapsed_seconds": time.perf_counter() - self.started, + "estimated_memory_bytes": self.estimated_memory, + } + + +@dataclass(frozen=True) +class _CorridorPath: + """One weighted shortest path retained by corridor discovery.""" + + edges: tuple[Any, ...] + vertices: tuple[Any, ...] + cost: float + coordinates: tuple[Any, ...] = () + + @dataclass class LoopSeriesCache: """Cache edge-loop geometry for a fixed pairwise tensor-network topology.""" @@ -184,14 +422,51 @@ def terms_for( allowed_tids, excluded_edges=(), ) -> tuple[LoopSeriesTerm, ...]: - """Return open generalized-loop terms for one rho support.""" + """Return open generalized-loop terms for one rho support. + + This eager compatibility method is retained for callers that inspect + the geometry directly. The public open-series contractions use + :meth:`iter_terms_for` so terms are generated and consumed lazily. + """ + return tuple( + self.iter_terms_for( + tn, + max_degree, + allowed_tids, + excluded_edges=excluded_edges, + ) + ) + + def iter_terms_for( + self, + tn, + max_degree: int, + allowed_tids, + excluded_edges=(), + *, + max_terms: int | None = None, + max_enumeration_time: float | None = None, + max_enumeration_memory: int | None = None, + ): + """Yield open terms lazily for one rho support. + + Newly discovered terms are streamed to the caller and retained in the + cache only after the complete discovery finishes. This means a + bounded call can stop before a large configuration set is materialized + while preserving the old eager cache behavior for completed calls. + """ self._check_topology(tn) max_degree = _validate_nonnegative_degree(max_degree) allowed_tids = frozenset(allowed_tids) excluded_edges = frozenset(excluded_edges) + limits = _OpenEnumerationLimits.validate( + max_terms=max_terms, + max_enumeration_time=max_enumeration_time, + max_enumeration_memory=max_enumeration_memory, + ) key = (max_degree, allowed_tids, excluded_edges) try: - return self.terms_by_key[key] + cached = self.terms_by_key[key] except KeyError: larger_keys = [ known_key @@ -201,18 +476,34 @@ def terms_for( ] if larger_keys: larger = self.terms_by_key[min(larger_keys)] - terms = tuple( + cached = tuple( term for term in larger if term.degree <= max_degree ) else: - terms = _enumerate_open_edge_loops( - tn, - max_degree, - allowed_tids=allowed_tids, - excluded_edges=excluded_edges, - ) - self.terms_by_key[key] = terms - return terms + cached = None + + if cached is not None: + guard = _OpenEnumerationGuard(limits) + for term in cached: + guard.accept(term) + yield term + return + + discovered = [] + try: + for term in _iter_open_edge_loops( + tn, + max_degree, + allowed_tids=allowed_tids, + excluded_edges=excluded_edges, + limits=limits, + ): + discovered.append(term) + yield term + except OpenLoopEnumerationLimitError: + raise + else: + self.terms_by_key[key] = tuple(discovered) @dataclass @@ -476,13 +767,752 @@ def _open_term_from_edges( ) -def _enumerate_open_edge_loops( +def _iter_open_support_paths(tn, edges, allowed_tids, max_degree, guard=None): + """Yield simple support-connecting paths in increasing length order.""" + if len(allowed_tids) < 2 or max_degree < 1: + return + + adjacency: dict[Any, list[tuple[Any, Any]]] = {} + for index, left, right in edges: + adjacency.setdefault(left, []).append((right, index)) + adjacency.setdefault(right, []).append((left, index)) + for neighbors in adjacency.values(): + neighbors.sort(key=lambda item: (repr(item[0]), repr(item[1]))) + + support = tuple(sorted(allowed_tids, key=repr)) + queue = [] + serial = 0 + for source_pos, source in enumerate(support): + for target in support[source_pos + 1 :]: + heapq.heappush( + queue, + (0, serial, source, target, source, frozenset((source,)), ()), + ) + serial += 1 + + seen = set() + while queue: + if guard is not None: + guard.check() + length, _, source, target, current, visited, path_edges = heapq.heappop( + queue + ) + if current == target and path_edges: + canonical = tuple(sorted(path_edges, key=repr)) + if canonical not in seen: + seen.add(canonical) + records = _edge_records(tn) + tids = set() + for edge in canonical: + tids.update(records[edge]) + yield LoopSeriesTerm(canonical, frozenset(tids)) + continue + if length >= max_degree: + continue + for neighbor, edge in adjacency.get(current, ()): + if neighbor in visited: + continue + heapq.heappush( + queue, + ( + length + 1, + serial, + source, + target, + neighbor, + visited | {neighbor}, + path_edges + (edge,), + ), + ) + serial += 1 + + +def _validate_corridor_options( + *, + corridor_width, + max_path_candidates, + loop_decoration_size, + corridor_segment_length, + loop_radius, + max_loop_clusters_per_segment, + max_corridor_edges, + corridor_max_bond, +): + """Validate bounded path/corridor controls.""" + if corridor_width is not None: + if ( + not isinstance(corridor_width, (int, np.integer)) + or corridor_width < 0 + ): + raise ValueError("corridor_width must be a non-negative integer") + corridor_width = int(corridor_width) + for name, value in ( + ("max_path_candidates", max_path_candidates), + ("loop_decoration_size", loop_decoration_size), + ("corridor_segment_length", corridor_segment_length), + ("max_loop_clusters_per_segment", max_loop_clusters_per_segment), + ): + if not isinstance(value, (int, np.integer)) or value < 1: + raise ValueError(f"{name} must be a positive integer") + value = int(value) + if name == "max_path_candidates": + max_path_candidates = value + elif name == "loop_decoration_size": + loop_decoration_size = value + elif name == "corridor_segment_length": + corridor_segment_length = value + else: + max_loop_clusters_per_segment = value + if loop_radius is None: + loop_radius = max(1, corridor_width or 1) + elif not isinstance(loop_radius, (int, np.integer)) or loop_radius < 1: + raise ValueError("loop_radius must be a positive integer or None") + else: + loop_radius = int(loop_radius) + if max_corridor_edges is not None: + if ( + not isinstance(max_corridor_edges, (int, np.integer)) + or max_corridor_edges < 1 + ): + raise ValueError( + "max_corridor_edges must be a positive integer or None" + ) + max_corridor_edges = int(max_corridor_edges) + if corridor_max_bond is not None: + if corridor_width is None: + raise ValueError( + "corridor_max_bond requires corridor_width" + ) + if ( + not isinstance(corridor_max_bond, (int, np.integer)) + or corridor_max_bond < 1 + ): + raise ValueError( + "corridor_max_bond must be a positive integer or None" + ) + corridor_max_bond = int(corridor_max_bond) + return { + "corridor_width": corridor_width, + "max_path_candidates": max_path_candidates, + "loop_decoration_size": loop_decoration_size, + "corridor_segment_length": corridor_segment_length, + "loop_radius": loop_radius, + "max_loop_clusters_per_segment": max_loop_clusters_per_segment, + "max_corridor_edges": max_corridor_edges, + "corridor_max_bond": corridor_max_bond, + } + + +def _corridor_adjacency(tn, edges, edge_weights=None): + """Build deterministic tensor-graph adjacency for corridor search.""" + weights = {} if edge_weights is None else dict(edge_weights) + adjacency = {} + records = {} + for index, left, right in edges: + weight = weights.get(index, 1.0) + if not isinstance(weight, (int, float, np.integer, np.floating)): + raise TypeError(f"path edge weight for {index!r} must be real") + if not np.isfinite(weight) or weight <= 0: + raise ValueError( + f"path edge weight for {index!r} must be finite and positive" + ) + weight = float(weight) + records[index] = (left, right) + adjacency.setdefault(left, []).append((right, index, weight)) + adjacency.setdefault(right, []).append((left, index, weight)) + for neighbors in adjacency.values(): + neighbors.sort(key=lambda item: (item[2], repr(item[0]), repr(item[1]))) + return adjacency, records + + +def _weighted_shortest_distances(adjacency, target, guard=None): + distances = {target: 0.0} + pending = [(0.0, 0, target)] + serial = 1 + while pending: + if guard is not None: + guard.check() + distance, _, current = heapq.heappop(pending) + if distance > distances[current] + 1e-12: + continue + for neighbor, _, weight in adjacency.get(current, ()): + candidate = distance + weight + if candidate + 1e-12 >= distances.get(neighbor, np.inf): + continue + distances[neighbor] = candidate + heapq.heappush(pending, (candidate, serial, neighbor)) + serial += 1 + return distances + + +def _grid_corridor_context(tn): + """Return lazy coordinate-neighbor access for rectangular PEPS graphs.""" + if not all(hasattr(tn, name) for name in ("Lx", "Ly", "has_site")): + return None + if not callable(getattr(tn, "site_tag", None)): + return None + cyclic_x = bool(tn.is_cyclic_x()) if hasattr(tn, "is_cyclic_x") else False + cyclic_y = bool(tn.is_cyclic_y()) if hasattr(tn, "is_cyclic_y") else False + tid_cache = {} + + def tid_at(coo): + if coo in tid_cache: + return tid_cache[coo] + if not tn.has_site(coo): + return None + tids = tuple( + tn._get_tids_from_tags([tn.site_tag(coo)], "any") + ) + if len(tids) != 1: + return None + tid_cache[coo] = tids[0] + return tids[0] + + def neighbors(coo): + x, y = coo + candidates = ((x - 1, y), (x + 1, y), (x, y - 1), (x, y + 1)) + seen = set() + for nx, ny in candidates: + if nx < 0 or nx >= tn.Lx: + if not cyclic_x: + continue + nx %= tn.Lx + if ny < 0 or ny >= tn.Ly: + if not cyclic_y: + continue + ny %= tn.Ly + neighbor = (nx, ny) + if neighbor in seen or not tn.has_site(neighbor): + continue + seen.add(neighbor) + neighbor_tid = tid_at(neighbor) + if neighbor_tid is not None: + yield neighbor, neighbor_tid + + def edge_between(tid, neighbor_tid): + left_inds = set(tn.tensor_map[tid].inds) + right_inds = set(tn.tensor_map[neighbor_tid].inds) + shared = tuple(left_inds & right_inds) + for index in shared: + if len(tn.ind_map[index]) == 2: + return index + return None + + def axis_distance(left, right, size, cyclic): + distance = abs(left - right) + return min(distance, size - distance) if cyclic else distance + + def distance(left, right): + return axis_distance(left[0], right[0], tn.Lx, cyclic_x) + axis_distance( + left[1], right[1], tn.Ly, cyclic_y + ) + + return { + "tid_at": tid_at, + "neighbors": neighbors, + "edge_between": edge_between, + "distance": distance, + } + + +def _discover_grid_corridor_paths( + tn, + support_coos, + *, + excluded_edges=(), + corridor_width=2, + max_path_candidates=8, + corridor_segment_length=32, + max_corridor_edges=100_000, + guard=None, +): + """Discover corridor paths lazily on a rectangular lattice.""" + context = _grid_corridor_context(tn) + if context is None: + return None + excluded_edges = frozenset(excluded_edges) + normalized_coos = [] + for coo in support_coos: + if not isinstance(coo, (tuple, list)) or len(coo) != 2: + return None + normalized_coos.append(tuple(coo)) + support_coos = tuple(dict.fromkeys(normalized_coos)) + if len(support_coos) < 2: + return (), frozenset(), { + "path_count": 0, + "path_lengths": (), + "shortest_path_length": None, + "corridor_vertices": 0, + "corridor_edges": 0, + } + + paths = [] + seen_paths = set() + beam_width = max(8, 4 * max_path_candidates) + for source_pos, source in enumerate(support_coos[:-1]): + for target in support_coos[source_pos + 1 :]: + if guard is not None: + guard.check() + if context["tid_at"](source) is None or context["tid_at"](target) is None: + continue + target_distance = context["distance"](source, target) + frontier = [(source, (source,), (), frozenset((source,)))] + found = [] + while frontier and len(found) < max_path_candidates: + if guard is not None: + guard.check() + next_frontier = [] + for current, coordinates, path_edges, visited in frontier: + current_distance = context["distance"](current, target) + if current == target: + edge_tuple = tuple(sorted(path_edges, key=repr)) + if edge_tuple not in seen_paths: + seen_paths.add(edge_tuple) + found.append( + _CorridorPath( + edge_tuple, + tuple( + context["tid_at"](coo) + for coo in coordinates + ), + float(len(path_edges)), + coordinates, + ) + ) + continue + if current_distance >= target_distance and current != source: + continue + for neighbor, neighbor_tid in context["neighbors"](current): + if neighbor in visited: + continue + edge = context["edge_between"]( + context["tid_at"](current), + neighbor_tid, + ) + if edge is None or edge in excluded_edges: + continue + if context["distance"](neighbor, target) != current_distance - 1: + continue + next_frontier.append( + ( + neighbor, + coordinates + (neighbor,), + path_edges + (edge,), + visited | {neighbor}, + ) + ) + next_frontier.sort( + key=lambda state: ( + tuple(map(repr, state[2])), + state[0], + ) + ) + frontier = next_frontier[:beam_width] + paths.extend(found) + + paths.sort(key=lambda path: (path.cost, tuple(map(repr, path.edges)))) + paths = paths[:max_path_candidates] + corridor_coos = set() + for path in paths: + corridor_coos.update(path.coordinates) + pending = deque((coo, 0) for coo in corridor_coos) + while pending: + if guard is not None: + guard.check() + coo, distance = pending.popleft() + if distance >= corridor_width: + continue + for neighbor, _ in context["neighbors"](coo): + if neighbor in corridor_coos: + continue + corridor_coos.add(neighbor) + pending.append((neighbor, distance + 1)) + + corridor_edges = set() + corridor_tids = { + context["tid_at"](coo) for coo in corridor_coos + } + for coo in corridor_coos: + for neighbor, neighbor_tid in context["neighbors"](coo): + if neighbor not in corridor_coos: + continue + edge = context["edge_between"]( + context["tid_at"](coo), + neighbor_tid, + ) + if edge is not None and edge not in excluded_edges: + corridor_edges.add(edge) + corridor_edges = frozenset(corridor_edges) + if ( + max_corridor_edges is not None + and len(corridor_edges) > max_corridor_edges + ): + raise OpenLoopEnumerationLimitError( + "max_corridor_edges", + max_corridor_edges, + len(corridor_edges), + ) + return tuple(paths), corridor_edges, { + "path_count": len(paths), + "path_lengths": tuple(len(path.edges) for path in paths), + "shortest_path_length": min( + (len(path.edges) for path in paths), + default=None, + ), + "path_costs": tuple(path.cost for path in paths), + "corridor_vertices": len(corridor_tids), + "corridor_edges": len(corridor_edges), + "corridor_width": corridor_width, + "segment_length": corridor_segment_length, + "search_backend": "rectangular_grid", + } + + +def _discover_corridor_paths( + tn, + allowed_tids, + *, + support_coos=None, + excluded_edges=(), + corridor_width=2, + max_path_candidates=8, + corridor_segment_length=32, + max_corridor_edges=100_000, + edge_weights=None, + guard=None, +): + """Discover a bounded set of weighted shortest support paths. + + The search follows the shortest-path DAG produced by Dijkstra and keeps a + small beam at each distance layer. It therefore never explores arbitrary + simple paths on the full lattice. The returned corridor is the graph + neighbourhood of the retained paths. + """ + if support_coos is not None and edge_weights is None: + grid_result = _discover_grid_corridor_paths( + tn, + support_coos, + excluded_edges=excluded_edges, + corridor_width=corridor_width, + max_path_candidates=max_path_candidates, + corridor_segment_length=corridor_segment_length, + max_corridor_edges=max_corridor_edges, + guard=guard, + ) + if grid_result is not None: + return grid_result + + excluded_edges = frozenset(excluded_edges) + all_edges = _pairwise_edges(tn, norm="2norm") + search_edges = tuple( + edge for edge in all_edges if edge[0] not in excluded_edges + ) + adjacency, records = _corridor_adjacency(tn, search_edges, edge_weights) + support = tuple(sorted(frozenset(allowed_tids), key=repr)) + if len(support) < 2: + return (), frozenset(), { + "path_count": 0, + "path_lengths": (), + "shortest_path_length": None, + "corridor_vertices": 0, + "corridor_edges": 0, + } + + path_records = [] + seen_paths = set() + beam_width = max(8, 4 * max_path_candidates) + for source_pos, source in enumerate(support[:-1]): + for target in support[source_pos + 1 :]: + if guard is not None: + guard.check() + distances = _weighted_shortest_distances( + adjacency, + target, + guard=guard, + ) + if source not in distances: + continue + shortest_cost = distances[source] + frontier = [ + (source, (source,), (), 0.0, frozenset((source,))) + ] + found = [] + while frontier and len(found) < max_path_candidates: + if guard is not None: + guard.check() + next_frontier = [] + for current, vertices, path_edges, cost, visited in frontier: + if current == target: + canonical = tuple(sorted(path_edges, key=repr)) + if canonical not in seen_paths: + seen_paths.add(canonical) + found.append( + _CorridorPath(canonical, vertices, cost) + ) + continue + for neighbor, edge, weight in adjacency.get(current, ()): + if neighbor in visited: + continue + remaining = distances.get(neighbor) + if remaining is None: + continue + new_cost = cost + weight + if abs(new_cost + remaining - shortest_cost) > 1e-10: + continue + next_frontier.append( + ( + neighbor, + vertices + (neighbor,), + path_edges + (edge,), + new_cost, + visited | {neighbor}, + ) + ) + next_frontier.sort( + key=lambda state: ( + state[3], + tuple(map(repr, state[2])), + repr(state[0]), + ) + ) + frontier = next_frontier[:beam_width] + path_records.extend(found) + + path_records.sort( + key=lambda path: (path.cost, len(path.edges), tuple(map(repr, path.edges))) + ) + path_records = path_records[:max_path_candidates] + path_vertices = set() + for path in path_records: + path_vertices.update(path.vertices) + + corridor_vertices = set(path_vertices) + pending = deque((vertex, 0) for vertex in path_vertices) + while pending: + current, distance = pending.popleft() + if distance >= corridor_width: + continue + for neighbor, _, _ in adjacency.get(current, ()): + if neighbor in corridor_vertices: + continue + corridor_vertices.add(neighbor) + pending.append((neighbor, distance + 1)) + + corridor_edges = frozenset( + index + for index, left, right in search_edges + if left in corridor_vertices and right in corridor_vertices + ) + if ( + max_corridor_edges is not None + and len(corridor_edges) > max_corridor_edges + ): + raise OpenLoopEnumerationLimitError( + "max_corridor_edges", + max_corridor_edges, + len(corridor_edges), + ) + diagnostics = { + "path_count": len(path_records), + "path_lengths": tuple(len(path.edges) for path in path_records), + "shortest_path_length": ( + min((len(path.edges) for path in path_records), default=None) + ), + "path_costs": tuple(path.cost for path in path_records), + "corridor_vertices": len(corridor_vertices), + "corridor_edges": len(corridor_edges), + "corridor_width": corridor_width, + "segment_length": corridor_segment_length, + "search_backend": "weighted_graph", + } + return tuple(path_records), corridor_edges, diagnostics + + +def _iter_corridor_loop_clusters( + tn, + corridor_edges, + *, + max_size, + path_records, + segment_length, + loop_radius, + max_per_segment, + guard=None, +): + """Yield bounded simple-cycle decorations near path segments. + + The corridor route intentionally searches cycles rather than arbitrary + connected edge subsets. This keeps loop discovery bounded by the local + radius and decoration size, while the exact global route retains its old + generalized-loop semantics. + """ + records = _edge_records(tn) + edge_list = tuple(sorted(corridor_edges, key=repr)) + by_vertex = {} + for edge in edge_list: + left, right = records[edge] + by_vertex.setdefault(left, []).append((right, edge)) + by_vertex.setdefault(right, []).append((left, edge)) + for neighbors in by_vertex.values(): + neighbors.sort(key=lambda item: (repr(item[0]), repr(item[1]))) + + loop_terms = {} + for path in path_records: + if guard is not None: + guard.check() + anchors = path.vertices[::segment_length] + if path.vertices and path.vertices[-1] not in anchors: + anchors = (*anchors, path.vertices[-1]) + for anchor in anchors: + if guard is not None: + guard.check() + local_vertices = {anchor} + pending = deque(((anchor, 0),)) + while pending: + if guard is not None: + guard.check() + vertex, distance = pending.popleft() + if distance >= loop_radius: + continue + for neighbor, _ in by_vertex.get(vertex, ()): + if neighbor in local_vertices: + continue + local_vertices.add(neighbor) + pending.append((neighbor, distance + 1)) + + discovered = 0 + for start in sorted(local_vertices, key=repr): + if discovered >= max_per_segment: + break + stack = [(start, (start,), (), frozenset((start,)))] + while stack and discovered < max_per_segment: + if guard is not None: + guard.check() + current, vertices, path_edges, visited = stack.pop() + if ( + len(path_edges) >= 3 + and len(path_edges) <= max_size + ): + for neighbor, edge in by_vertex.get(current, ()): + if neighbor != start or edge in path_edges: + continue + canonical = tuple(sorted((*path_edges, edge), key=repr)) + if canonical not in loop_terms: + loop_terms[canonical] = LoopSeriesTerm( + canonical, + frozenset(vertices), + ) + discovered += 1 + break + if len(path_edges) >= max_size: + continue + for neighbor, edge in reversed( + by_vertex.get(current, ()) + ): + if neighbor not in local_vertices or neighbor in visited: + continue + if repr(neighbor) < repr(start): + continue + stack.append( + ( + neighbor, + vertices + (neighbor,), + path_edges + (edge,), + visited | {neighbor}, + ) + ) + + return tuple( + sorted( + loop_terms.values(), + key=lambda term: (term.degree, tuple(map(repr, term.edges))), + ) + ) + + +def _iter_corridor_open_terms( + tn, + path_records, + corridor_edges, + *, + allowed_tids, + excluded_edges, + total_edge_cutoff, + loop_decoration_size, + corridor_segment_length, + loop_radius, + max_loop_clusters_per_segment, + limits, + guard=None, +): + """Yield path-first corridor terms and one connected loop decoration.""" + guard = _OpenEnumerationGuard(limits) if guard is None else guard + seen = set() + path_terms = [] + for path in path_records: + if ( + total_edge_cutoff is not None + and len(path.edges) > total_edge_cutoff + ): + continue + term = LoopSeriesTerm(path.edges, frozenset(path.vertices)) + if term.edges in seen: + continue + guard.accept(term) + seen.add(term.edges) + path_terms.append(term) + yield term + + loop_terms = _iter_corridor_loop_clusters( + tn, + corridor_edges, + max_size=loop_decoration_size, + path_records=path_records, + segment_length=corridor_segment_length, + loop_radius=loop_radius, + max_per_segment=max_loop_clusters_per_segment, + guard=guard, + ) + for loop in loop_terms: + if ( + total_edge_cutoff is not None + and loop.degree > total_edge_cutoff + ): + continue + if loop.edges in seen: + continue + guard.accept(loop) + seen.add(loop.edges) + yield loop + + for path in path_terms: + for loop in loop_terms: + union = tuple(sorted(set(path.edges) | set(loop.edges), key=repr)) + if union == path.edges or union in seen: + continue + if ( + total_edge_cutoff is not None + and len(union) > total_edge_cutoff + ): + continue + term = _open_term_from_edges( + tn, + union, + allowed_tids=allowed_tids, + excluded_edges=excluded_edges, + ) + guard.accept(term) + seen.add(term.edges) + yield term + +def _iter_open_edge_loops( tn, max_degree: int, *, allowed_tids, excluded_edges=(), -) -> tuple[LoopSeriesTerm, ...]: + limits: _OpenEnumerationLimits | None = None, +): """Enumerate open and closed Q-edge configurations for a rho support. The cutoff is the number of excited Q edges. Every non-support tensor @@ -493,7 +1523,7 @@ def _enumerate_open_edge_loops( """ max_degree = _validate_nonnegative_degree(max_degree) if max_degree == 0: - return () + return excluded_edges = frozenset(excluded_edges) edges = tuple( @@ -503,6 +1533,26 @@ def _enumerate_open_edge_loops( ) max_degree = min(max_degree, len(edges)) allowed_tids = frozenset(allowed_tids) + limits = limits or _OpenEnumerationLimits.validate() + guard = _OpenEnumerationGuard(limits) + + # The first stream is deliberately path-first. This makes a bounded + # call useful for long-range observables: the smallest support-connecting + # configurations are seen before the much larger closed-loop tail. The + # exhaustive fallback below still retains every admissible path-plus-loop + # and disconnected-loop configuration when no limit is reached. + path_terms = set() + for term in _iter_open_support_paths( + tn, + edges, + allowed_tids, + max_degree, + guard=guard, + ): + guard.accept(term) + path_terms.add(term.edges) + yield term + remaining = Counter() for _, left, right in edges: remaining[left] += 1 @@ -510,8 +1560,6 @@ def _enumerate_open_edge_loops( degrees: Counter[Any] = Counter() selected = [] - terms = [] - def has_closed_dangling_vertex(): return any( remaining[tid] == 0 @@ -521,14 +1569,16 @@ def has_closed_dangling_vertex(): ) def visit(edge_pos, selected_count): + guard.check() if edge_pos == len(edges): if selected and not has_closed_dangling_vertex(): - terms.append( - LoopSeriesTerm( - tuple(sorted(selected, key=repr)), - frozenset(degrees), - ) + term = LoopSeriesTerm( + tuple(sorted(selected, key=repr)), + frozenset(degrees), ) + if term.edges not in path_terms: + guard.accept(term) + yield term return _, left, right = edges[edge_pos] @@ -536,14 +1586,14 @@ def visit(edge_pos, selected_count): remaining[right] -= 1 if not has_closed_dangling_vertex(): - visit(edge_pos + 1, selected_count) + yield from visit(edge_pos + 1, selected_count) if selected_count < max_degree: selected.append(edges[edge_pos][0]) degrees[left] += 1 degrees[right] += 1 if not has_closed_dangling_vertex(): - visit(edge_pos + 1, selected_count + 1) + yield from visit(edge_pos + 1, selected_count + 1) degrees[left] -= 1 degrees[right] -= 1 selected.pop() @@ -551,10 +1601,30 @@ def visit(edge_pos, selected_count): remaining[left] += 1 remaining[right] += 1 - visit(0, 0) + yield from visit(0, 0) + + +def _enumerate_open_edge_loops( + tn, + max_degree: int, + *, + allowed_tids, + excluded_edges=(), + limits: _OpenEnumerationLimits | None = None, +) -> tuple[LoopSeriesTerm, ...]: + """Eager compatibility wrapper around :func:`_iter_open_edge_loops`.""" return tuple( - sorted(terms, key=lambda term: (term.degree, tuple(map(repr, term.edges)))) + sorted( + _iter_open_edge_loops( + tn, + max_degree, + allowed_tids=allowed_tids, + excluded_edges=excluded_edges, + limits=limits, + ), + key=lambda term: (term.degree, tuple(map(repr, term.edges))), + ) ) @@ -638,12 +1708,113 @@ def _use_native_fermionic_cluster_open_route(bp, where): ) -def _connected_term_from_edges(tn, edges, *, tids=()): - """Validate and canonicalize an explicit edge-resolved term.""" - records = _edge_records(tn) - edges = tuple(edges) - if not edges: - raise ValueError("a loop-series term must contain at least one edge") +def _minimum_support_graph_distance(tn, where): + """Estimate the shortest pair distance without enumerating paths.""" + sites = tuple(where) + if len(sites) < 2: + return 0 + context = _grid_corridor_context(tn) + if context is not None and all( + isinstance(site, (tuple, list)) and len(site) == 2 for site in sites + ): + distances = [ + context["distance"](tuple(left), tuple(right)) + for left, right in combinations(sites, 2) + ] + return min(distances, default=0) + + tags = [tn.site_tag(site) for site in sites] + support_tids = tuple( + frozenset(tn._get_tids_from_tags([tag], "any")) for tag in tags + ) + adjacency, _ = _corridor_adjacency( + tn, + _pairwise_edges(tn, norm="2norm"), + edge_weights=None, + ) + best = np.inf + for left_pos, left_tids in enumerate(support_tids[:-1]): + targets = support_tids[left_pos + 1] + pending = deque((tid, 0) for tid in left_tids) + visited = set(left_tids) + while pending: + current, distance = pending.popleft() + if current in targets: + best = min(best, distance) + break + for neighbor, _, _ in adjacency.get(current, ()): + if neighbor in visited: + continue + visited.add(neighbor) + pending.append((neighbor, distance + 1)) + return int(best) if np.isfinite(best) else None + + +def _resolve_open_route( + bp, + where, + *, + mode, + corridor_width, + auto_corridor_distance, +): + """Select the safe open-series route before term discovery.""" + if mode not in {"exact", "corridor", "auto"}: + raise ValueError("mode must be 'exact', 'corridor', or 'auto'") + if auto_corridor_distance is None: + auto_corridor_distance = 32 + if ( + not isinstance(auto_corridor_distance, (int, np.integer)) + or auto_corridor_distance < 1 + ): + raise ValueError("auto_corridor_distance must be a positive integer or None") + native_cluster_route = _use_native_fermionic_cluster_open_route(bp, where) + distance = _minimum_support_graph_distance(bp.tn, where) + + if native_cluster_route: + if mode == "corridor" or (mode == "exact" and corridor_width is not None): + raise ValueError( + "corridor mode is for explicit dense/tree open terms; cyclic " + "native fermionic observables use cluster_size" + ) + return { + "route": "graded_cluster_compatible", + "corridor_width": None, + "native_cluster_route": True, + "support_distance": distance, + "auto_corridor_distance": int(auto_corridor_distance), + } + + use_corridor = mode == "corridor" + if mode == "auto": + use_corridor = ( + corridor_width is not None + or (distance is not None and distance > auto_corridor_distance) + ) + if mode == "exact" and corridor_width is not None: + # Preserve the pre-mode API: specifying corridor_width was already the + # explicit request to use the bounded route. + use_corridor = True + if use_corridor: + corridor_width = 2 if corridor_width is None else corridor_width + route = "corridor" + else: + route = "exact" + return { + "route": route, + "corridor_width": corridor_width if use_corridor else None, + "native_cluster_route": False, + "support_distance": distance, + "auto_corridor_distance": int(auto_corridor_distance), + } + + +def _connected_term_from_edges(tn, edges, *, tids=()): + """Validate and canonicalize an explicit edge-resolved term.""" + records = _edge_records(tn) + edges = tuple(edges) + if not edges: + raise ValueError("a loop-series term must contain at least one edge") if len(set(edges)) != len(edges): raise ValueError("a loop-series term cannot contain duplicate edges") unknown = set(edges).difference(records) @@ -1025,6 +2196,7 @@ def _get_d2_edge_partial_trace_excited( gate_as_operator=False, projector_index_order="bra-ket", fermionic_q=False, + index_namespace=None, ): """Build a D2 local RDM network with explicit P/Q edge choices. @@ -1055,12 +2227,24 @@ def _get_d2_edge_partial_trace_excited( boundary_inds = [] gate_index_map = {} + def make_index(role, tid, index): + if index_namespace is None: + import quimb.tensor as qtn + + return qtn.rand_uuid() + # Compressed contraction treats tuple-valued labels as structured + # index groups. Use a deterministic string instead so regional path + # reuse and boundary compression see an ordinary scalar index label. + return "__pepsy_open__" + repr( + (repr(index_namespace), role, repr(tid), repr(index)) + ) + for index, region_tids in stn.ind_map.items(): region_tids = tuple(region_tids) if index in bp.output_inds: if gate_as_operator and index in gate_inds: (tid,) = region_tids - kix = qtn.rand_uuid() + kix = make_index("gate-ket", tid, index) kixmaps[tid][index] = kix # ``tensor_network_gate_inds`` represents a gate with its # original physical labels on the first (bra/output) legs @@ -1072,13 +2256,13 @@ def _get_d2_edge_partial_trace_excited( (tid,) = region_tids bixmaps[tid][index] = partial_trace_map[index] elif index in exclude: - bix = qtn.rand_uuid() + bix = make_index("excluded-bra", region_tids[0], index) for tid in region_tids: bixmaps[tid][index] = bix elif index in stn._inner_inds: for tid in region_tids: - kix = qtn.rand_uuid() - bix = qtn.rand_uuid() + kix = make_index("ket", tid, index) + bix = make_index("bra", tid, index) kixmaps[tid][index] = kix bixmaps[tid][index] = bix if projector_index_order == "bra-ket": @@ -1087,8 +2271,8 @@ def _get_d2_edge_partial_trace_excited( projector_inds.setdefault(index, {})[tid] = (kix, bix) else: (tid,) = region_tids - kix = qtn.rand_uuid() - bix = qtn.rand_uuid() + kix = make_index("boundary-ket", tid, index) + bix = make_index("boundary-bra", tid, index) kixmaps[tid][index] = kix bixmaps[tid][index] = bix boundary_inds.append((index, tid)) @@ -1213,11 +2397,23 @@ def _validate_contraction_cost_limits( ) -def _contract_cost_record(tree): +def _contract_cost_record(tree, *, max_bond=None): """Extract the standard Cotengra log-cost diagnostics from a tree.""" + if max_bond is None: + flops = tree.total_flops(log=10) + peak = tree.peak_size(log=2) + else: + try: + flops = tree.total_flops(chi=max_bond, log=10) + peak = tree.peak_size(chi=max_bond, log=2) + except TypeError: + # Older cotengra trees expose only exact-tree diagnostics. Keep + # the budget conservative rather than failing the corridor route. + flops = tree.total_flops(log=10) + peak = tree.peak_size(log=2) return { - "flops_log10": float(tree.total_flops(log=10)), - "peak_memory_log2": float(tree.peak_size(log=2)), + "flops_log10": float(flops), + "peak_memory_log2": float(peak), } @@ -1228,6 +2424,9 @@ def _contract_with_cost_limits( contract_opts, max_flops_log10=None, max_peak_memory_log2=None, + path_cache=None, + path_cache_key=None, + compress_opts=None, ): """Contract ``network`` or return its cost record when over budget. @@ -1236,7 +2435,77 @@ def _contract_with_cost_limits( so path search is not repeated. ``peak_memory_log2`` follows Cotengra's convention: log2 of the largest concurrently live scalar tensor size. """ - if max_flops_log10 is None and max_peak_memory_log2 is None: + if compress_opts is not None: + if path_cache is not None and path_cache_key is not None: + # Exact contraction trees are not valid compressed-contraction + # paths, so regional path reuse is intentionally separate here. + path_cache_key = None + if "get" in contract_opts: + raise TypeError( + "contract_opts['get'] cannot be combined with compressed " + "corridor contraction" + ) + cost = None + if ( + max_flops_log10 is not None + or max_peak_memory_log2 is not None + ): + tree = network.contract(get="tree", optimize=optimize) + cost = _contract_cost_record( + tree, + max_bond=compress_opts.get("max_bond"), + ) + accepted = ( + ( + max_flops_log10 is None + or cost["flops_log10"] <= max_flops_log10 + ) + and ( + max_peak_memory_log2 is None + or cost["peak_memory_log2"] <= max_peak_memory_log2 + ) + ) + if not accepted: + return False, None, cost + compressed_opts = dict(compress_opts) + compressed_contract_opts = dict(contract_opts) + compressed_contract_opts.setdefault("output_inds", ()) + value = network.contract_compressed( + optimize, + **compressed_opts, + **compressed_contract_opts, + ) + return True, value, cost + + cached = None + if path_cache is not None and path_cache_key is not None: + cached = path_cache.get(path_cache_key) + + if cached is not None: + tree, cost = cached + accepted = ( + ( + max_flops_log10 is None + or cost["flops_log10"] <= max_flops_log10 + ) + and ( + max_peak_memory_log2 is None + or cost["peak_memory_log2"] <= max_peak_memory_log2 + ) + ) + if not accepted: + return False, None, cost + return ( + True, + network.contract(optimize=tree, **contract_opts), + cost if max_flops_log10 is not None or max_peak_memory_log2 is not None else None, + ) + + if ( + max_flops_log10 is None + and max_peak_memory_log2 is None + and path_cache is None + ): return ( True, network.contract(optimize=optimize, **contract_opts), @@ -1254,6 +2523,8 @@ def _contract_with_cost_limits( **contract_opts, ) cost = _contract_cost_record(tree) + if path_cache is not None and path_cache_key is not None: + path_cache[path_cache_key] = (tree, cost) accepted = ( ( max_flops_log10 is None @@ -1289,6 +2560,73 @@ def _term_sites(tn, where): return sites +def _resolve_open_observable_operator(operator): + """Resolve the small descriptor forms accepted by open measurement.""" + if isinstance(operator, OpenLoopObservableTerm): + return operator.operator + if isinstance(operator, dict) and "operator" in operator: + value = operator["operator"] + fermion = operator.get("fermion") + if fermion is not None and isinstance(value, str): + observable = getattr(fermion, "observable", None) + if not callable(observable): + raise TypeError( + "an observable descriptor with a named operator must " + "contain a Fermion-like object with observable(name)" + ) + return observable(value) + return value + if ( + isinstance(operator, (tuple, list)) + and len(operator) == 2 + and isinstance(operator[1], str) + and callable(getattr(operator[0], "observable", None)) + ): + # Convenient form: ``(fermion, "number")``. + return operator[0].observable(operator[1]) + if operator.__class__.__name__ == "Fermion": + raise TypeError( + "a Fermion helper is not itself an observable. Use, for example, " + "fermion.observable('number') or " + "fermion.operator_term(...), or pass " + "OpenLoopObservableTerm(where, operator)" + ) + return operator + + +def _normalize_open_observable_terms(terms): + """Normalize mappings and explicit ``(where, operator)`` term records.""" + if isinstance(terms, OpenLoopObservableTerm): + records = ((terms.where, terms.operator),) + elif hasattr(terms, "items"): + records = tuple(terms.items()) + else: + try: + records = tuple(terms) + except TypeError as exc: + raise TypeError( + "terms must be a mapping, OpenLoopObservableTerm, or an " + "iterable of (where, operator) pairs" + ) from exc + + if not records: + raise ValueError("terms must contain at least one operator") + normalized = [] + for item in records: + if isinstance(item, OpenLoopObservableTerm): + key = item.where + operator = item.operator + else: + try: + key, operator = item + except (TypeError, ValueError) as exc: + raise TypeError( + "observable terms must have form (where, operator)" + ) from exc + normalized.append((key, _resolve_open_observable_operator(operator))) + return tuple(normalized) + + def _partial_trace_loop_series( bp, where, @@ -1323,7 +2661,6 @@ def _partial_trace_loop_series( tids = frozenset(bp.tn._get_tids_from_tags(tags, "any")) if not tids: raise ValueError("where must contain at least one site in the network") - kix = [bp.tn.site_ind(coo) for coo in where] import quimb.tensor as qtn @@ -1725,58 +3062,59 @@ def _edge_series_terms_for_support(bp, tids, gloops, *, cache): return terms, inner_bonds -def _open_edge_series_terms_for_support(bp, tids, gloops, *, cache): - """Parse open rho terms, allowing dangling Q edges at ``tids``.""" - inner_bonds = frozenset(bp.tn._select_tids(tids).inner_inds()) - allowed_tids = frozenset(tids) - - if isinstance(gloops, (int, np.integer)): - cutoff = _validate_nonnegative_degree(gloops) - if cache is None: - terms = _enumerate_open_edge_loops( - bp.tn, - cutoff, - allowed_tids=allowed_tids, - excluded_edges=inner_bonds, +def _resolve_open_cutoffs( + legacy_gloops, + *, + edge_cutoff, + cluster_size, + native_cluster_route, +): + """Resolve explicit open-series and cluster cutoffs without ambiguity.""" + if edge_cutoff is not None and cluster_size is not None: + raise TypeError("pass only one of edge_cutoff and cluster_size") + + if legacy_gloops is not None: + if edge_cutoff is not None or cluster_size is not None: + raise TypeError( + "gloops is a legacy alias; do not combine it with " + "edge_cutoff or cluster_size" ) + if native_cluster_route: + cluster_size = legacy_gloops else: - terms = cache.terms_for( - bp.tn, - cutoff, - allowed_tids, - excluded_edges=inner_bonds, + edge_cutoff = legacy_gloops + + if native_cluster_route: + if edge_cutoff is not None: + raise ValueError( + "cyclic native fermionic open observables use the graded " + "cluster route; pass cluster_size instead of edge_cutoff" ) - return terms, inner_bonds + return None, cluster_size - if gloops is None: - max_degree = sum( - index not in inner_bonds - for index, _, _ in _pairwise_edges(bp.tn, norm="2norm") + if cluster_size is not None: + raise ValueError( + "cluster_size is only valid for the cyclic native fermionic " + "cluster route; pass edge_cutoff for explicit open-edge terms" ) - if cache is None: - terms = _enumerate_open_edge_loops( - bp.tn, - max_degree, - allowed_tids=allowed_tids, - excluded_edges=inner_bonds, - ) - else: - terms = cache.terms_for( - bp.tn, - max_degree, - allowed_tids, - excluded_edges=inner_bonds, - ) - return terms, inner_bonds + return edge_cutoff, None + +def _explicit_open_terms_iterator( + bp, + tids, + gloops, + *, + inner_bonds, +): + """Lazily validate explicit open-edge terms supplied by the caller.""" edge_labels = { index for index, _, _ in _pairwise_edges(bp.tn, norm="2norm") if index not in inner_bonds } - terms = [] seen = set() - for item in tuple(gloops): + for item in gloops: if isinstance(item, LoopSeriesTerm): edges = item.edges elif hasattr(item, "edges"): @@ -1798,17 +3136,521 @@ def _open_edge_series_terms_for_support(bp, tids, gloops, *, cache): term = _open_term_from_edges( bp.tn, edges, - allowed_tids=allowed_tids, + allowed_tids=tids, excluded_edges=inner_bonds, ) if term in seen: raise ValueError(f"duplicate open rho term: {term.edges!r}") seen.add(term) - terms.append(term) + yield term - return tuple( - sorted(terms, key=lambda term: (term.degree, tuple(map(repr, term.edges)))), - ), inner_bonds + +def _open_edge_series_terms_for_support( + bp, + tids, + edge_cutoff, + *, + cache, + max_terms=None, + max_enumeration_time=None, + max_enumeration_memory=None, + corridor_width=None, + max_path_candidates=8, + loop_decoration_size=4, + corridor_segment_length=32, + loop_radius=None, + max_loop_clusters_per_segment=8, + max_corridor_edges=100_000, + path_edge_weights=None, + support_coos=None, + corridor_info=None, +): + """Parse open rho terms, allowing dangling Q edges at ``tids``.""" + inner_bonds = frozenset(bp.tn._select_tids(tids).inner_inds()) + allowed_tids = frozenset(tids) + limits = _OpenEnumerationLimits.validate( + max_terms=max_terms, + max_enumeration_time=max_enumeration_time, + max_enumeration_memory=max_enumeration_memory, + ) + + if corridor_width is not None: + if edge_cutoff is not None and not isinstance( + edge_cutoff, (int, np.integer) + ): + raise TypeError( + "corridor mode accepts an integer edge_cutoff or None; " + "explicit edge subsets are not corridor paths" + ) + options = _validate_corridor_options( + corridor_width=corridor_width, + max_path_candidates=max_path_candidates, + loop_decoration_size=loop_decoration_size, + corridor_segment_length=corridor_segment_length, + loop_radius=loop_radius, + max_loop_clusters_per_segment=max_loop_clusters_per_segment, + max_corridor_edges=max_corridor_edges, + corridor_max_bond=None, + ) + total_edge_cutoff = ( + None + if edge_cutoff is None + else _validate_nonnegative_degree(edge_cutoff) + ) + + def corridor_terms(): + if limits.max_terms == 0: + raise OpenLoopEnumerationLimitError( + "max_terms", limits.max_terms, 1 + ) + guard = _OpenEnumerationGuard(limits) + paths, corridor_edges, diagnostics = _discover_corridor_paths( + bp.tn, + allowed_tids, + support_coos=support_coos, + excluded_edges=inner_bonds, + corridor_width=options["corridor_width"], + max_path_candidates=options["max_path_candidates"], + corridor_segment_length=options["corridor_segment_length"], + max_corridor_edges=options["max_corridor_edges"], + edge_weights=path_edge_weights, + guard=guard, + ) + diagnostics.update( + { + "loop_decoration_size": options[ + "loop_decoration_size" + ], + "loop_radius": options["loop_radius"], + "max_path_candidates": options["max_path_candidates"], + "max_loop_clusters_per_segment": options[ + "max_loop_clusters_per_segment" + ], + "max_edge_cutoff": total_edge_cutoff, + "approximation": "path_plus_connected_loop_decorations", + } + ) + if corridor_info is not None: + corridor_info.clear() + corridor_info.update(diagnostics) + if ( + total_edge_cutoff is not None + and diagnostics["shortest_path_length"] is not None + and diagnostics["shortest_path_length"] > total_edge_cutoff + ): + raise ValueError( + "edge_cutoff is smaller than the shortest corridor path: " + f"{total_edge_cutoff} < " + f"{diagnostics['shortest_path_length']}" + ) + yield from _iter_corridor_open_terms( + bp.tn, + paths, + corridor_edges, + allowed_tids=allowed_tids, + excluded_edges=inner_bonds, + total_edge_cutoff=total_edge_cutoff, + loop_decoration_size=options["loop_decoration_size"], + corridor_segment_length=options["corridor_segment_length"], + loop_radius=options["loop_radius"], + max_loop_clusters_per_segment=options[ + "max_loop_clusters_per_segment" + ], + limits=limits, + guard=guard, + ) + + return corridor_terms(), inner_bonds + + if isinstance(edge_cutoff, (int, np.integer)): + cutoff = _validate_nonnegative_degree(edge_cutoff) + if cache is None: + terms = _iter_open_edge_loops( + bp.tn, + cutoff, + allowed_tids=allowed_tids, + excluded_edges=inner_bonds, + limits=limits, + ) + else: + terms = cache.iter_terms_for( + bp.tn, + cutoff, + allowed_tids, + excluded_edges=inner_bonds, + max_terms=limits.max_terms, + max_enumeration_time=limits.max_enumeration_time, + max_enumeration_memory=limits.max_enumeration_memory, + ) + return terms, inner_bonds + + if edge_cutoff is None: + max_degree = sum( + index not in inner_bonds + for index, _, _ in _pairwise_edges(bp.tn, norm="2norm") + ) + if cache is None: + terms = _iter_open_edge_loops( + bp.tn, + max_degree, + allowed_tids=allowed_tids, + excluded_edges=inner_bonds, + limits=limits, + ) + else: + terms = cache.iter_terms_for( + bp.tn, + max_degree, + allowed_tids, + excluded_edges=inner_bonds, + max_terms=limits.max_terms, + max_enumeration_time=limits.max_enumeration_time, + max_enumeration_memory=limits.max_enumeration_memory, + ) + return terms, inner_bonds + + def limited_terms(): + guard = _OpenEnumerationGuard(limits) + for term in _explicit_open_terms_iterator( + bp, + allowed_tids, + edge_cutoff, + inner_bonds=inner_bonds, + ): + guard.accept(term) + yield term + + return limited_terms(), inner_bonds + + +def _log10_sum_costs(costs): + """Sum positive costs represented in base-10 logarithmic form.""" + values = [ + 10.0 ** float(cost["flops_log10"]) + for cost in costs + if cost is not None and np.isfinite(cost["flops_log10"]) + ] + if not values: + return None + return float(np.log10(sum(values))) + + +def _cost_within_limits(cost, max_flops_log10, max_peak_memory_log2): + """Return whether a diagnostic cost passes both optional budgets.""" + return ( + cost is not None + and ( + max_flops_log10 is None + or cost["flops_log10"] <= max_flops_log10 + ) + and ( + max_peak_memory_log2 is None + or cost["peak_memory_log2"] <= max_peak_memory_log2 + ) + ) + + +def _diagnose_network_cost( + network, + *, + optimize, + contract_opts, + max_bond=None, +): + """Build a contraction tree and return costs without contracting data.""" + if "get" in contract_opts: + raise TypeError( + "contract_opts['get'] cannot be combined with open-series " + "diagnostics; diagnostics always build a contraction tree" + ) + tree = network.contract(get="tree", optimize=optimize, **contract_opts) + return _contract_cost_record(tree, max_bond=max_bond) + + +def _open_diagnostic_key( + sites, + gate, + *, + route, + edge_cutoff, + cluster_size, + corridor_options, + max_terms, + max_enumeration_time, + max_enumeration_memory, + max_flops_log10, + max_peak_memory_log2, + path_edge_weights, +): + """Build a stable cache key for geometry and cost diagnostics.""" + try: + shape = tuple(ar.do("shape", gate)) + except Exception: + shape = repr(type(gate)) + try: + dtype = repr(ar.do("dtype", gate)) + except Exception: + dtype = repr(type(gate)) + return ( + tuple(sites), + route, + repr(edge_cutoff), + repr(cluster_size), + tuple(sorted((key, repr(value)) for key, value in corridor_options.items())), + max_terms, + max_enumeration_time, + max_enumeration_memory, + max_flops_log10, + max_peak_memory_log2, + repr(path_edge_weights), + shape, + dtype, + ) + + +def _diagnose_open_scalar_support( + bp, + where, + gate, + gloops, + *, + edge_cutoff, + cluster_size, + route_selection, + normalized, + optimize, + contract_opts, + cache, + max_flops_log10, + max_peak_memory_log2, + max_terms, + max_enumeration_time, + max_enumeration_memory, + corridor_options, + path_edge_weights, +): + """Diagnose one scalar support without contracting numerical values.""" + _align_symmray_d2bp_messages(bp) + bp.normalize_message_pairs() + bp.normalize_tensors() + tags = [bp.tn.site_tag(coo) for coo in where] + tids = frozenset(bp.tn._get_tids_from_tags(tags, "any")) + if not tids: + raise ValueError("where must contain at least one site in the network") + inner_bonds = frozenset(bp.tn._select_tids(tids).inner_inds()) + where_key = tuple(where) + fermionic_q = _uses_symmray(bp.tn) and _gate_needs_fermionic_open_q(gate) + route = route_selection["route"] + total_costs = [] + + if route == "graded_cluster_compatible": + from quimb.tensor.belief_propagation import gen_region_counts + + regions = tuple( + bp.tn.get_local_gloops( + tids=tids, + gloops=cluster_size, + grow_from="alldangle", + strict_size=False, + ) + ) + region_costs = {} + for region, _ in gen_region_counts(regions, autocomplete=True): + region = frozenset(region) + norm_cost = _diagnose_network_cost( + _get_d2_cluster_norm(bp, region), + optimize=optimize, + contract_opts=contract_opts, + ) + gate_cost = _diagnose_network_cost( + _get_d2_cluster_norm( + bp, + region, + gate=gate, + gate_inds=[bp.tn.site_ind(coo) for coo in where], + ), + optimize=optimize, + contract_opts=contract_opts, + ) + record = { + "norm": norm_cost, + "gate": gate_cost, + "flops_log10": max( + norm_cost["flops_log10"], gate_cost["flops_log10"] + ), + "peak_memory_log2": max( + norm_cost["peak_memory_log2"], + gate_cost["peak_memory_log2"], + ), + } + region_costs[region] = record + total_costs.extend((norm_cost, gate_cost)) + return { + "route": route, + "terms": (), + "requested_terms": (), + "term_costs": {}, + "skipped_terms": {}, + "cluster_region_costs": region_costs, + "corridor": {}, + "base_cost": None, + "fermionic_q": fermionic_q, + "support_distance": route_selection["support_distance"], + "total_flops_log10": _log10_sum_costs(total_costs), + "peak_memory_log2": max( + (cost["peak_memory_log2"] for cost in total_costs), + default=None, + ), + } + + corridor_info = {} + terms, inner_bonds = _open_edge_series_terms_for_support( + bp, + tids, + edge_cutoff, + cache=cache, + max_terms=max_terms, + max_enumeration_time=max_enumeration_time, + max_enumeration_memory=max_enumeration_memory, + corridor_width=( + corridor_options["corridor_width"] if route == "corridor" else None + ), + max_path_candidates=corridor_options["max_path_candidates"], + loop_decoration_size=corridor_options["loop_decoration_size"], + corridor_segment_length=corridor_options["corridor_segment_length"], + loop_radius=corridor_options["loop_radius"], + max_loop_clusters_per_segment=corridor_options[ + "max_loop_clusters_per_segment" + ], + max_corridor_edges=corridor_options["max_corridor_edges"], + path_edge_weights=path_edge_weights, + support_coos=where, + corridor_info=corridor_info, + ) + requested_terms = tuple(terms) + kix = [bp.tn.site_ind(coo) for coo in where] + term_costs = {} + skipped_terms = {} + compressed_max_bond = ( + corridor_options["corridor_max_bond"] if route == "corridor" else None + ) + compressed = compressed_max_bond is not None + for term in requested_terms: + region = frozenset((*tids, *term.tids)) + norm_network = _get_d2_edge_partial_trace_excited( + bp, + region, + excited_edges=term.edges, + exclude=inner_bonds, + projector_layout="open" if _uses_symmray(bp.tn) else "series", + fermionic_q=fermionic_q, + index_namespace=("diagnostic", "norm", where_key, term.edges), + ) + gate_network = _get_d2_edge_partial_trace_excited( + bp, + region, + excited_edges=term.edges, + exclude=inner_bonds, + gate=gate, + gate_inds=kix, + projector_layout="open" if _uses_symmray(bp.tn) else "series", + gate_as_operator=True, + fermionic_q=fermionic_q, + index_namespace=("diagnostic", "gate", where_key, term.edges), + ) + norm_cost = _diagnose_network_cost( + norm_network, + optimize=optimize, + contract_opts=contract_opts, + max_bond=compressed_max_bond if compressed else None, + ) + gate_cost = _diagnose_network_cost( + gate_network, + optimize=optimize, + contract_opts=contract_opts, + max_bond=compressed_max_bond if compressed else None, + ) + record = { + "norm": norm_cost, + "gate": gate_cost, + "flops_log10": max( + norm_cost["flops_log10"], gate_cost["flops_log10"] + ), + "peak_memory_log2": max( + norm_cost["peak_memory_log2"], + gate_cost["peak_memory_log2"], + ), + } + term_costs[term.edges] = record + total_costs.extend((norm_cost, gate_cost)) + if not _cost_within_limits( + norm_cost, max_flops_log10, max_peak_memory_log2 + ) or not _cost_within_limits( + gate_cost, max_flops_log10, max_peak_memory_log2 + ): + skipped_terms[term.edges] = record + + base_norm_network = _get_d2_edge_partial_trace_excited( + bp, + tids, + exclude=inner_bonds, + projector_layout="open" if _uses_symmray(bp.tn) else "series", + fermionic_q=fermionic_q, + index_namespace=("diagnostic", "base-norm", where_key), + ) + base_gate_network = _get_d2_edge_partial_trace_excited( + bp, + tids, + exclude=inner_bonds, + gate=gate, + gate_inds=kix, + projector_layout="open" if _uses_symmray(bp.tn) else "series", + gate_as_operator=True, + fermionic_q=fermionic_q, + index_namespace=("diagnostic", "base-gate", where_key), + ) + base_norm_cost = _diagnose_network_cost( + base_norm_network, + optimize=optimize, + contract_opts=contract_opts, + max_bond=compressed_max_bond if compressed else None, + ) + base_gate_cost = _diagnose_network_cost( + base_gate_network, + optimize=optimize, + contract_opts=contract_opts, + max_bond=compressed_max_bond if compressed else None, + ) + total_costs.extend((base_norm_cost, base_gate_cost)) + base_cost = { + "norm": base_norm_cost, + "gate": base_gate_cost, + "flops_log10": max( + base_norm_cost["flops_log10"], base_gate_cost["flops_log10"] + ), + "peak_memory_log2": max( + base_norm_cost["peak_memory_log2"], + base_gate_cost["peak_memory_log2"], + ), + } + return { + "route": route, + "terms": requested_terms, + "requested_terms": requested_terms, + "term_costs": term_costs, + "skipped_terms": skipped_terms, + "cluster_region_costs": {}, + "corridor": corridor_info, + "base_cost": base_cost, + "fermionic_q": fermionic_q, + "support_distance": route_selection["support_distance"], + "total_flops_log10": _log10_sum_costs(total_costs), + "peak_memory_log2": max( + (cost["peak_memory_log2"] for cost in total_costs), + default=None, + ), + "inner_bonds": inner_bonds, + } def _edge_series_suppression( @@ -1938,6 +3780,8 @@ def _partial_trace_open_loop_series( where, gloops, *, + edge_cutoff, + cluster_size, normalized, optimize, contract_opts, @@ -1945,6 +3789,21 @@ def _partial_trace_open_loop_series( info, max_flops_log10, max_peak_memory_log2, + max_terms, + max_enumeration_time, + max_enumeration_memory, + corridor_width, + max_path_candidates, + loop_decoration_size, + corridor_segment_length, + loop_radius, + max_loop_clusters_per_segment, + max_corridor_edges, + path_edge_weights, + corridor_max_bond, + mode, + auto_corridor_distance, + diagnostic_support, ): """Contract the explicit open-edge rho loop-series expansion.""" if bp.__class__.__name__ != "D2BP": @@ -1967,6 +3826,7 @@ def _partial_trace_open_loop_series( tids = frozenset(bp.tn._get_tids_from_tags(tags, "any")) if not tids: raise ValueError("where must contain at least one site in the network") + inner_bonds = frozenset(bp.tn._select_tids(tids).inner_inds()) kix = [bp.tn.site_ind(coo) for coo in where] import quimb.tensor as qtn @@ -1989,19 +3849,58 @@ def _partial_trace_open_loop_series( partial_trace_maps[map_key] = partial_trace_map bix = [partial_trace_map[index] for index in kix] output_inds = (*kix, *bix) - terms, inner_bonds = _open_edge_series_terms_for_support( - bp, - tids, - gloops, - cache=cache, - ) - requested_terms = terms max_flops_log10, max_peak_memory_log2 = _validate_contraction_cost_limits( max_flops_log10, max_peak_memory_log2, ) - if _use_native_fermionic_cluster_open_route(bp, where): + if diagnostic_support is not None and mode == "exact": + planned_route = diagnostic_support.get("route") + if planned_route == "corridor" and corridor_width is None: + mode = "corridor" + corridor_width = diagnostic_support.get("corridor_options", {}).get( + "corridor_width", 2 + ) + route_selection = _resolve_open_route( + bp, + where, + mode=mode, + corridor_width=corridor_width, + auto_corridor_distance=auto_corridor_distance, + ) + native_cluster_route = route_selection["native_cluster_route"] + if ( + diagnostic_support is not None + and gloops is None + and edge_cutoff is None + and cluster_size is None + ): + edge_cutoff = diagnostic_support.get("edge_cutoff") + cluster_size = diagnostic_support.get("cluster_size") + edge_cutoff, cluster_size = _resolve_open_cutoffs( + gloops, + edge_cutoff=edge_cutoff, + cluster_size=cluster_size, + native_cluster_route=native_cluster_route, + ) + + corridor_options = _validate_corridor_options( + corridor_width=route_selection["corridor_width"], + max_path_candidates=max_path_candidates, + loop_decoration_size=loop_decoration_size, + corridor_segment_length=corridor_segment_length, + loop_radius=loop_radius, + max_loop_clusters_per_segment=max_loop_clusters_per_segment, + max_corridor_edges=max_corridor_edges, + corridor_max_bond=corridor_max_bond, + ) + if native_cluster_route and route_selection["corridor_width"] is not None: + raise ValueError( + "corridor mode is for explicit dense/tree open terms; cyclic " + "native fermionic observables use cluster_size" + ) + + if native_cluster_route: # See the scalar counterpart below. This preserves a native # fermionic rho on cyclic graphs while avoiding the unsupported mixed # P/Q contraction path in Symmray. @@ -2009,7 +3908,7 @@ def _partial_trace_open_loop_series( rho = _partial_trace_loop_cluster( bp, where, - gloops, + cluster_size, combine="sum", normalized=normalized, autocomplete=True, @@ -2021,12 +3920,6 @@ def _partial_trace_open_loop_series( max_flops_log10=max_flops_log10, max_peak_memory_log2=max_peak_memory_log2, ) - term_families = { - term.edges: _open_term_family(bp.tn, term) - for term in requested_terms - } - family_counts = Counter(term_families.values()) - family_weights = {family: 0.0 for family in family_counts} if info is not None: cluster_region_costs = { (where_key, region): cost @@ -2040,8 +3933,8 @@ def _partial_trace_open_loop_series( "cluster_rho_skipped_terms", {} ).items() } - info["open_rho_requested_terms"] = requested_terms - info["open_rho_terms_list"] = requested_terms + info["open_rho_requested_terms"] = () + info["open_rho_terms_list"] = () info["open_rho_term_costs"] = dict( cluster_info.get("cluster_rho_term_costs", {}) ) @@ -2059,15 +3952,84 @@ def _partial_trace_open_loop_series( "open_rho_cluster_region_skipped_terms" ] = cluster_region_skipped info["open_rho_weights"] = {} - info["open_rho_term_families"] = term_families - info["open_rho_family_counts"] = dict(family_counts) - info["open_rho_family_weights"] = family_weights + info["open_rho_term_families"] = {} + info["open_rho_family_counts"] = {} + info["open_rho_family_weights"] = {} info["open_rho_base_weight"] = _rho_trace(rho) info["open_rho_support_tids"] = tids info["open_rho_excluded_edges"] = inner_bonds info["open_rho_native_route"] = "graded_cluster_compatible" + info["open_rho_edge_cutoff"] = None + info["open_rho_cluster_size"] = cluster_size + info["open_rho_enumeration_limits"] = { + "max_terms": max_terms, + "max_enumeration_time": max_enumeration_time, + "max_enumeration_memory": max_enumeration_memory, + } + info["open_rho_mode"] = mode + info["open_rho_support_distance"] = route_selection[ + "support_distance" + ] + info["open_rho_diagnostic"] = ( + None + if diagnostic_support is None + else dict(diagnostic_support) + ) return rho + corridor_info = ( + None + if info is None + else info.setdefault("open_rho_corridor", {}) + ) + if diagnostic_support is not None: + if diagnostic_support.get("route") != route_selection["route"]: + raise ValueError( + "the supplied open-series diagnostic does not match the " + "route selected for this rho support" + ) + terms = iter(diagnostic_support.get("terms", ())) + if corridor_info is not None: + corridor_info.clear() + corridor_info.update(diagnostic_support.get("corridor", {})) + inner_bonds = frozenset( + diagnostic_support.get( + "inner_bonds", bp.tn._select_tids(tids).inner_inds() + ) + ) + else: + terms, inner_bonds = _open_edge_series_terms_for_support( + bp, + tids, + edge_cutoff, + cache=cache, + max_terms=max_terms, + max_enumeration_time=max_enumeration_time, + max_enumeration_memory=max_enumeration_memory, + corridor_width=corridor_options["corridor_width"], + max_path_candidates=corridor_options["max_path_candidates"], + loop_decoration_size=corridor_options["loop_decoration_size"], + corridor_segment_length=corridor_options["corridor_segment_length"], + loop_radius=corridor_options["loop_radius"], + max_loop_clusters_per_segment=corridor_options[ + "max_loop_clusters_per_segment" + ], + max_corridor_edges=corridor_options["max_corridor_edges"], + path_edge_weights=path_edge_weights, + support_coos=where, + corridor_info=corridor_info, + ) + requested_terms = [] + + compressed_corridor_opts = None + if corridor_options["corridor_max_bond"] is not None: + compressed_corridor_opts = { + "max_bond": corridor_options["corridor_max_bond"], + "tree_gauge_distance": corridor_options[ + "corridor_segment_length" + ], + } + term_cache = {} if info is None else info.setdefault("open_rho_terms", {}) term_cost_cache = ( {} if info is None else info.setdefault("open_rho_term_costs", {}) @@ -2075,13 +4037,26 @@ def _partial_trace_open_loop_series( skipped_terms = ( {} if info is None else info.setdefault("open_rho_skipped_terms", {}) ) + path_cache = ( + {} + if info is None + else info.setdefault("open_rho_region_path_cache", {}) + ) rho_terms = {} accepted_terms = [] for term in terms: + requested_terms.append(term) if term.edges in skipped_terms: continue region = frozenset((*tids, *term.tids)) cache_key = (term.edges, region, where_key) + region_key = ( + "rho", + where_key, + tuple(sorted(region, key=repr)), + tuple(output_inds), + tuple(sorted(inner_bonds, key=repr)), + ) try: rho_e = term_cache[cache_key] except KeyError: @@ -2092,6 +4067,7 @@ def _partial_trace_open_loop_series( partial_trace_map=partial_trace_map, exclude=inner_bonds, projector_layout="open", + index_namespace=region_key, ) accepted, rho_e, cost = _contract_with_cost_limits( rho_network, @@ -2099,6 +4075,9 @@ def _partial_trace_open_loop_series( contract_opts={"output_inds": output_inds, **contract_opts}, max_flops_log10=max_flops_log10, max_peak_memory_log2=max_peak_memory_log2, + path_cache=path_cache, + path_cache_key=region_key, + compress_opts=compressed_corridor_opts, ) if not accepted: skipped_terms[term.edges] = cost @@ -2124,6 +4103,20 @@ def _partial_trace_open_loop_series( partial_trace_map=partial_trace_map, exclude=inner_bonds, projector_layout="open", + index_namespace=( + "rho-base", + where_key, + tuple(sorted(tids, key=repr)), + tuple(output_inds), + tuple(sorted(inner_bonds, key=repr)), + ), + ) + base_path_key = ( + "rho-base", + where_key, + tuple(sorted(tids, key=repr)), + tuple(output_inds), + tuple(sorted(inner_bonds, key=repr)), ) accepted, base, base_cost = _contract_with_cost_limits( base_network, @@ -2131,6 +4124,9 @@ def _partial_trace_open_loop_series( contract_opts={"output_inds": output_inds, **contract_opts}, max_flops_log10=max_flops_log10, max_peak_memory_log2=max_peak_memory_log2, + path_cache=path_cache, + path_cache_key=base_path_key, + compress_opts=compressed_corridor_opts, ) if not accepted: raise ValueError( @@ -2182,8 +4178,8 @@ def _partial_trace_open_loop_series( } info["open_rho_cluster_region_costs"] = {} info["open_rho_cluster_region_skipped_terms"] = {} - info["open_rho_requested_terms"] = requested_terms - info["open_rho_terms_list"] = terms + info["open_rho_requested_terms"] = tuple(requested_terms) + info["open_rho_terms_list"] = tuple(accepted_terms) info["open_rho_term_costs"] = dict(term_cost_cache) info["open_rho_skipped_terms"] = dict(skipped_terms) info["open_rho_cost_limits"] = { @@ -2195,8 +4191,26 @@ def _partial_trace_open_loop_series( info["open_rho_family_counts"] = dict(family_counts) info["open_rho_family_weights"] = family_weights info["open_rho_base_weight"] = _rho_trace(base) + info["open_rho_bp_baseline"] = base / _rho_trace(base) info["open_rho_support_tids"] = tids info["open_rho_excluded_edges"] = inner_bonds + info["open_rho_edge_cutoff"] = edge_cutoff + info["open_rho_cluster_size"] = None + info["open_rho_enumeration_limits"] = { + "max_terms": max_terms, + "max_enumeration_time": max_enumeration_time, + "max_enumeration_memory": max_enumeration_memory, + } + info["open_rho_corridor_options"] = dict(corridor_options) + info["open_rho_mode"] = mode + info["open_rho_support_distance"] = route_selection[ + "support_distance" + ] + info["open_rho_diagnostic"] = ( + None + if diagnostic_support is None + else dict(diagnostic_support) + ) return rho @@ -2330,6 +4344,8 @@ def _local_expectation_open_loop_series( gate, gloops, *, + edge_cutoff, + cluster_size, normalized, optimize, contract_opts, @@ -2337,6 +4353,21 @@ def _local_expectation_open_loop_series( info, max_flops_log10, max_peak_memory_log2, + max_terms, + max_enumeration_time, + max_enumeration_memory, + corridor_width, + max_path_candidates, + loop_decoration_size, + corridor_segment_length, + loop_radius, + max_loop_clusters_per_segment, + max_corridor_edges, + path_edge_weights, + corridor_max_bond, + mode, + auto_corridor_distance, + diagnostic_support, ): """Contract a gate through the explicit open-edge loop series.""" if normalized == "prod": @@ -2354,22 +4385,74 @@ def _local_expectation_open_loop_series( tids = frozenset(bp.tn._get_tids_from_tags(tags, "any")) if not tids: raise ValueError("where must contain at least one site in the network") + inner_bonds = frozenset(bp.tn._select_tids(tids).inner_inds()) kix = [bp.tn.site_ind(coo) for coo in where] - terms, inner_bonds = _open_edge_series_terms_for_support( - bp, - tids, - gloops, - cache=cache, - ) max_flops_log10, max_peak_memory_log2 = _validate_contraction_cost_limits( max_flops_log10, max_peak_memory_log2, ) where_key = tuple(where) fermionic_q = _uses_symmray(bp.tn) and _gate_needs_fermionic_open_q(gate) + if diagnostic_support is not None and mode == "exact": + planned_route = diagnostic_support.get("route") + if planned_route == "corridor" and corridor_width is None: + mode = "corridor" + corridor_width = diagnostic_support.get("corridor_options", {}).get( + "corridor_width", 2 + ) + route_selection = _resolve_open_route( + bp, + where, + mode=mode, + corridor_width=corridor_width, + auto_corridor_distance=auto_corridor_distance, + ) + if diagnostic_support is not None: + planned_route = diagnostic_support.get("route") + if planned_route is not None and planned_route != route_selection["route"]: + raise ValueError( + "the supplied open-series diagnostic was built for route " + f"{planned_route!r}, but the current measurement selected " + f"{route_selection['route']!r}" + ) + if planned_route == "corridor": + corridor_width = diagnostic_support["corridor_options"][ + "corridor_width" + ] + native_cluster_route = route_selection["native_cluster_route"] + if ( + diagnostic_support is not None + and gloops is None + and edge_cutoff is None + and cluster_size is None + ): + edge_cutoff = diagnostic_support.get("edge_cutoff") + cluster_size = diagnostic_support.get("cluster_size") + edge_cutoff, cluster_size = _resolve_open_cutoffs( + gloops, + edge_cutoff=edge_cutoff, + cluster_size=cluster_size, + native_cluster_route=native_cluster_route, + ) - if _use_native_fermionic_cluster_open_route(bp, where): + corridor_options = _validate_corridor_options( + corridor_width=route_selection["corridor_width"], + max_path_candidates=max_path_candidates, + loop_decoration_size=loop_decoration_size, + corridor_segment_length=corridor_segment_length, + loop_radius=loop_radius, + max_loop_clusters_per_segment=max_loop_clusters_per_segment, + max_corridor_edges=max_corridor_edges, + corridor_max_bond=corridor_max_bond, + ) + if native_cluster_route and corridor_width is not None: + raise ValueError( + "corridor mode is for explicit dense/tree open terms; cyclic " + "native fermionic observables use cluster_size" + ) + + if native_cluster_route: # The explicit open-edge decomposition is exact for dense networks # and for fermionic trees. On a cyclic native Symmray graph, however, # a mixed P/Q network can require a non-pairwise fermionic contraction @@ -2381,7 +4464,7 @@ def _local_expectation_open_loop_series( bp, where, gate, - gloops, + cluster_size, combine="sum", normalized=normalized, autocomplete=True, @@ -2393,10 +4476,6 @@ def _local_expectation_open_loop_series( max_flops_log10=max_flops_log10, max_peak_memory_log2=max_peak_memory_log2, ) - term_families = { - term.edges: _open_term_family(bp.tn, term) for term in terms - } - family_counts = Counter(term_families.values()) if info is not None: cluster_region_costs = { (where_key, region): cost @@ -2410,8 +4489,8 @@ def _local_expectation_open_loop_series( "cluster_scalar_skipped_terms", {} ).items() } - info["open_scalar_requested_terms"] = terms - info["open_scalar_terms"] = terms + info["open_scalar_requested_terms"] = () + info["open_scalar_terms"] = () info["open_scalar_skipped_terms"] = dict( cluster_info.get("cluster_scalar_skipped_terms", {}) ) @@ -2420,8 +4499,8 @@ def _local_expectation_open_loop_series( ) info["open_scalar_norm_weights"] = {} info["open_scalar_gate_terms"] = {} - info["open_scalar_term_families"] = term_families - info["open_scalar_family_counts"] = dict(family_counts) + info["open_scalar_term_families"] = {} + info["open_scalar_family_counts"] = {} info["open_scalar_family_weights"] = {} info["open_scalar_base_weight"] = normalization info["open_scalar_numerator"] = value * normalization @@ -2441,8 +4520,71 @@ def _local_expectation_open_loop_series( info[ "open_scalar_cluster_region_skipped_terms" ] = cluster_region_skipped + info["open_scalar_edge_cutoff"] = None + info["open_scalar_cluster_size"] = cluster_size + info["open_scalar_enumeration_limits"] = { + "max_terms": max_terms, + "max_enumeration_time": max_enumeration_time, + "max_enumeration_memory": max_enumeration_memory, + } + info["open_scalar_mode"] = mode + info["open_scalar_support_distance"] = route_selection[ + "support_distance" + ] + info["open_scalar_diagnostic"] = ( + None + if diagnostic_support is None + else dict(diagnostic_support) + ) return value, normalization + corridor_info = ( + None + if info is None + else info.setdefault("open_scalar_corridor", {}) + ) + if diagnostic_support is not None: + terms = iter(diagnostic_support.get("terms", ())) + if corridor_info is not None: + corridor_info.clear() + corridor_info.update(diagnostic_support.get("corridor", {})) + inner_bonds = frozenset( + diagnostic_support.get( + "inner_bonds", bp.tn._select_tids(tids).inner_inds() + ) + ) + else: + terms, inner_bonds = _open_edge_series_terms_for_support( + bp, + tids, + edge_cutoff, + cache=cache, + max_terms=max_terms, + max_enumeration_time=max_enumeration_time, + max_enumeration_memory=max_enumeration_memory, + corridor_width=corridor_options["corridor_width"], + max_path_candidates=corridor_options["max_path_candidates"], + loop_decoration_size=corridor_options["loop_decoration_size"], + corridor_segment_length=corridor_options["corridor_segment_length"], + loop_radius=corridor_options["loop_radius"], + max_loop_clusters_per_segment=corridor_options[ + "max_loop_clusters_per_segment" + ], + max_corridor_edges=corridor_options["max_corridor_edges"], + path_edge_weights=path_edge_weights, + support_coos=where, + corridor_info=corridor_info, + ) + + compressed_corridor_opts = None + if corridor_options["corridor_max_bond"] is not None: + compressed_corridor_opts = { + "max_bond": corridor_options["corridor_max_bond"], + "tree_gauge_distance": corridor_options[ + "corridor_segment_length" + ], + } + norm_cache = ( {} if info is None else info.setdefault("open_scalar_norm_terms", {}) ) @@ -2464,17 +4606,38 @@ def _local_expectation_open_loop_series( norm_terms = {} gate_terms = {} accepted_terms = [] + requested_terms = [] term_costs = {} if info is None else info.setdefault( "open_scalar_term_costs", {} ) skipped_terms = {} if info is None else info.setdefault( "open_scalar_skipped_terms", {} ) + path_cache = ( + {} + if info is None + else info.setdefault("open_scalar_region_path_cache", {}) + ) for term in terms: + requested_terms.append(term) if (where_key, term.edges) in skipped_terms: continue region = frozenset((*tids, *term.tids)) cache_key = (term.edges, region, where_key, fermionic_q) + norm_path_key = ( + "norm", + where_key, + tuple(sorted(region, key=repr)), + tuple(sorted(inner_bonds, key=repr)), + fermionic_q, + ) + gate_path_key = ( + "gate", + where_key, + tuple(sorted(region, key=repr)), + tuple(sorted(inner_bonds, key=repr)), + fermionic_q, + ) try: norm_e = norm_cache[cache_key] norm_cost = norm_cost_cache.get((where_key, term.edges)) @@ -2488,6 +4651,7 @@ def _local_expectation_open_loop_series( "open" if _uses_symmray(bp.tn) else "series" ), fermionic_q=fermionic_q, + index_namespace=("open-scalar", *norm_path_key), ) accepted, norm_e, norm_cost = _contract_with_cost_limits( norm_network, @@ -2495,6 +4659,9 @@ def _local_expectation_open_loop_series( contract_opts=contract_opts, max_flops_log10=max_flops_log10, max_peak_memory_log2=max_peak_memory_log2, + path_cache=path_cache, + path_cache_key=norm_path_key, + compress_opts=compressed_corridor_opts, ) if not accepted: skipped_terms[(where_key, term.edges)] = {"norm": norm_cost} @@ -2520,6 +4687,7 @@ def _local_expectation_open_loop_series( ), gate_as_operator=True, fermionic_q=fermionic_q, + index_namespace=("open-scalar", *gate_path_key), ) accepted, gate_e, gate_cost = _contract_with_cost_limits( gate_network, @@ -2527,6 +4695,9 @@ def _local_expectation_open_loop_series( contract_opts=contract_opts, max_flops_log10=max_flops_log10, max_peak_memory_log2=max_peak_memory_log2, + path_cache=path_cache, + path_cache_key=gate_path_key, + compress_opts=compressed_corridor_opts, ) if accepted: gate_cache[gate_cache_key] = gate_e @@ -2559,6 +4730,13 @@ def _local_expectation_open_loop_series( gate_terms[term.edges] = gate_e base_key = (where_key, tuple(kix), tids, inner_bonds, fermionic_q) + base_norm_path_key = ( + "base-norm", + where_key, + tuple(sorted(tids, key=repr)), + tuple(sorted(inner_bonds, key=repr)), + fermionic_q, + ) base_cache = ( {} if info is None else info.setdefault("open_scalar_base_terms", {}) ) @@ -2573,6 +4751,7 @@ def _local_expectation_open_loop_series( "open" if _uses_symmray(bp.tn) else "series" ), fermionic_q=fermionic_q, + index_namespace=("open-scalar", *base_norm_path_key), ) accepted, base_norm, base_cost = _contract_with_cost_limits( base_network, @@ -2580,6 +4759,9 @@ def _local_expectation_open_loop_series( contract_opts=contract_opts, max_flops_log10=max_flops_log10, max_peak_memory_log2=max_peak_memory_log2, + path_cache=path_cache, + path_cache_key=base_norm_path_key, + compress_opts=compressed_corridor_opts, ) if not accepted: raise ValueError( @@ -2601,6 +4783,21 @@ def _local_expectation_open_loop_series( ), gate_as_operator=True, fermionic_q=fermionic_q, + index_namespace=( + "open-scalar", + "base-gate", + where_key, + tuple(sorted(tids, key=repr)), + tuple(sorted(inner_bonds, key=repr)), + fermionic_q, + ), + ) + base_gate_path_key = ( + "base-gate", + where_key, + tuple(sorted(tids, key=repr)), + tuple(sorted(inner_bonds, key=repr)), + fermionic_q, ) accepted, base_value, base_gate_cost = _contract_with_cost_limits( base_gate_network, @@ -2608,6 +4805,9 @@ def _local_expectation_open_loop_series( contract_opts=contract_opts, max_flops_log10=max_flops_log10, max_peak_memory_log2=max_peak_memory_log2, + path_cache=path_cache, + path_cache_key=base_gate_path_key, + compress_opts=compressed_corridor_opts, ) if not accepted: raise ValueError( @@ -2638,7 +4838,7 @@ def _local_expectation_open_loop_series( ) for family in family_counts } - info["open_scalar_requested_terms"] = terms + info["open_scalar_requested_terms"] = tuple(requested_terms) info["open_scalar_terms"] = tuple(accepted_terms) info["open_scalar_edge_term_costs"] = dict(term_costs) info["open_scalar_edge_skipped_terms"] = dict(skipped_terms) @@ -2656,6 +4856,10 @@ def _local_expectation_open_loop_series( info["open_scalar_base_weight"] = base_norm info["open_scalar_numerator"] = raw_value info["open_scalar_denominator"] = norm + info["open_scalar_bp_baseline"] = base_value / base_norm + info["open_scalar_corridor_correction"] = value - ( + base_value / base_norm + ) info["open_scalar_excluded_edges"] = inner_bonds info["open_scalar_native_route"] = ( "graded_open_projectors" @@ -2663,6 +4867,23 @@ def _local_expectation_open_loop_series( else "dense_open_projectors" ) info["open_scalar_fermionic_q_phase"] = fermionic_q + info["open_scalar_edge_cutoff"] = edge_cutoff + info["open_scalar_cluster_size"] = None + info["open_scalar_enumeration_limits"] = { + "max_terms": max_terms, + "max_enumeration_time": max_enumeration_time, + "max_enumeration_memory": max_enumeration_memory, + } + info["open_scalar_corridor_options"] = dict(corridor_options) + info["open_scalar_mode"] = mode + info["open_scalar_support_distance"] = route_selection[ + "support_distance" + ] + info["open_scalar_diagnostic"] = ( + None + if diagnostic_support is None + else dict(diagnostic_support) + ) return value, norm @@ -3382,7 +5603,7 @@ def partial_trace_edge_loop_series_expand( def partial_trace_open_loop_series_sweep( tn, supports, - cutoffs, + cutoffs=None, *, messages=None, gauges=None, @@ -3402,6 +5623,22 @@ def partial_trace_open_loop_series_sweep( optimize: Any = "auto-hq", max_flops_log10: float | None = None, max_peak_memory_log2: float | None = None, + edge_cutoffs=None, + cluster_sizes=None, + max_terms: int | None = None, + max_enumeration_time: float | None = None, + max_enumeration_memory: int | None = None, + mode: str = "exact", + auto_corridor_distance: int | None = 32, + corridor_width: int | None = None, + max_path_candidates: int = 8, + loop_decoration_size: int = 4, + corridor_segment_length: int = 32, + loop_radius: int | None = None, + max_loop_clusters_per_segment: int = 8, + max_corridor_edges: int | None = 100_000, + path_edge_weights=None, + corridor_max_bond: int | None = None, contract_opts: dict[str, Any] | None = None, **bp_opts, ) -> OpenLoopSeriesSweepResult: @@ -3415,7 +5652,10 @@ def partial_trace_open_loop_series_sweep( The physical sites to retain for each rho. List and tuple site coordinates are normalized to hashable tuples. cutoffs : iterable of int or int - Maximum numbers of excited Q edges to evaluate, in order. + Legacy cutoff alias. Prefer ``edge_cutoffs`` for explicit edge terms + or ``cluster_sizes`` for the cyclic native fermionic route. + edge_cutoffs, cluster_sizes : iterable of int or int, optional + Route-specific cutoff sweep. Pass only one of these. messages, gauges, run_bp, ... BP controls matching :func:`partial_trace_open_loop_series_expand`. @@ -3432,6 +5672,22 @@ def partial_trace_open_loop_series_sweep( underlying open-rho expansion. For native fermionic PEPS, the returned rhos are diagnostics; evaluate operators through the graded scalar APIs. """ + if edge_cutoffs is not None and cluster_sizes is not None: + raise TypeError("pass only one of edge_cutoffs and cluster_sizes") + if edge_cutoffs is not None: + cutoffs = edge_cutoffs + cutoff_kind = "edge_cutoff" + elif cluster_sizes is not None: + cutoffs = cluster_sizes + cutoff_kind = "cluster_size" + else: + cutoff_kind = "legacy" + + if cutoffs is None: + raise TypeError( + "pass cutoffs, edge_cutoffs, or cluster_sizes" + ) + if isinstance(cutoffs, (int, np.integer)): cutoffs = (int(cutoffs),) else: @@ -3493,10 +5749,15 @@ def partial_trace_open_loop_series_sweep( support_rhos = {} support_diagnostics = {} for cutoff in cutoffs: + edge_cutoff = cutoff if cutoff_kind == "edge_cutoff" else None + cluster_size = cutoff if cutoff_kind == "cluster_size" else None + legacy_cutoff = cutoff if cutoff_kind == "legacy" else None rho = _partial_trace_open_loop_series( bp, support, - cutoff, + legacy_cutoff, + edge_cutoff=edge_cutoff, + cluster_size=cluster_size, normalized=normalized, optimize=optimize, max_flops_log10=max_flops_log10, @@ -3504,6 +5765,21 @@ def partial_trace_open_loop_series_sweep( contract_opts=contract_opts, cache=cache, info=support_info, + max_terms=max_terms, + max_enumeration_time=max_enumeration_time, + max_enumeration_memory=max_enumeration_memory, + corridor_width=corridor_width, + max_path_candidates=max_path_candidates, + loop_decoration_size=loop_decoration_size, + corridor_segment_length=corridor_segment_length, + loop_radius=loop_radius, + max_loop_clusters_per_segment=max_loop_clusters_per_segment, + max_corridor_edges=max_corridor_edges, + path_edge_weights=path_edge_weights, + corridor_max_bond=corridor_max_bond, + mode=mode, + auto_corridor_distance=auto_corridor_distance, + diagnostic_support=None, ) support_rhos[cutoff] = rho support_diagnostics[cutoff] = { @@ -3537,6 +5813,12 @@ def partial_trace_open_loop_series_sweep( "open_rho_cluster_region_skipped_terms" ] ), + "enumeration_limits": dict( + support_info["open_rho_enumeration_limits"] + ), + "corridor": dict( + support_info.get("open_rho_corridor", {}) + ), "cost_limits": dict(support_info["open_rho_cost_limits"]), } rhos[support_key] = support_rhos @@ -3578,13 +5860,30 @@ def partial_trace_open_loop_series_expand( optimize: Any = "auto-hq", max_flops_log10: float | None = None, max_peak_memory_log2: float | None = None, + edge_cutoff=None, + cluster_size=None, + max_terms: int | None = None, + max_enumeration_time: float | None = None, + max_enumeration_memory: int | None = None, + mode: str = "exact", + auto_corridor_distance: int | None = 32, + corridor_width: int | None = None, + max_path_candidates: int = 8, + loop_decoration_size: int = 4, + corridor_segment_length: int = 32, + loop_radius: int | None = None, + max_loop_clusters_per_segment: int = 8, + max_corridor_edges: int | None = 100_000, + path_edge_weights=None, + corridor_max_bond: int | None = None, + diagnostic: OpenLoopSeriesDiagnostic | None = None, info: dict[str, Any] | None = None, contract_opts: dict[str, Any] | None = None, **bp_opts, ): """Compute a long-range local rho from an explicit open-edge series. - The integer ``gloops`` cutoff counts excited ``Q`` virtual edges. A + ``edge_cutoff`` counts excited ``Q`` virtual edges. A retained edge subset may have degree one only at one of the selected physical rho sites; every other touched tensor must have either zero or at least two excited edges. This keeps the open excitation paths connecting @@ -3601,12 +5900,27 @@ def partial_trace_open_loop_series_expand( all retained terms have unit coefficient and are normalized only after summation. + On cyclic native fermionic networks, use ``cluster_size`` instead. The + native route is selected before edge geometry is enumerated and contracts + graded message-closed tensor regions rather than mixed open ``P/Q`` edge + networks. ``gloops`` remains accepted as a legacy alias, but new code + should use the route-specific parameter. + ``messages`` can be supplied from a previously converged D2BP run with ``run_bp=False``. This is the intended route for measuring many - long-range rho supports after one BP solve. ``gloops=None`` enumerates up + long-range rho supports after one BP solve. ``edge_cutoff=None`` enumerates up to the number of eligible pairwise virtual bonds and can be expensive; use an integer cutoff for practical calculations. + Set ``corridor_width`` to activate the large-separation approximation. + It keeps a bounded beam of weighted shortest support paths, inflates them + into a graph corridor, and adds only connected loop decorations sampled + near corridor segments. This route is deliberately approximate: it does + not enumerate disconnected products of distant loop clusters. Supply + ``corridor_max_bond`` to use compressed boundary contraction for long + corridors; ``corridor_segment_length`` controls its local compression + scale and loop sampling stride. + Parameters ---------- tn : TensorNetwork @@ -3614,7 +5928,46 @@ def partial_trace_open_loop_series_expand( where : sequence The physical sites to retain in the reduced density matrix. gloops : int or iterable, optional + Legacy alias for ``edge_cutoff`` on dense/tree networks and for + ``cluster_size`` on cyclic native fermionic networks. + edge_cutoff : int or iterable, optional Maximum number of excited Q edges, or explicit virtual-edge subsets. + cluster_size : int or iterable, optional + Tensor-region cutoff or explicit cluster regions for the cyclic native + fermionic compatibility route. + max_terms : int, optional + Hard limit on discovered explicit edge terms. Exceeding it raises + :class:`OpenLoopEnumerationLimitError`; partial sums are never + returned silently. + max_enumeration_time : float, optional + Maximum edge-geometry discovery time in seconds. + max_enumeration_memory : int, optional + Approximate Python geometry memory budget in bytes. + corridor_width : int, optional + Graph distance used to inflate retained shortest paths. ``None`` keeps + the exact global edge-series route; a non-negative integer activates + corridor mode. + max_path_candidates : int, optional + Beam size for retained weighted shortest paths. + loop_decoration_size : int, optional + Maximum Q-edge size of one connected loop decoration. + corridor_segment_length : int, optional + Path stride for local loop sampling and compressed boundary gauges. + loop_radius : int, optional + Graph radius searched around each sampled path segment for loop + decorations. Defaults to one site or the corridor width, whichever is + larger. + max_loop_clusters_per_segment : int, optional + Maximum connected loop decorations retained near each sampled segment. + max_corridor_edges : int, optional + Hard corridor-geometry limit. Exceeding it raises + :class:`OpenLoopEnumerationLimitError`. + path_edge_weights : mapping, optional + Positive edge costs used to rank shortest paths. Unspecified edges + have unit cost. + corridor_max_bond : int, optional + If supplied, use compressed boundary contraction with this maximum + bond dimension for corridor terms. normalized : bool or {"prod", "separate"}, optional Whether to normalize the final explicit configuration sum. optimize : str or path optimizer, optional @@ -3675,6 +6028,11 @@ def partial_trace_open_loop_series_expand( "estimate" ) + diagnostic_support = ( + None + if diagnostic is None + else diagnostic.supports.get(tuple(where)) + ) return _partial_trace_open_loop_series( bp, where, @@ -3686,6 +6044,23 @@ def partial_trace_open_loop_series_expand( contract_opts=contract_opts, cache=cache or OpenLoopSeriesCache(), info=info, + edge_cutoff=edge_cutoff, + cluster_size=cluster_size, + max_terms=max_terms, + max_enumeration_time=max_enumeration_time, + max_enumeration_memory=max_enumeration_memory, + corridor_width=corridor_width, + max_path_candidates=max_path_candidates, + loop_decoration_size=loop_decoration_size, + corridor_segment_length=corridor_segment_length, + loop_radius=loop_radius, + max_loop_clusters_per_segment=max_loop_clusters_per_segment, + max_corridor_edges=max_corridor_edges, + path_edge_weights=path_edge_weights, + corridor_max_bond=corridor_max_bond, + mode=mode, + auto_corridor_distance=auto_corridor_distance, + diagnostic_support=diagnostic_support, ) @@ -4110,6 +6485,200 @@ def compute_local_expectation_edge_loop_series( return functools.reduce(operator.add, expecs.values()) +def diagnose_open_loop_series( + tn, + terms, + gloops=None, + *, + messages=None, + gauges=None, + run_bp: bool = True, + bp_runner: str = "plain", + relay_opts: dict[str, Any] | None = None, + max_iterations: int = 1000, + tol: float = 5e-6, + tol_abs: float | None = None, + tol_rolling_diff: float | None = 0.0, + diis: bool | dict[str, Any] = False, + damping: float = 0.0, + update: str = "sequential", + require_fixed_point: bool = True, + cache: OpenLoopSeriesCache | None = None, + diagnostic_cache: OpenLoopSeriesDiagnosticCache | None = None, + optimize: Any = "auto-hq", + max_flops_log10: float | None = None, + max_peak_memory_log2: float | None = None, + edge_cutoff=None, + cluster_size=None, + max_terms: int | None = None, + max_enumeration_time: float | None = None, + max_enumeration_memory: int | None = None, + mode: str = "auto", + auto_corridor_distance: int | None = 32, + corridor_width: int | None = None, + max_path_candidates: int = 8, + loop_decoration_size: int = 4, + corridor_segment_length: int = 32, + loop_radius: int | None = None, + max_loop_clusters_per_segment: int = 8, + max_corridor_edges: int | None = 100_000, + path_edge_weights=None, + corridor_max_bond: int | None = None, + contract_opts: dict[str, Any] | None = None, + **bp_opts, +): + """Diagnose open-series geometry and contraction costs without measuring. + + ``terms`` accepts the normal ``{where: operator}`` mapping, an iterable of + ``(where, operator)`` pairs, or :class:`OpenLoopObservableTerm`. Operators + produced by :class:`pepsy.tensors.Fermion` remain native and are never + Jordan--Wigner substituted. The result can be passed as + ``diagnostic=...`` to :func:`compute_local_expectation_open_loop_series`. + + The diagnostic phase performs path/loop discovery and Cotengra tree-cost + estimation only. It does not contract tensor values. For ``mode="auto"`` + native cyclic fermionic supports select the graded cluster-compatible + route first; long separated dense/tree supports select the bounded + corridor route, with ``auto_corridor_distance`` controlling that switch. + """ + records = _normalize_open_observable_terms(terms) + contract_opts = {} if contract_opts is None else dict(contract_opts) + if require_fixed_point and run_bp and ( + not isinstance(max_iterations, (int, np.integer)) or max_iterations < 1 + ): + raise ValueError("max_iterations must be a positive integer when run_bp=True") + + bp, bp_info = _build_bp( + tn, + norm="2norm", + messages=messages, + gauges=gauges, + run_bp=run_bp, + bp_runner=bp_runner, + relay_opts=relay_opts, + max_iterations=max_iterations, + tol=tol, + tol_abs=tol_abs, + tol_rolling_diff=tol_rolling_diff, + diis=diis, + damping=damping, + update=update, + optimize=optimize, + bp_opts=bp_opts, + progbar=False, + ) + if require_fixed_point and run_bp and not bp_info.get("converged", False): + raise RuntimeError( + "diagnose_open_loop_series requires converged BP messages; pass " + "require_fixed_point=False for an exploratory estimate" + ) + + cache = cache or OpenLoopSeriesCache() + diagnostic_cache = diagnostic_cache or OpenLoopSeriesDiagnosticCache() + max_flops_log10, max_peak_memory_log2 = _validate_contraction_cost_limits( + max_flops_log10, + max_peak_memory_log2, + ) + support_reports = {} + cache_hits = 0 + for where_key, gate in records: + sites = _term_sites(bp.tn, where_key) + selection = _resolve_open_route( + bp, + sites, + mode=mode, + corridor_width=corridor_width, + auto_corridor_distance=auto_corridor_distance, + ) + edge_value, cluster_value = _resolve_open_cutoffs( + gloops, + edge_cutoff=edge_cutoff, + cluster_size=cluster_size, + native_cluster_route=selection["native_cluster_route"], + ) + corridor_options = _validate_corridor_options( + corridor_width=selection["corridor_width"], + max_path_candidates=max_path_candidates, + loop_decoration_size=loop_decoration_size, + corridor_segment_length=corridor_segment_length, + loop_radius=loop_radius, + max_loop_clusters_per_segment=max_loop_clusters_per_segment, + max_corridor_edges=max_corridor_edges, + corridor_max_bond=corridor_max_bond, + ) + key = _open_diagnostic_key( + sites, + gate, + route=selection["route"], + edge_cutoff=edge_value, + cluster_size=cluster_value, + corridor_options=corridor_options, + max_terms=max_terms, + max_enumeration_time=max_enumeration_time, + max_enumeration_memory=max_enumeration_memory, + max_flops_log10=max_flops_log10, + max_peak_memory_log2=max_peak_memory_log2, + path_edge_weights=path_edge_weights, + ) + report = diagnostic_cache.get(bp.tn, key) + if report is not None: + cache_hits += 1 + support_reports[tuple(sites)] = report.supports[tuple(sites)] + continue + report_data = _diagnose_open_scalar_support( + bp, + sites, + gate, + gloops, + edge_cutoff=edge_value, + cluster_size=cluster_value, + route_selection=selection, + normalized=True, + optimize=optimize, + contract_opts=contract_opts, + cache=cache, + max_flops_log10=max_flops_log10, + max_peak_memory_log2=max_peak_memory_log2, + max_terms=max_terms, + max_enumeration_time=max_enumeration_time, + max_enumeration_memory=max_enumeration_memory, + corridor_options=corridor_options, + path_edge_weights=path_edge_weights, + ) + report_data.update( + { + "mode": mode, + "edge_cutoff": edge_value, + "cluster_size": cluster_value, + "corridor_options": corridor_options, + "cache_key": key, + } + ) + support_report = OpenLoopSeriesDiagnostic({tuple(sites): report_data}) + diagnostic_cache.put(bp.tn, key, support_report) + support_reports[tuple(sites)] = report_data + + costs = [ + report.get("total_flops_log10") + for report in support_reports.values() + if report.get("total_flops_log10") is not None + ] + peaks = [ + report.get("peak_memory_log2") + for report in support_reports.values() + if report.get("peak_memory_log2") is not None + ] + total_flops = None if not costs else float( + np.log10(sum(10.0 ** value for value in costs)) + ) + return OpenLoopSeriesDiagnostic( + supports=support_reports, + total_flops_log10=total_flops, + peak_memory_log2=max(peaks, default=None), + cache_hits=cache_hits, + ) + + def compute_local_expectation_open_loop_series( tn, terms, @@ -4133,6 +6702,24 @@ def compute_local_expectation_open_loop_series( optimize: Any = "auto-hq", max_flops_log10: float | None = None, max_peak_memory_log2: float | None = None, + edge_cutoff=None, + cluster_size=None, + max_terms: int | None = None, + max_enumeration_time: float | None = None, + max_enumeration_memory: int | None = None, + mode: str = "exact", + auto_corridor_distance: int | None = 32, + corridor_width: int | None = None, + max_path_candidates: int = 8, + loop_decoration_size: int = 4, + corridor_segment_length: int = 32, + loop_radius: int | None = None, + max_loop_clusters_per_segment: int = 8, + max_corridor_edges: int | None = 100_000, + path_edge_weights=None, + corridor_max_bond: int | None = None, + diagnostic_cache: OpenLoopSeriesDiagnosticCache | None = None, + diagnostic: OpenLoopSeriesDiagnostic | None = None, info: dict[str, Any] | None = None, return_all: bool = False, contract_opts: dict[str, Any] | None = None, @@ -4140,7 +6727,10 @@ def compute_local_expectation_open_loop_series( ): """Compute fermion-safe expectations from open-edge loop terms. - The terms mapping has the usual site-or-sites to gate form. Unlike the + The terms mapping has the usual site-or-sites to gate form. It can also be + an iterable of ``(where, operator)`` pairs or + :class:`OpenLoopObservableTerm`; native operators produced by + :class:`pepsy.tensors.Fermion` are accepted directly. Unlike the open rho API, this function evaluates the observable as a scalar numerator and identity denominator. Dense networks use direct gate insertion over the explicit open-path, closed-loop, and path-plus-loop configurations; @@ -4148,9 +6738,19 @@ def compute_local_expectation_open_loop_series( the gate in the ket/bra contraction, so the graded gate ordering is preserved without materializing a diagnostic rho. - Integer gloops counts excited virtual edges. For native fermionic PEPS, - this is the preferred route for even two-site operators such as hopping, - pairing, and density terms. + ``edge_cutoff`` counts excited virtual edges on dense networks and + fermionic trees. For cyclic native fermionic PEPS, pass ``cluster_size``; + the safe graded cluster route is selected before edge terms are generated. + ``gloops`` remains a legacy route-dependent alias. + + ``mode="auto"`` selects the graded cluster-compatible route for native + cyclic fermionic supports and the bounded corridor route for supports + farther apart than ``auto_corridor_distance``. ``mode="exact"`` is the + compatibility default; an explicitly supplied ``corridor_width`` still + activates corridor mode. For a strict two-phase workflow, call + :func:`diagnose_open_loop_series` first and pass its result as + ``diagnostic=...``. That reuses the discovered terms and records the + pre-contraction FLOP and peak-memory estimates in ``info``. ``optimize`` is forwarded unchanged to Quimb's ``TensorNetwork.contract`` calls. Callers can pass the reusable Cotengra @@ -4165,11 +6765,21 @@ def compute_local_expectation_open_loop_series( ``info["open_scalar_cluster_region_skipped_terms"]`` depending on the native route. The older ``open_scalar_skipped_terms`` field is a route-specific compatibility alias. + + ``max_terms``, ``max_enumeration_time``, and + ``max_enumeration_memory`` bound explicit edge-geometry discovery. A + bound violation raises :class:`OpenLoopEnumerationLimitError` rather than + returning a silently incomplete observable. + + ``corridor_width`` activates the bounded long-separation route: weighted + shortest paths are retained in a small beam, the paths are inflated into + a corridor, and connected loop decorations are sampled near its segments. + Set ``corridor_max_bond`` for compressed boundary contraction. This is an + approximation and intentionally omits disconnected products of distant + loop clusters; widen the corridor or increase the decoration controls for + convergence diagnostics. """ - if not hasattr(terms, "items"): - raise TypeError("terms must be a mapping from sites to operators") - if not terms: - raise ValueError("terms must contain at least one operator") + records = _normalize_open_observable_terms(terms) if normalized == "prod": normalized = True if normalized not in (True, False, "separate"): @@ -4204,6 +6814,97 @@ def compute_local_expectation_open_loop_series( ) cache = cache or OpenLoopSeriesCache() + cached_diagnostic_supports = {} + if diagnostic is None and (diagnostic_cache is not None or mode == "auto"): + diagnostic_flops, diagnostic_peak = _validate_contraction_cost_limits( + max_flops_log10, + max_peak_memory_log2, + ) + for where_key, gate in records: + sites = _term_sites(bp.tn, where_key) + selection = _resolve_open_route( + bp, + sites, + mode=mode, + corridor_width=corridor_width, + auto_corridor_distance=auto_corridor_distance, + ) + edge_value, cluster_value = _resolve_open_cutoffs( + gloops, + edge_cutoff=edge_cutoff, + cluster_size=cluster_size, + native_cluster_route=selection["native_cluster_route"], + ) + corridor_options = _validate_corridor_options( + corridor_width=selection["corridor_width"], + max_path_candidates=max_path_candidates, + loop_decoration_size=loop_decoration_size, + corridor_segment_length=corridor_segment_length, + loop_radius=loop_radius, + max_loop_clusters_per_segment=max_loop_clusters_per_segment, + max_corridor_edges=max_corridor_edges, + corridor_max_bond=corridor_max_bond, + ) + key = _open_diagnostic_key( + sites, + gate, + route=selection["route"], + edge_cutoff=edge_value, + cluster_size=cluster_value, + corridor_options=corridor_options, + max_terms=max_terms, + max_enumeration_time=max_enumeration_time, + max_enumeration_memory=max_enumeration_memory, + max_flops_log10=max_flops_log10, + max_peak_memory_log2=max_peak_memory_log2, + path_edge_weights=path_edge_weights, + ) + cached = ( + None + if diagnostic_cache is None + else diagnostic_cache.get(bp.tn, key) + ) + if cached is not None: + cached_diagnostic_supports[tuple(sites)] = cached.supports[ + tuple(sites) + ] + elif mode == "auto": + report_data = _diagnose_open_scalar_support( + bp, + sites, + gate, + gloops, + edge_cutoff=edge_value, + cluster_size=cluster_value, + route_selection=selection, + normalized=True, + optimize=optimize, + contract_opts=contract_opts, + cache=cache, + max_flops_log10=diagnostic_flops, + max_peak_memory_log2=diagnostic_peak, + max_terms=max_terms, + max_enumeration_time=max_enumeration_time, + max_enumeration_memory=max_enumeration_memory, + corridor_options=corridor_options, + path_edge_weights=path_edge_weights, + ) + report_data.update( + { + "mode": mode, + "edge_cutoff": edge_value, + "cluster_size": cluster_value, + "corridor_options": corridor_options, + "cache_key": key, + } + ) + cached_diagnostic_supports[tuple(sites)] = report_data + if diagnostic_cache is not None: + diagnostic_cache.put( + bp.tn, + key, + OpenLoopSeriesDiagnostic({tuple(sites): report_data}), + ) term_info = ( {} if info is None @@ -4215,8 +6916,13 @@ def compute_local_expectation_open_loop_series( else info.setdefault("open_scalar_supports", {}) ) expecs = {} - for where, gate in terms.items(): + for where, gate in records: sites = _term_sites(bp.tn, where) + diagnostic_support = ( + cached_diagnostic_supports.get(tuple(sites)) + if diagnostic is None + else diagnostic.supports.get(tuple(sites)) + ) value, normalization = _local_expectation_open_loop_series( bp, sites, @@ -4226,12 +6932,34 @@ def compute_local_expectation_open_loop_series( optimize=optimize, max_flops_log10=max_flops_log10, max_peak_memory_log2=max_peak_memory_log2, + edge_cutoff=edge_cutoff, + cluster_size=cluster_size, + max_terms=max_terms, + max_enumeration_time=max_enumeration_time, + max_enumeration_memory=max_enumeration_memory, + corridor_width=corridor_width, + max_path_candidates=max_path_candidates, + loop_decoration_size=loop_decoration_size, + corridor_segment_length=corridor_segment_length, + loop_radius=loop_radius, + max_loop_clusters_per_segment=max_loop_clusters_per_segment, + max_corridor_edges=max_corridor_edges, + path_edge_weights=path_edge_weights, + corridor_max_bond=corridor_max_bond, + mode=mode, + auto_corridor_distance=auto_corridor_distance, + diagnostic_support=diagnostic_support, contract_opts=contract_opts, cache=cache, info=info, ) - term_info[where] = normalization - expecs[where] = value + result_key = where + try: + hash(result_key) + except TypeError: + result_key = tuple(sites) + term_info[result_key] = normalization + expecs[result_key] = value if info is not None: support_key = tuple(sites) support_edge_costs = { @@ -4275,6 +7003,14 @@ def compute_local_expectation_open_loop_series( "cluster_region_costs": support_cluster_costs, "family_counts": dict(info["open_scalar_family_counts"]), "family_weights": dict(info["open_scalar_family_weights"]), + "corridor": dict( + info.get("open_scalar_corridor", {}) + ), + "diagnostic": ( + None + if diagnostic_support is None + else dict(diagnostic_support) + ), } if return_all: return expecs diff --git a/src/pepsy/operators/gates.py b/src/pepsy/operators/gates.py index 3d50fcd..5697316 100644 --- a/src/pepsy/operators/gates.py +++ b/src/pepsy/operators/gates.py @@ -18,7 +18,7 @@ infer_backend_converter_from_sample, resolve_backend_sample_data_from_tn, ) -from ..tensors.core import add_cycle, id_to_mpo, id_to_pepo +from ..tensors.core import OneDMap, add_cycle, id_to_mpo, id_to_pepo __all__ = [ "gate", @@ -2798,6 +2798,119 @@ def _apply_gate_3d( return tn +def _native_gate_stream_info(gate_list, *, allow_charged=False): + """Validate a native Symmray fermionic gate stream and describe its legs.""" + native_flags = [ + _is_block_sparse_array(gate_op) + and "FermionicArray" in type(gate_op).__name__ + for gate_op in gate_list + ] + if not any(native_flags): + return None + if not all(native_flags): + raise TypeError( + "Native FermionicArray gate builders cannot mix native and dense gates." + ) + + from ..tensors.symmetric import ( # pylint: disable=import-outside-toplevel + _expanded_index_charges, + _normalize_group_charge, + _zero_like_charge, + ) + + first = gate_list[0] + rank = len(first.indices) + if rank not in {2, 4}: + raise ValueError( + "Native gate builders support one- and two-site FermionicArray gates." + ) + n_sites = rank // 2 + symmetry = str(getattr(first, "symmetry", "")) + first_charge = _normalize_group_charge( + getattr(first, "charge", 0), symmetry + ) + zero = _zero_like_charge(first_charge) + gate_charges = [ + _normalize_group_charge(getattr(gate_op, "charge", zero), symmetry) + for gate_op in gate_list + ] + if not allow_charged and any(charge != zero for charge in gate_charges): + raise ValueError( + "Native gate builders require charge-neutral gates by default; " + "pass allow_charged=True to accumulate a charged operator stream." + ) + + output_maps = tuple( + tuple(_expanded_index_charges(index)) + for index in first.indices[:n_sites] + ) + input_maps = tuple( + tuple(_expanded_index_charges(index)) + for index in first.indices[n_sites:] + ) + if output_maps != input_maps: + raise ValueError( + "Native gate builders require matching upper/lower physical charge maps." + ) + if any(site_map != output_maps[0] for site_map in output_maps[1:]): + raise ValueError( + "Native gate builders require one physical charge map per site." + ) + + for gate_op, gate_charge in zip(gate_list[1:], gate_charges[1:]): + if len(gate_op.indices) != rank: + raise ValueError("All native gates must act on the same number of sites.") + if str(getattr(gate_op, "symmetry", "")) != symmetry: + raise ValueError("All native gates must use the same Abelian symmetry.") + if not allow_charged and gate_charge != zero: + raise ValueError("All native gates must be charge-neutral.") + gate_maps = tuple( + tuple(_expanded_index_charges(index)) + for index in gate_op.indices + ) + if gate_maps[:n_sites] != output_maps or gate_maps[n_sites:] != input_maps: + raise ValueError( + "All native gates must use the same physical charge maps." + ) + + return { + "fermionic": True, + "n_sites": n_sites, + "phys_map": list(output_maps[0]), + "symmetry": symmetry, + "zero": zero, + "dtype": getattr(first, "dtype", np.dtype("complex128")), + "gate_charges": tuple(gate_charges), + } + + +def _native_identity_mpo(length, info, *, max_bond, cutoff): + """Build an identity MPO whose tensors use native fermionic sectors.""" + from ..tensors.symmetric import ( # pylint: disable=import-outside-toplevel + _assemble_symmray_mpo, + ) + + start = ("start",) + done = ("done",) + channels = [ + [(start, info["zero"]), (done, info["zero"])] + for _ in range(max(length - 1, 0)) + ] + return _assemble_symmray_mpo( + L=length, + channels=channels, + transitions=[[] for _ in range(length)], + phys_map=info["phys_map"], + symmetry=info["symmetry"], + zero=info["zero"], + dtype=info["dtype"], + max_bond=max_bond, + cutoff=cutoff, + compress=False, + fermionic=True, + ) + + def build_pepo_from_gates( gates, wheres=None, @@ -2810,6 +2923,8 @@ def build_pepo_from_gates( sequence="auto", contract="reduce-split", ind_id="k{},{}", + mapper=None, + allow_charged=False, ): """Build a PEPO from gate-style input on top of a PEPO identity. @@ -2846,6 +2961,13 @@ def build_pepo_from_gates( path, which is usually cheaper than ``"split"`` for PEPO/PEPS tensors. ind_id : str, default="k{},{}" Physical index format used for PEPO ket-family indices. + mapper : OneDMap | None, optional + Optional lattice-to-chain mapping used when native Symmray gates are + supplied. PEPO conversion supports ``snake`` and + ``snake-row-major`` mappings. + allow_charged : bool, default=False + Allow native gates with nonzero operator charge. The returned PEPO + then carries the accumulated charge of the sequential gate product. Returns ------- @@ -2856,19 +2978,70 @@ def build_pepo_from_gates( gate_list = [g for g, _ in entries] where_list = [w for _, w in entries] + native_info = _native_gate_stream_info( + gate_list, + allow_charged=allow_charged, + ) + coords = [c for w in where_list for c in w] - Lx = max(i for i, _ in coords) + 1 - Ly = max(j for _, j in coords) + 1 + if mapper is None: + Lx = max(i for i, _ in coords) + 1 + Ly = max(j for _, j in coords) + 1 + else: + if not isinstance(mapper, OneDMap): + raise TypeError("mapper must be a 2D OneDMap instance or None.") + try: + mapper_shape = tuple(mapper.shape) + except (AttributeError, TypeError, ValueError) as exc: + raise TypeError("mapper must be a 2D OneDMap instance or None.") from exc + if len(mapper_shape) != 2: + raise ValueError("build_pepo_from_gates requires a 2D OneDMap.") + Lx, Ly = mapper_shape - pepo = pepo_.copy() if pepo_ is not None else id_to_pepo(Lx, Ly, dtype=dtype) - if pepo_ is None and cyclic: - pepo = add_cycle(pepo, 1) + if native_info is not None and pepo_ is None: + from .hamiltonians import ham_tn # pylint: disable=import-outside-toplevel - for tensor in pepo: - tensor.modify(data=ar.do("array", tensor.data, like=gate_list[0])) + builder = ham_tn( + Lx=Lx, + Ly=Ly, + mapper=mapper, + max_bond=256 if max_bond is None else max_bond, + cutoff=cutoff, + data_type=native_info["dtype"], + ) + mpo = _native_identity_mpo( + builder.L, + native_info, + max_bond=max_bond, + cutoff=cutoff, + ) + pepo = builder.mpo_to_pepo( + mpo, + cycle_peps=cyclic, + cycle_bond_dim=1, + inplace=True, + ) + elif native_info is not None: + pepo = pepo_.copy() + if any(not _is_block_sparse_array(tensor.data) for tensor in pepo): + raise TypeError( + "Native FermionicArray gates require a native Symmray PEPO." + ) + else: + pepo = pepo_.copy() if pepo_ is not None else id_to_pepo(Lx, Ly, dtype=dtype) + if pepo_ is None and cyclic: + pepo = add_cycle(pepo, 1) + + if native_info is None: + for tensor in pepo: + tensor.modify(data=ar.do("array", tensor.data, like=gate_list[0])) for gate_op, where_norm in zip(gate_list, where_list): - gate_use = _to_ket_gate_layout(gate_op, len(where_norm)) + gate_use = ( + gate_op + if native_info is not None + else _to_ket_gate_layout(gate_op, len(where_norm)) + ) gate( pepo, gate_use, where_norm, @@ -2925,6 +3098,7 @@ def build_mpo_from_gates( max_bond=16, contract="reduce-split", ind_id="k{}", + allow_charged=False, ): """Build an MPO from gate-style input on top of an MPO identity. @@ -2957,6 +3131,9 @@ def build_mpo_from_gates( path, which is usually cheaper than ``"split"`` for MPO tensors. ind_id : str, default="k{}" Physical index format used for MPO ket-family indices. + allow_charged : bool, default=False + Allow native gates with nonzero operator charge. The returned MPO + then carries the accumulated charge of the sequential gate product. Returns ------- @@ -2967,18 +3144,42 @@ def build_mpo_from_gates( gate_list = [g for g, _ in entries] where_list = [w for _, w in entries] + native_info = _native_gate_stream_info( + gate_list, + allow_charged=allow_charged, + ) + coords = [int(i) for w in where_list for i in w] L = max(coords) + 1 - mpo = mpo_.copy() if mpo_ is not None else id_to_mpo( - L, phys_dim=2, dtype=dtype, cyclic=cyclic - ) + if native_info is not None and mpo_ is None: + mpo = _native_identity_mpo( + L, + native_info, + max_bond=max_bond, + cutoff=cutoff, + ) + else: + mpo = mpo_.copy() if mpo_ is not None else id_to_mpo( + L, phys_dim=2, dtype=dtype, cyclic=cyclic + ) - for tensor in mpo: - tensor.modify(data=ar.do("array", tensor.data, like=gate_list[0])) + if native_info is not None and any( + not _is_block_sparse_array(tensor.data) for tensor in mpo + ): + raise TypeError( + "Native FermionicArray gates require a native Symmray MPO." + ) + if native_info is None: + for tensor in mpo: + tensor.modify(data=ar.do("array", tensor.data, like=gate_list[0])) for gate_op, where_norm in zip(gate_list, where_list): - gate_use = _to_ket_gate_layout(gate_op, len(where_norm)) + gate_use = ( + gate_op + if native_info is not None + else _to_ket_gate_layout(gate_op, len(where_norm)) + ) gate( mpo, diff --git a/src/pepsy/operators/hamiltonians.py b/src/pepsy/operators/hamiltonians.py index 7aed622..7262dea 100644 --- a/src/pepsy/operators/hamiltonians.py +++ b/src/pepsy/operators/hamiltonians.py @@ -3,6 +3,7 @@ from __future__ import annotations import warnings +from collections.abc import Mapping from numbers import Integral @@ -476,6 +477,7 @@ def build_mpo( fermion=None, edges=None, fermionic=None, + charge_sectors=False, **model_params, ): """Build MPO from user interactions. @@ -515,6 +517,10 @@ def build_mpo( Native graded encoding flag for the fermion-model form. ``None`` and ``False`` select the Jordan-Wigner-compatible MPO builder; ``True`` selects ``Fermion.to_mpo(...)``. + charge_sectors : bool, default=False + When native construction is enabled, return one MPO per operator + charge as ``{charge: mpo}`` instead of requiring one homogeneous + charge for the whole collection. **model_params Explicit fermion couplings such as ``t``, ``U``/``V``, and ``mu``. @@ -553,6 +559,8 @@ def build_mpo( cutoff_use = self.cutoff if cutoff is None else float(cutoff) mapper_use = self.mapper if mapper is None else mapper fermionic_use = False if fermionic is None else bool(fermionic) + if charge_sectors and not fermionic_use: + raise ValueError("charge_sectors=True requires fermionic=True.") mpo_builder = fermion.to_mpo if fermionic_use else fermion.build_mpo return mpo_builder( ints, @@ -563,10 +571,11 @@ def build_mpo( compress=bool(compress_each), dtype=dtype, fermionic=fermionic_use, + charge_sectors=charge_sectors, **model_params, ) - if edges is not None or model_params or fermionic is not None: + if edges is not None or model_params or fermionic is not None or charge_sectors: raise TypeError( "edges, fermion model parameters, and fermionic encoding are " "only valid with fermion=... ." @@ -709,7 +718,7 @@ def mpo_to_pepo( def build_pepo( self, - ints, + ints=None, *, phys_dim=2, max_bond=None, @@ -718,8 +727,21 @@ def build_pepo( compress_each=True, cycle_peps=False, cycle_bond_dim=1, + mapper=None, + fermion=None, + edges=None, + fermionic=None, + charge_sectors=False, + **model_params, ): - """Build PEPO directly from interaction terms.""" + """Build a PEPO from interactions or a native fermion model. + + The ``fermion=...``/``edges=...`` form mirrors :meth:`build_mpo` and + forwards ``mapper=OneDMap(...)`` and ``fermionic=True`` to the native + fermion MPO builder before converting the result to a PEPO. + With ``charge_sectors=True``, return ``{charge: pepo}`` for a mixed + native operator. + """ self._require_2d("build_pepo") mpo = self.build_mpo( ints, @@ -728,7 +750,23 @@ def build_pepo( cutoff=cutoff, data_type=data_type, compress_each=compress_each, + mapper=mapper, + fermion=fermion, + edges=edges, + fermionic=fermionic, + charge_sectors=charge_sectors, + **model_params, ) + if isinstance(mpo, Mapping): + return { + charge: self.mpo_to_pepo( + sector_mpo, + cycle_peps=cycle_peps, + cycle_bond_dim=cycle_bond_dim, + inplace=False, + ) + for charge, sector_mpo in mpo.items() + } return self.mpo_to_pepo( mpo, cycle_peps=cycle_peps, diff --git a/src/pepsy/tensors/symmetric.py b/src/pepsy/tensors/symmetric.py index fbcf74a..bb16450 100644 --- a/src/pepsy/tensors/symmetric.py +++ b/src/pepsy/tensors/symmetric.py @@ -5474,7 +5474,13 @@ def _assemble_symmray_mpo( site_tag_id="I{}", to_backend=None, fermionic=False, + operator_charge=None, ): + operator_charge = ( + zero + if operator_charge is None + else _normalize_group_charge(operator_charge, symmetry) + ) channel_pos = [ {channel_id: pos for pos, (channel_id, _) in enumerate(cut_channels)} for cut_channels in channels @@ -5540,7 +5546,10 @@ def _assemble_symmray_mpo( index_maps=index_maps, duals=duals, fermionic=bool(fermionic), - charge=zero, + charge=operator_charge if site == L - 1 else zero, + label=(site if fermionic and site == L - 1 and + _charged_op_needs_fermion_string(operator_charge) + else None), ) ) @@ -5756,8 +5765,9 @@ def _add_native_term_to_mpo( symmetry, dtype, zero, + operator_charge=None, ): - """Add one neutral native fermion term as graded MPO transitions. + """Add one homogeneous native fermion term as graded MPO transitions. The term is split by operator Schmidt decompositions over the ordered support sites. This preserves Symmray's fermionic bond phases while @@ -5800,10 +5810,15 @@ def _add_native_term_to_mpo( getattr(term, "charge", zero), symmetry, ) - if term_charge != zero: + expected_charge = ( + zero + if operator_charge is None + else _normalize_group_charge(operator_charge, symmetry) + ) + if term_charge != expected_charge: raise ValueError( - "MPO Hamiltonian terms must be neutral under the selected " - f"symmetry; term {term_pos} has charge {term_charge!r}." + "Native MPO terms must share one homogeneous operator charge " + f"{expected_charge!r}; term {term_pos} has charge {term_charge!r}." ) if n_sites == 1: @@ -5833,8 +5848,13 @@ def _add_native_term_to_mpo( else: left_group = (0, 1) right_group = tuple(range(2, ndim)) + # Absorb the singular values into the left factor for a charged + # operator. This leaves the total operator charge on the final site + # factor, where the open MPO boundary can carry it, while all + # preceding tensors remain neutral and can propagate identity paths. + absorb = "left" if expected_charge != zero else "right" left, _, right = current.fuse(left_group, right_group).svd( - absorb="right" + absorb=absorb ) if pos == 0: factors.append(left.unfuse(0).transpose((2, 0, 1))) @@ -5855,6 +5875,11 @@ def _add_native_term_to_mpo( bond_map = _expanded_index_charges(factor.indices[bond_axis]) channel_ids = [] for bond_pos, bond_charge in enumerate(bond_map): + if expected_charge != zero: + # With absorb="left", the factor bond is dual on the side + # that becomes the MPO's outgoing bond. Reverse its charge + # when installing the common MPO bond orientation. + bond_charge = _charge_neg(bond_charge, symmetry) channel_id = ("native", term_pos, interval, bond_pos) channel_ids.append(channel_id) for cut in range(sites[interval], sites[interval + 1]): @@ -5953,15 +5978,33 @@ def _generic_symhamiltonian_to_mpo( hamiltonian.symmetry, ) zero = _zero_like_charge(first_charge) + operator_charge = first_charge if fermionic else zero + if not fermionic and first_charge != zero: + raise ValueError( + "Charged native operator terms require fermionic=True; the " + "Jordan-Wigner compatibility MPO is neutral-only." + ) start = ("start",) done = ("done",) - channels = [[(start, zero), (done, zero)] for _ in range(max(L - 1, 0))] + channels = [ + [(start, zero), (done, _charge_neg(operator_charge, hamiltonian.symmetry))] + for _ in range(max(L - 1, 0)) + ] transitions = [[] for _ in range(L)] phys_map = None for term_pos, (raw_where, where) in enumerate(zip(raw_wheres, mapped_wheres)): term = hamiltonian.terms[raw_where] term_is_fermionic = _is_fermionic_symmray_array(term) + term_charge = _normalize_group_charge( + getattr(term, "charge", zero), + hamiltonian.symmetry, + ) + if not fermionic and term_charge != zero: + raise ValueError( + "Charged native operator terms require fermionic=True; the " + "Jordan-Wigner compatibility MPO is neutral-only." + ) if fermionic and not term_is_fermionic: raise TypeError( "Native fermionic MPO construction requires every Hamiltonian " @@ -5978,6 +6021,7 @@ def _generic_symhamiltonian_to_mpo( symmetry=hamiltonian.symmetry, dtype=dtype, zero=zero, + operator_charge=operator_charge, ) if phys_map is None: phys_map = term_phys @@ -6082,9 +6126,22 @@ def _generic_symhamiltonian_to_mpo( site_tag_id=site_tag_id, to_backend=to_backend, fermionic=fermionic, + operator_charge=operator_charge, ) +def _group_symhamiltonian_terms_by_charge(hamiltonian): + """Group native Hamiltonian terms into homogeneous charge sectors.""" + sectors = {} + for where, term in hamiltonian.terms.items(): + charge = _normalize_group_charge( + getattr(term, "charge", 0), + hamiltonian.symmetry, + ) + sectors.setdefault(charge, {})[where] = term + return sectors + + @dataclass(frozen=True) class SymHamiltonian: """Container for Symmray local Hamiltonian terms.""" @@ -6178,6 +6235,7 @@ def to_mpo( to_backend=None, dtype=None, fermionic=False, + charge_sectors=False, ): """Build a symmetry-preserving MPS-chain MPO for this Hamiltonian. @@ -6186,8 +6244,46 @@ def to_mpo( Fermionic compatibility paths include parity strings along non-adjacent mapped hopping channels. With ``fermionic=True``, native Symmray ``FermionicArray`` tensors are built directly from arbitrary - neutral one- or multi-site terms. + homogeneous-charge one- or multi-site terms. The open MPO boundary + carries a nonzero operator charge when required. + + With ``charge_sectors=True``, return a mapping from each operator + charge to its own homogeneous native MPO. This is the explicit way to + represent a mixed-charge operator such as ``I + c^\u2020`` without + converting it to a dense or non-symmetric tensor. """ + if charge_sectors: + if not fermionic: + raise ValueError("charge_sectors=True requires fermionic=True.") + sectors = _group_symhamiltonian_terms_by_charge(self) + if not sectors: + raise ValueError( + "At least one Hamiltonian term is required to build an MPO." + ) + return { + charge: type(self).from_terms( + self.model, + self.symmetry, + terms, + parameters=self.parameters, + ).to_mpo( + L=L, + mapper=mapper, + idx2coo=idx2coo, + coo2idx=coo2idx, + max_bond=max_bond, + cutoff=cutoff, + compress=compress, + upper_ind_id=upper_ind_id, + lower_ind_id=lower_ind_id, + site_tag_id=site_tag_id, + to_backend=to_backend, + dtype=dtype, + fermionic=True, + charge_sectors=False, + ) + for charge, terms in sectors.items() + } if self.explicit_terms or fermionic: return _generic_symhamiltonian_to_mpo( self, @@ -6359,6 +6455,77 @@ def to_mpo( to_backend=to_backend, ) + def to_pepo( + self, + Lx=None, + Ly=None, + *, + mapper=None, + max_bond=None, + cutoff=1e-12, + compress=True, + cyclic=False, + cycle_bond_dim=1, + dtype=None, + fermionic=True, + to_backend=None, + charge_sectors=False, + ): + """Build a 2D PEPO from this Hamiltonian's native local terms. + + The Hamiltonian is first assembled as an MPO using ``mapper`` and is + then embedded with the same snake-style ordering into a PEPO. Native + ``fermionic=True`` construction preserves homogeneous neutral or + nonzero operator charge and Symmray grading metadata. Set + ``fermionic=False`` to request the compatibility MPO path. + With ``charge_sectors=True``, return ``{charge: PEPO}`` for mixed + charge collections. + """ + if Lx is None or Ly is None: + raise TypeError("to_pepo requires both Lx and Ly.") + + from ..operators.hamiltonians import ham_tn + + builder = ham_tn( + Lx=Lx, + Ly=Ly, + mapper=mapper, + max_bond=256 if max_bond is None else max_bond, + cutoff=cutoff, + data_type=( + _dtype_from_hamiltonian_terms(self.terms) + if dtype is None + else dtype + ), + ) + mpo = self.to_mpo( + L=builder.L, + mapper=builder.mapper, + max_bond=max_bond, + cutoff=cutoff, + compress=compress, + dtype=dtype, + fermionic=fermionic, + charge_sectors=charge_sectors, + to_backend=to_backend, + ) + if charge_sectors: + return { + charge: builder.mpo_to_pepo( + sector_mpo, + cycle_peps=cyclic, + cycle_bond_dim=cycle_bond_dim, + inplace=True, + ) + for charge, sector_mpo in mpo.items() + } + return builder.mpo_to_pepo( + mpo, + cycle_peps=cyclic, + cycle_bond_dim=cycle_bond_dim, + inplace=True, + ) + def jw_trotter_gates( self, dt, @@ -9759,6 +9926,7 @@ def build_mpo( site_tag_id="I{}", dtype=None, fermionic=False, + charge_sectors=False, to_backend=None, **params, ): @@ -9794,7 +9962,53 @@ def build_mpo( site_tag_id=site_tag_id, dtype=dtype, fermionic=fermionic, + charge_sectors=charge_sectors, + to_backend=to_backend, + ) + + def build_pepo( + self, + terms_or_edges=None, + *, + hamiltonian=None, + Lx=None, + Ly=None, + mapper=None, + max_bond=None, + cutoff=1e-12, + compress=True, + cyclic=False, + cycle_bond_dim=1, + dtype=None, + fermionic=True, + charge_sectors=False, + to_backend=None, + **params, + ): + """Build a 2D PEPO from this fermion model. + + This is the model-facing shorthand for :meth:`to_pepo`. Native + graded construction is selected by default; pass ``fermionic=False`` + for the compatibility Jordan--Wigner MPO before PEPO embedding. + Coordinate-keyed explicit terms can be supplied with a + ``mapper=OneDMap(...)``. + """ + return self.to_pepo( + terms_or_edges, + hamiltonian=hamiltonian, + Lx=Lx, + Ly=Ly, + mapper=mapper, + max_bond=max_bond, + cutoff=cutoff, + compress=compress, + cyclic=cyclic, + cycle_bond_dim=cycle_bond_dim, + dtype=dtype, + fermionic=fermionic, + charge_sectors=charge_sectors, to_backend=to_backend, + **params, ) def to_mpo( @@ -9814,6 +10028,7 @@ def to_mpo( site_tag_id="I{}", dtype=None, fermionic=True, + charge_sectors=False, to_backend=None, **params, ): @@ -9821,9 +10036,11 @@ def to_mpo( ``terms_or_edges`` may be lattice edges for the built-in model or a mapping such as ``{(0, 2, 4): fermion.operator_term(...)}``. The - latter supports arbitrary neutral term support, including + latter supports arbitrary homogeneous-charge term support, including non-contiguous sites. Pass an existing :class:`SymHamiltonian` with ``hamiltonian=`` when the terms have already been assembled. + Set ``charge_sectors=True`` to return one native MPO per charge for a + mixed-charge term collection. The native path is selected by default and returns MPO tensors backed by Symmray ``FermionicArray`` objects. Set ``fermionic=False`` to @@ -9872,9 +10089,133 @@ def to_mpo( site_tag_id=site_tag_id, dtype=dtype, fermionic=fermionic, + charge_sectors=charge_sectors, to_backend=to_backend, ) + def to_pepo( + self, + terms_or_edges=None, + *, + hamiltonian=None, + Lx=None, + Ly=None, + mapper=None, + max_bond=None, + cutoff=1e-12, + compress=True, + cyclic=False, + cycle_bond_dim=1, + dtype=None, + fermionic=True, + charge_sectors=False, + to_backend=None, + **params, + ): + """Build a native fermionic PEPO on a 2D lattice. + + The operator terms can be keyed by lattice coordinates, for example + ``{((0, 1), (2, 2)): term}``, where ``term`` is a native + :class:`symmray.FermionicArray` returned by :meth:`operator_term`. + Use ``{((0, 1),): term}`` for a one-site coordinate term so it is not + confused with a one-dimensional ``(i, j)`` edge. + The native graded MPO assembler is used internally with the supplied + one-dimensional map, and the result is embedded as a snake-style + PEPO. This keeps the fermionic charge and grading metadata intact, + including homogeneous nonzero operator charge and odd-parity dummy + modes; it does not pass through a dense or Jordan--Wigner + representation when ``fermionic=True``. + + Parameters + ---------- + terms_or_edges : mapping, sequence, or SymHamiltonian + Explicit coordinate-keyed native terms, built-in model edges, or + an already assembled Hamiltonian. + hamiltonian : SymHamiltonian, optional + Existing Hamiltonian. Pass either this or ``terms_or_edges``. + Lx, Ly : int + Dimensions of the 2D PEPO lattice. + mapper : OneDMap, optional + One-dimensional ordering used for the native fermionic channels. + The PEPO embedding currently requires ``snake`` or + ``snake-row-major`` ordering. + max_bond, cutoff, compress + Forwarded to native MPO construction before PEPO embedding. + cyclic : bool, optional + Add dimension-``cycle_bond_dim`` PEPO bonds around both lattice + directions after embedding. + dtype, fermionic, to_backend + Forwarded to :meth:`to_mpo`. Keep ``fermionic=True`` for the + native graded path; ``False`` explicitly selects the compatibility + MPO path. + charge_sectors : bool, optional + Return ``{charge: PEPO}`` for mixed-charge term collections. + + Returns + ------- + qtn.PEPO + A PEPO with coordinate tags ``I{x},{y}``, input indices + ``k{x},{y}``, and output indices ``b{x},{y}``. + + Notes + ----- + Native terms must be homogeneous: all terms in one operator + collection must carry the same Abelian charge, unless + ``charge_sectors=True`` is requested. Neutral and nonzero charges are + both supported with ``fermionic=True``. Odd-parity terms + should be created with an explicit ``label=`` in + :meth:`operator_term` so their dummy-mode phase metadata is retained. + The Jordan--Wigner compatibility path remains neutral-only. + The current implementation uses the MPO ordering as the fermionic + ordering. Thus arbitrary two-site and non-contiguous terms are + supported, but the PEPO's nontrivial operator bonds follow the + selected snake-style chain; the added transverse lattice bonds have + dimension one unless ``cyclic=True``. + """ + if Lx is None or Ly is None: + raise TypeError("to_pepo requires both Lx and Ly.") + + from ..operators.hamiltonians import ham_tn + + builder = ham_tn( + Lx=Lx, + Ly=Ly, + mapper=mapper, + max_bond=256 if max_bond is None else max_bond, + cutoff=cutoff, + data_type=self.dtype if dtype is None else dtype, + ) + mpo = self.to_mpo( + terms_or_edges, + hamiltonian=hamiltonian, + L=builder.L, + mapper=builder.mapper, + max_bond=max_bond, + cutoff=cutoff, + compress=compress, + dtype=dtype, + fermionic=fermionic, + charge_sectors=charge_sectors, + to_backend=to_backend, + **params, + ) + if charge_sectors: + return { + charge: builder.mpo_to_pepo( + sector_mpo, + cycle_peps=cyclic, + cycle_bond_dim=cycle_bond_dim, + inplace=True, + ) + for charge, sector_mpo in mpo.items() + } + return builder.mpo_to_pepo( + mpo, + cycle_peps=cyclic, + cycle_bond_dim=cycle_bond_dim, + inplace=True, + ) + def local_terms(self, edges, *, layout="site", **params): """Return native local terms for site or qMERA energy workflows. diff --git a/tests/test_bp_open_series.py b/tests/test_bp_open_series.py index 4e90f88..9f6a19b 100644 --- a/tests/test_bp_open_series.py +++ b/tests/test_bp_open_series.py @@ -4,12 +4,17 @@ from itertools import combinations import numpy as np +import pytest import quimb.tensor as qtn from pepsy.bp import ( + OpenLoopEnumerationLimitError, + OpenLoopObservableTerm, OpenLoopSeriesCache, + OpenLoopSeriesDiagnosticCache, OpenLoopSeriesSweepResult, compute_local_expectation_open_loop_series, + diagnose_open_loop_series, partial_trace_open_loop_series_expand, partial_trace_open_loop_series_sweep, two_norm_bp, @@ -80,6 +85,256 @@ def test_open_rho_series_keeps_the_long_range_path_and_is_exact_on_a_tree(): np.testing.assert_allclose(np.trace(rho), 1.0, atol=1e-12) +def test_open_terms_stream_shortest_support_path_first(): + state = qtn.PEPS.rand( + 1, + 4, + bond_dim=2, + phys_dim=2, + seed=1908, + dtype="complex128", + ) + where = ((0, 0), (0, 3)) + tags = [state.site_tag(coo) for coo in where] + tids = frozenset(state._get_tids_from_tags(tags, "any")) + excluded_edges = frozenset(state._select_tids(tids).inner_inds()) + iterator = OpenLoopSeriesCache().iter_terms_for( + state, + 3, + tids, + excluded_edges=excluded_edges, + ) + + first = next(iterator) + assert first.degree == 3 + assert _open_term_family(state, first) == "open_path" + + +def test_open_series_enumeration_limits_raise_before_partial_contraction(): + state = qtn.PEPS.rand( + 1, + 4, + bond_dim=2, + phys_dim=2, + seed=1918, + dtype="complex128", + ) + where = ((0, 0), (0, 3)) + with pytest.raises(OpenLoopEnumerationLimitError, match="max_terms"): + partial_trace_open_loop_series_expand( + state, + where, + edge_cutoff=3, + max_terms=0, + max_iterations=200, + tol=1e-10, + diis=False, + ) + with pytest.raises( + OpenLoopEnumerationLimitError, + match="max_enumeration_memory", + ): + partial_trace_open_loop_series_expand( + state, + where, + edge_cutoff=3, + max_enumeration_memory=1, + max_iterations=200, + tol=1e-10, + diis=False, + ) + + +def test_edge_cutoff_is_the_explicit_name_for_dense_open_series(): + state = qtn.PEPS.rand( + 1, + 4, + bond_dim=2, + phys_dim=2, + seed=1919, + dtype="complex128", + ) + where = ((0, 0), (0, 3)) + named = partial_trace_open_loop_series_expand( + state, + where, + edge_cutoff=3, + max_iterations=200, + tol=1e-10, + diis=False, + ) + legacy = partial_trace_open_loop_series_expand( + state, + where, + gloops=3, + max_iterations=200, + tol=1e-10, + diis=False, + ) + np.testing.assert_allclose(named, legacy, rtol=1e-10, atol=1e-12) + + +def test_open_series_rejects_mixed_edge_and_cluster_cutoffs(): + state = qtn.PEPS.rand(1, 4, bond_dim=2, phys_dim=2, seed=1921) + with pytest.raises(TypeError, match="edge_cutoff and cluster_size"): + partial_trace_open_loop_series_expand( + state, + ((0, 0), (0, 3)), + edge_cutoff=3, + cluster_size=4, + max_iterations=100, + ) + + +def test_corridor_mode_keeps_shortest_paths_and_local_loop_decorations(): + state = qtn.PEPS.rand( + 3, + 3, + bond_dim=2, + phys_dim=2, + seed=1936, + dtype="complex128", + ) + info = {} + partial_trace_open_loop_series_expand( + state, + ((0, 0), (2, 2)), + corridor_width=1, + max_path_candidates=4, + loop_decoration_size=4, + loop_radius=2, + corridor_segment_length=2, + max_loop_clusters_per_segment=3, + max_iterations=200, + tol=1e-10, + diis=False, + info=info, + ) + + corridor = info["open_rho_corridor"] + assert corridor["path_count"] == 4 + assert corridor["shortest_path_length"] == 4 + assert corridor["search_backend"] == "rectangular_grid" + assert corridor["corridor_edges"] < len(state.ind_map) + assert info["open_rho_family_counts"]["open_path"] == 4 + assert info["open_rho_family_counts"]["closed_loop"] + assert info["open_rho_family_counts"]["path_plus_loop"] + + +def test_corridor_mode_limits_geometry_before_contraction(): + state = qtn.PEPS.rand(3, 3, bond_dim=2, phys_dim=2, seed=1937) + with pytest.raises(OpenLoopEnumerationLimitError, match="max_corridor_edges"): + partial_trace_open_loop_series_expand( + state, + ((0, 0), (2, 2)), + corridor_width=1, + max_corridor_edges=1, + max_iterations=200, + tol=1e-10, + diis=False, + ) + + +def test_corridor_mode_can_use_compressed_boundary_contraction(): + state = qtn.PEPS.rand( + 1, + 4, + bond_dim=2, + phys_dim=2, + seed=1938, + dtype="complex128", + ) + where = ((0, 0), (0, 3)) + gate = np.diag([1.0, 2.0, 3.0, 4.0]) + info = {} + value = compute_local_expectation_open_loop_series( + state, + {where: gate}, + corridor_width=0, + max_path_candidates=1, + loop_decoration_size=4, + corridor_max_bond=4, + max_iterations=200, + tol=1e-10, + diis=False, + info=info, + ) + + assert np.isfinite(value) + assert info["open_scalar_corridor"]["path_count"] == 1 + + +def test_open_measurement_diagnostic_selects_auto_route_and_reuses_terms(): + state = qtn.PEPS.rand( + 1, + 4, + bond_dim=2, + phys_dim=2, + seed=1944, + dtype="complex128", + ) + where = ((0, 0), (0, 3)) + gate = np.diag([1.0, 2.0, 3.0, 4.0]) + diagnostic_cache = OpenLoopSeriesDiagnosticCache() + diagnostic = diagnose_open_loop_series( + state, + OpenLoopObservableTerm(where, gate), + edge_cutoff=3, + mode="auto", + auto_corridor_distance=1, + max_path_candidates=1, + loop_decoration_size=2, + max_loop_clusters_per_segment=1, + max_iterations=200, + tol=1e-10, + diis=False, + diagnostic_cache=diagnostic_cache, + ) + + support = diagnostic.for_support(where) + assert diagnostic.routes == {where: "corridor"} + assert support["terms"] + assert support["term_costs"] + assert diagnostic.total_flops_log10 is not None + assert diagnostic.peak_memory_log2 is not None + + info = {} + value = compute_local_expectation_open_loop_series( + state, + {where: gate}, + edge_cutoff=3, + mode="auto", + auto_corridor_distance=1, + max_path_candidates=1, + loop_decoration_size=2, + max_loop_clusters_per_segment=1, + diagnostic=diagnostic, + max_iterations=200, + tol=1e-10, + diis=False, + info=info, + ) + assert np.isfinite(value) + assert info["open_scalar_mode"] == "auto" + assert tuple(info["open_scalar_requested_terms"]) == support["terms"] + + cached = diagnose_open_loop_series( + state, + {where: gate}, + edge_cutoff=3, + mode="auto", + auto_corridor_distance=1, + max_path_candidates=1, + loop_decoration_size=2, + max_loop_clusters_per_segment=1, + max_iterations=200, + tol=1e-10, + diis=False, + diagnostic_cache=diagnostic_cache, + ) + assert cached.cache_hits == 1 + + def test_open_scalar_series_inserts_a_two_site_gate_and_normalizes_after_sum(): state = qtn.PEPS.rand( 1, @@ -305,6 +560,9 @@ def test_open_rho_series_reports_term_families_and_reuses_one_bp_run(): } assert info["open_rho_edge_term_costs"] assert not info["open_rho_cluster_region_costs"] + assert len(info["open_rho_region_path_cache"]) <= len( + info["open_rho_terms_list"] + ) np.testing.assert_allclose(np.trace(rho), 1.0, atol=1e-12) other_info = {} @@ -322,6 +580,31 @@ def test_open_rho_series_reports_term_families_and_reuses_one_bp_run(): assert len(cache.terms_by_key) == 2 +def test_open_series_reuses_contraction_paths_for_shared_regions(): + state = qtn.PEPS.rand( + 3, + 2, + bond_dim=2, + phys_dim=2, + cyclic=(True, True), + seed=1922, + dtype="complex128", + ) + info = {} + partial_trace_open_loop_series_expand( + state, + ((0, 0), (2, 1)), + edge_cutoff=4, + max_iterations=200, + tol=1e-10, + diis=False, + info=info, + ) + assert len(info["open_rho_region_path_cache"]) < len( + info["open_rho_terms_list"] + ) + + def test_open_rho_series_reuses_one_d2bp_message_set(): state = qtn.PEPS.rand( 1, @@ -432,6 +715,28 @@ def test_open_rho_series_sweep_reuses_bp_across_supports_and_cutoffs(): assert result.diagnostics[tuple(support)][cutoff]["term_count"] >= 0 +def test_open_rho_series_sweep_accepts_route_specific_edge_cutoffs(): + state = qtn.PEPS.rand( + 1, + 4, + bond_dim=2, + phys_dim=2, + seed=1923, + dtype="complex128", + ) + result = partial_trace_open_loop_series_sweep( + state, + (((0, 0), (0, 3)),), + edge_cutoffs=(0, 3), + max_iterations=200, + tol=1e-10, + diis=False, + ) + + assert result.get_rho(((0, 0), (0, 3)), 0) is not None + assert result.get_rho(((0, 0), (0, 3)), 3) is not None + + def test_open_rho_series_is_exact_for_a_tree_with_a_multi_site_support(): state = qtn.PEPS.rand( 1, diff --git a/tests/test_bp_symmray.py b/tests/test_bp_symmray.py index 793dad3..ec2f59c 100644 --- a/tests/test_bp_symmray.py +++ b/tests/test_bp_symmray.py @@ -4,6 +4,7 @@ import numpy as np import pytest +import pepsy.bp.series as bp_series sr = pytest.importorskip("symmray") qtn = pytest.importorskip("quimb.tensor") @@ -32,8 +33,10 @@ two_norm_bp, weight_pass, ) +from pepsy.operators.gates import gate_simple # noqa: E402 from pepsy.tensors import ( # noqa: E402 Fermion, + OneDMap, SymPEPS, ps_to_peps, site_charge_alternating, @@ -1125,6 +1128,60 @@ def test_fermionic_u1_loop_cluster_runs_on_3x4_peps(bond_dim): ) +@pytest.mark.parametrize("symmetry", ["U1U1", "U1", "Z2"]) +def test_native_fermionic_pepo_simple_update_norm_and_loop_cluster(symmetry): + """Native PEPO evolution and D2BP norm correction compose cleanly.""" + fermion = Fermion(spinful=True, symmetry=symmetry) + left, right = (0, 0), (0, 1) + term = fermion.operator_term( + [(1.0, ((left, "create_up"), (right, "annihilate_up")))], + sites=(left, right), + add_hc=True, + ) + pepo = fermion.to_pepo( + {(left, right): term}, + Lx=2, + Ly=2, + mapper=OneDMap(2, 2, mode="snake-row-major"), + max_bond=16, + compress=False, + ) + gauges = {} + pepo.gauge_all_simple_(gauges=gauges, progbar=False) + norm_before = complex(pepo.norm()) + + evolved = gate_simple( + pepo, + fermion.hopping_gate(0.001, t=1.0).H, + where=(left, right), + gauges=gauges, + max_bond=16, + cutoff=1e-10, + contract="split", + renorm=False, + inplace=False, + ) + norm_after = complex(evolved.norm()) + correction = loop_cluster_expand( + evolved, + gloops=2, + norm="2norm", + max_iterations=100, + tol=1e-8, + diis=False, + progbar=False, + ) + + assert len(gauges) > 0 + assert np.isfinite(norm_before.real) + assert np.isfinite(norm_after.real) + assert np.isfinite(float(np.real(correction.estimate))) + assert all( + type(tensor.data).__name__.endswith("FermionicArray") + for tensor in evolved + ) + + def _fermionic_symmetry_cases(): return ( ( @@ -2006,6 +2063,64 @@ def test_cyclic_native_open_series_honors_contraction_cost_limits(): assert cost["peak_memory_log2"] <= 30.0 +def test_cyclic_native_route_skips_open_edge_enumeration(monkeypatch): + """Cluster fallback is selected before combinatorial edge discovery.""" + state = SymPEPS.random( + 2, + 3, + symmetry="U1", + bond_dim=2, + phys_dim=2, + fermionic=True, + cyclic=(True, True), + seed=2201, + dtype="complex128", + ) + where = ((0, 0), (1, 2)) + bp = two_norm_bp( + state.tn, + max_iterations=100, + tol=1e-9, + diis=False, + ) + + def fail_if_enumerated(*args, **kwargs): + raise AssertionError("cyclic native route enumerated edge terms") + + monkeypatch.setattr( + bp_series, + "_iter_open_edge_loops", + fail_if_enumerated, + ) + rho_info = {} + rho = partial_trace_open_loop_series_expand( + state.tn, + where, + cluster_size=3, + max_terms=0, + messages=bp.messages, + run_bp=False, + info=rho_info, + ) + assert rho_info["open_rho_native_route"] == "graded_cluster_compatible" + np.testing.assert_allclose(np.trace(rho.to_dense()), 1.0, atol=1e-12) + + scalar_info = {} + value = compute_local_expectation_open_loop_series( + state.tn, + {where: Fermion(spinful=False, symmetry="U1").density_operator()}, + cluster_size=3, + max_terms=0, + messages=bp.messages, + run_bp=False, + info=scalar_info, + ) + assert scalar_info["open_scalar_native_route"] == ( + "graded_cluster_compatible" + ) + assert np.isfinite(value) + + def test_explicit_edge_loop_series_preserves_dense_edge_degree_terms(): """Edge-degree terms are distinct from the local-region cutoff API.""" state = qtn.PEPS.rand(2, 2, bond_dim=2, seed=1906, dtype="complex128") diff --git a/tests/test_gate.py b/tests/test_gate.py index bd07d75..f655195 100644 --- a/tests/test_gate.py +++ b/tests/test_gate.py @@ -22,6 +22,7 @@ y, z, ) +from pepsy.tensors.core import OneDMap def _dense_numpy(tn, out_inds): @@ -1087,6 +1088,85 @@ def test_build_mpo_from_gates_accepts_single_gate_where(): assert mpo.max_bond() >= 1 +@pytest.mark.parametrize("symmetry", ["U1U1", "U1", "Z2"]) +def test_native_fermion_gate_builders_preserve_symmetry(symmetry): + """Gate-to-MPO/PEPO builders accept native fermionic gates.""" + pytest.importorskip("symmray") + from pepsy.tensors.symmetric import Fermion + + fermion = Fermion(spinful=True, symmetry=symmetry) + hopping_gate = fermion.hopping_gate(0.01, t=1.0) + mpo = build_mpo_from_gates( + hopping_gate, + where=(0, 1), + max_bond=8, + contract="split", + ) + pepo = build_pepo_from_gates( + hopping_gate, + where=((0, 0), (0, 1)), + mapper=OneDMap(2, 2, mode="snake-row-major"), + max_bond=8, + contract="split", + ) + + assert all(type(tensor.data).__name__.endswith("FermionicArray") for tensor in mpo) + assert all(type(tensor.data).__name__.endswith("FermionicArray") for tensor in pepo) + assert pepo.Lx == 2 + assert pepo.Ly == 2 + + +@pytest.mark.parametrize("symmetry", ["U1U1", "U1", "Z2"]) +def test_native_charged_gate_builders_require_opt_in(symmetry): + """Charged native gate streams work when explicitly enabled.""" + pytest.importorskip("symmray") + from pepsy.tensors.symmetric import Fermion, symmray_mpo_summary + + fermion = Fermion(spinful=True, symmetry=symmetry) + charged_gate = fermion.operator_term( + [(1.0, ((0, "double"), (1, "annihilate_up")))], + sites=(0, 1), + label="charged_gate_builder", + ) + + with pytest.raises(ValueError, match="allow_charged=True"): + build_mpo_from_gates( + charged_gate, + where=(0, 1), + max_bond=16, + contract="split", + ) + with pytest.raises(ValueError, match="allow_charged=True"): + build_pepo_from_gates( + charged_gate, + where=((0, 0), (0, 1)), + mapper=OneDMap(2, 2, mode="snake-row-major"), + max_bond=16, + contract="split", + ) + + mpo = build_mpo_from_gates( + charged_gate, + where=(0, 1), + allow_charged=True, + max_bond=16, + contract="split", + ) + pepo = build_pepo_from_gates( + charged_gate, + where=((0, 0), (0, 1)), + mapper=OneDMap(2, 2, mode="snake-row-major"), + allow_charged=True, + max_bond=16, + contract="split", + ) + + assert symmray_mpo_summary(mpo)["total_charge"] == charged_gate.charge + assert all(type(tensor.data).__name__.endswith("FermionicArray") for tensor in mpo) + assert all(type(tensor.data).__name__.endswith("FermionicArray") for tensor in pepo) + assert any(tensor.data.charge != fermion.zero_charge for tensor in pepo) + + def test_build_mpo_from_gates_forwards_max_bond_to_gate(monkeypatch): """MPO builder should use the public max_bond gate API.""" calls = [] diff --git a/tests/test_ham.py b/tests/test_ham.py index 22457f7..2880947 100644 --- a/tests/test_ham.py +++ b/tests/test_ham.py @@ -49,6 +49,83 @@ def test_build_mpo_accepts_mapper_override(): assert np.allclose(mpo_from_override.to_dense(), mpo_from_mapper_builder.to_dense()) +def test_build_mpo_and_pepo_accept_native_fermion_terms_with_mapper(): + """Hamiltonian builders forward OneDMap and native fermion terms together.""" + pytest.importorskip("symmray") + fermion = py.Fermion(spinful=True, symmetry="U1U1") + left = (0, 0) + right = (2, 1) + term = fermion.operator_term( + [(1.0, ((left, "double"), (right, "annihilate_up")))], + sites=(left, right), + label="ham_builder_charged", + ) + mapper = OneDMap(3, 2, mode="snake-row-major") + builder = py.ham_tn( + Lx=3, + Ly=2, + mapper=mapper, + data_type="complex128", + ) + + mpo = builder.build_mpo( + {(left, right): term}, + fermion=fermion, + fermionic=True, + compress_each=False, + ) + pepo = builder.build_pepo( + {(left, right): term}, + fermion=fermion, + fermionic=True, + compress_each=False, + ) + + assert mpo.L == 6 + assert pepo.Lx == 3 + assert pepo.Ly == 2 + assert list(pepo)[-1].data.charge == term.charge + assert all(type(tensor.data).__name__.endswith("FermionicArray") for tensor in pepo) + + +def test_build_mpo_and_pepo_return_mixed_charge_sectors(): + """Hamiltonian builders expose mixed native charges explicitly.""" + pytest.importorskip("symmray") + fermion = py.Fermion(spinful=True, symmetry="U1U1") + left = (0, 0) + middle = (0, 1) + right = (1, 1) + neutral = fermion.hopping_operator() + charged = fermion.operator_term( + [(1.0, ((middle, "double"), (right, "annihilate_up")))], + sites=(middle, right), + label="ham_builder_mixed_charge", + ) + terms = {(left, middle): neutral, (middle, right): charged} + mapper = OneDMap(2, 2, mode="snake-row-major") + builder = py.ham_tn(Lx=2, Ly=2, mapper=mapper, data_type="complex128") + + mpo_sectors = builder.build_mpo( + terms, + fermion=fermion, + fermionic=True, + charge_sectors=True, + compress_each=False, + ) + pepo_sectors = builder.build_pepo( + terms, + fermion=fermion, + fermionic=True, + charge_sectors=True, + compress_each=False, + ) + + assert set(mpo_sectors) == {fermion.zero_charge, charged.charge} + assert set(pepo_sectors) == {fermion.zero_charge, charged.charge} + assert all(mpo.L == 4 for mpo in mpo_sectors.values()) + assert all(pepo.Lx == 2 and pepo.Ly == 2 for pepo in pepo_sectors.values()) + + def test_build_mpo_uses_canonical_ops_sites_coeff_order(): """build_mpo should accept the canonical (ops, sites, coeff) term order.""" builder = py.ham_tn(Lx=2, Ly=2, data_type="complex128") diff --git a/tests/test_symmetric_tensors.py b/tests/test_symmetric_tensors.py index 1e6809c..f2823e4 100644 --- a/tests/test_symmetric_tensors.py +++ b/tests/test_symmetric_tensors.py @@ -1115,6 +1115,238 @@ def test_fermion_explicit_coordinate_terms_preserve_peps_locations(): assert mpo.L == 4 +def test_symhamiltonian_and_fermion_build_pepo_accept_mapper(): + """Hamiltonian and model shorthands share the native PEPO route.""" + pytest.importorskip("symmray") + fermion = Fermion(spinful=True, symmetry="U1U1") + left = (0, 0) + right = (1, 1) + term = fermion.operator_term( + [(1.0, ((left, "double"), (right, "annihilate_up")))], + sites=(left, right), + label="symhamiltonian_pepo", + ) + hamiltonian = fermion.hamiltonian({(left, right): term}) + mapper = OneDMap(2, 2, mode="snake-row-major") + + pepo_from_hamiltonian = hamiltonian.to_pepo( + 2, + 2, + mapper=mapper, + max_bond=16, + compress=False, + ) + pepo_from_model = fermion.build_pepo( + hamiltonian=hamiltonian, + Lx=2, + Ly=2, + mapper=mapper, + max_bond=16, + compress=False, + ) + + for pepo in (pepo_from_hamiltonian, pepo_from_model): + assert pepo.Lx == 2 + assert pepo.Ly == 2 + assert list(pepo)[-1].data.charge == term.charge + assert all( + type(tensor.data).__name__.endswith("FermionicArray") + for tensor in pepo + ) + + +@pytest.mark.parametrize("symmetry", ["U1U1", "U1", "Z2"]) +def test_mixed_native_charges_return_explicit_mpo_and_pepo_sectors(symmetry): + """Mixed native operators decompose into homogeneous charge sectors.""" + pytest.importorskip("symmray") + fermion = Fermion(spinful=True, symmetry=symmetry) + left = (0, 0) + middle = (0, 1) + right = (1, 1) + neutral = fermion.hopping_operator() + charged = fermion.operator_term( + [(1.0, ((middle, "double"), (right, "annihilate_up")))], + sites=(middle, right), + label="mixed_charge_sector", + ) + hamiltonian = fermion.hamiltonian( + { + (left, middle): neutral, + (middle, right): charged, + } + ) + mapper = OneDMap(2, 2, mode="snake-row-major") + + mpo_sectors = hamiltonian.to_mpo( + mapper=mapper, + fermionic=True, + charge_sectors=True, + compress=False, + ) + pepo_sectors = hamiltonian.to_pepo( + 2, + 2, + mapper=mapper, + fermionic=True, + charge_sectors=True, + compress=False, + ) + + assert set(mpo_sectors) == {fermion.zero_charge, charged.charge} + assert set(pepo_sectors) == {fermion.zero_charge, charged.charge} + for charge, mpo in mpo_sectors.items(): + assert all(type(tensor.data).__name__.endswith("FermionicArray") for tensor in mpo) + assert list(mpo)[-1].data.charge == charge + for pepo in pepo_sectors.values(): + assert pepo.Lx == 2 + assert pepo.Ly == 2 + assert all( + type(tensor.data).__name__.endswith("FermionicArray") + for tensor in pepo + ) + + +@pytest.mark.parametrize("symmetry", ["U1U1", "U1", "Z2"]) +def test_fermion_to_pepo_builds_native_coordinate_terms(symmetry): + """Fermion.to_pepo preserves native grading for supported symmetries.""" + fermion = Fermion(spinful=True, symmetry=symmetry) + left = (0, 1) + right = (2, 2) + hopping = fermion.operator_term( + [(1.0, ((left, "create_up"), (right, "annihilate_up")))], + sites=(left, right), + add_hc=True, + ) + + pepo = fermion.to_pepo( + {(left, right): hopping}, + Lx=3, + Ly=3, + max_bond=16, + compress=False, + ) + + assert pepo.Lx == 3 + assert pepo.Ly == 3 + assert set(pepo.outer_inds()) == { + f"k{x},{y}" for x in range(3) for y in range(3) + } | { + f"b{x},{y}" for x in range(3) for y in range(3) + } + assert all(type(tensor.data).__name__.endswith("FermionicArray") for tensor in pepo) + + +def test_fermion_to_pepo_native_result_supports_reverse_simple_update(): + """Native PEPO output can take an adjoint gate through operator SU.""" + fermion = Fermion(spinful=True, symmetry="U1U1") + left = (0, 0) + right = (0, 1) + term = fermion.operator_term( + [(1.0, ((left, "create_up"), (right, "annihilate_up")))], + sites=(left, right), + add_hc=True, + ) + pepo = fermion.to_pepo( + {(left, right): term}, + Lx=2, + Ly=2, + max_bond=16, + compress=False, + ) + gauges = {} + pepo.gauge_all_simple_(gauges=gauges, progbar=False) + + out = gate_simple( + pepo, + fermion.hopping_gate(0.001, t=1.0).H, + where=(left, right), + gauges=gauges, + max_bond=16, + cutoff=1e-10, + contract="split", + renorm=False, + inplace=False, + ) + + assert out.max_bond() <= 16 + assert all(type(tensor.data).__name__.endswith("FermionicArray") for tensor in out) + assert len(gauges) > 0 + + +@pytest.mark.parametrize("symmetry", ["U1U1", "U1", "Z2"]) +def test_fermion_to_pepo_supports_charged_odd_native_terms(symmetry): + """Charged odd terms retain their native charge and dummy mode.""" + fermion = Fermion(spinful=True, symmetry=symmetry) + left = (0, 1) + right = (2, 2) + charged = fermion.operator_term( + [(1.0, ((left, "double"), (right, "annihilate_up")))], + sites=(left, right), + label="charged_pepo", + ) + + pepo = fermion.to_pepo( + {(left, right): charged}, + Lx=3, + Ly=3, + max_bond=16, + compress=False, + ) + tensors = list(pepo) + + assert charged.charge != fermion.zero_charge + assert tensors[-1].data.charge == charged.charge + assert tensors[-1].data.label is not None + assert tensors[-1].data.dummy_modes + assert all(type(tensor.data).__name__.endswith("FermionicArray") for tensor in tensors) + with pytest.raises(ValueError, match="fermionic=True"): + fermion.to_mpo( + {(1, 8): charged}, + L=9, + fermionic=False, + ) + + +@pytest.mark.parametrize("symmetry", ["U1U1", "U1", "Z2"]) +def test_charged_native_pepo_supports_reverse_simple_update(symmetry): + """Neutral reverse evolution preserves a charged PEPO operator sector.""" + fermion = Fermion(spinful=True, symmetry=symmetry) + left = (0, 0) + right = (0, 1) + charged = fermion.operator_term( + [(1.0, ((left, "double"), (right, "annihilate_up")))], + sites=(left, right), + label="charged_pepo", + ) + pepo = fermion.to_pepo( + {(left, right): charged}, + Lx=2, + Ly=2, + max_bond=16, + compress=False, + ) + gauges = {} + pepo.gauge_all_simple_(gauges=gauges, progbar=False) + + out = gate_simple( + pepo, + fermion.hopping_gate(0.001, t=1.0).H, + where=(left, right), + gauges=gauges, + max_bond=16, + cutoff=1e-10, + contract="split", + renorm=False, + inplace=False, + ) + + tensors = list(out) + assert tensors[-1].data.charge == charged.charge + assert tensors[-1].data.dummy_modes + assert out.max_bond() <= 16 + assert len(gauges) > 0 + + def test_unified_fermion_peps_energy_accepts_boundary_chi(): """The shared Fermion Hamiltonian can use SymPEPS boundary measurement.""" peps = SymPEPS.for_model( From e84cd1f1d5c4ed855b21fee25ba9f6b38303150c Mon Sep 17 00:00:00 2001 From: rezaquant Date: Fri, 31 Jul 2026 11:06:11 -0700 Subject: [PATCH 40/70] Improve tree quality layout search --- docs/api/optimizers/tree.md | 83 +++-- docs/api/tensors/symmetric.md | 20 + docs/development/modules/tensors.md | 20 +- src/pepsy/optimizers/tree/layout.py | 465 +++++++++++++++++++++--- src/pepsy/optimizers/tree/optimizer.py | 484 ++++++++++++++++++------- src/pepsy/tensors/maps.py | 188 +++++++++- tests/test_ham.py | 40 ++ tests/test_optimize_tree.py | 198 ++++++++-- 8 files changed, 1251 insertions(+), 247 deletions(-) diff --git a/docs/api/optimizers/tree.md b/docs/api/optimizers/tree.md index 02eb779..340d36f 100644 --- a/docs/api/optimizers/tree.md +++ b/docs/api/optimizers/tree.md @@ -428,17 +428,22 @@ every hierarchical scale, not only the root cut. It enables bounded subtree reconfiguration and simulated annealing by default; override these with `topology_refine="subtree"`, `topology_budget=`, `search="anneal"`, and `search_budget=`. The result is still a cheap layout proxy rather than a real -TTN replay, so the state-aware pilot remains the final accuracy check. +TTN replay, so the state-aware pilot remains the final accuracy check. The +default `chi=None` leaves this as a static, chi-blind objective; supplying +`chi` only adds cap-aware ranking and does not change the no-tensor nature of +layout discovery. Use `order="quality"` with `finder.run()` (or set it on the finder) for the -MPS-style higher-quality offline search. It enables bounded greedy leaf -refinement, bounded binary-tree nearest-neighbor-interchange (NNI) topology -refinement, and opportunistic Nevergrad refinement when Nevergrad is installed; -otherwise it uses the deterministic NNI and leaf stages. NNI changes the -internal grouping itself, so quality mode can improve a tree even when the -best leaf labels are already fixed. The zero-argument `run()` path remains the -fast deterministic candidate selection. Disable the topology stage explicitly -with `topology_refine=None`, or bound it with `topology_budget=`. +MPS-style high-quality offline search. Quality mode now means +`objective="full_tree"`: it evaluates every hierarchy scale, enables bounded +greedy leaf refinement and all-scale subtree topology refinement, and runs a +hybrid topology-annealing/Nevergrad search when Nevergrad is available. It +falls back explicitly to dependency-free simulated annealing otherwise. A +finder without `order="quality"` keeps the fast deterministic zero-argument +`run()` path. +Disable stages explicitly with `refine=None`, `topology_refine=None`, or +`search=None`, or bound them with `refine_budget=`, `topology_budget=`, and +`search_budget=`. For a stream whose locality changes over time, pass `time_decay=` and/or `time_window=` to `TreeLayoutFinder`. A decay in `(0, 1]` weights an event by @@ -587,11 +592,14 @@ well-aligned physical span `r` into a path with `O(log r)` tree hops, so the hybrid score uses path length as a replay-cost proxy while edge loads estimate the accuracy/bond-dimension cost. -For offline quality searches, add `search="nevergrad"`. Nevergrad starts from -the spectral/greedy plan, proposes leaf orders, and keeps its result only when -it improves the same chi-aware objective. It never acts on a live optimizer. -For `objective="full_tree"`, use `search="anneal"` to explore subtree -replacements at multiple scales without the optional Nevergrad dependency. +For offline quality searches, `search="nevergrad"` starts from the +spectral/greedy plan, proposes leaf orders, and keeps its result only when it +improves the same objective. `search="anneal"` explores subtree replacements +at multiple scales without optional dependencies. For the highest-quality +`objective="full_tree"` search, use `search="hybrid"`: it splits the bounded +budget between topology annealing and Nevergrad leaf refinement. Quality mode +selects this hybrid automatically when Nevergrad is installed and falls back +to annealing otherwise. None of these stages allocates or replays a TTN. Install the optional dependency with `pip install pepsy[layout]`: ```python @@ -615,7 +623,7 @@ tree_plan = finder.run( topology_budget=64, refine="greedy", refine_budget=64, - search="nevergrad", + search="hybrid", search_budget=128, seed=0, nevergrad_optimizer="OnePlusOne", @@ -663,13 +671,36 @@ opt = py.TreeOptimizer(gate_stream, tree=choice["plan"], chi=chi) The pilot replays candidates on independent copies with the real tree update kernels and returns measured infidelity, final bond, truncation count, and runtime under `choice["pilot"]`. By default one bounded `order="quality"` -candidate (greedy leaf refinement plus binary NNI topology refinement) is +candidate (greedy leaf refinement plus topology refinement; for +`full_tree`, this also includes bounded subtree/hybrid search) is reserved a pilot slot, so it cannot be rejected before state-aware replay. Use `include_quality=False` for the static-only candidate set. The original optimizer is unchanged unless `install=True` is passed. Installation is restricted to product initial states; an entangled TTN cannot generally be relaid out exactly. +For the recommended closed-loop choice, use the high-level optimizer helper: + +```python +choice = opt.optimize_layout( + objective="full_tree", + rounds=2, + pilot_candidates=4, + pilot_steps=64, + topology_budget=32, + search_budget=64, +) +``` + +This keeps the finder tensor-free, pilots the selected candidates with the +actual tree replay kernels, and uses each round's per-edge truncation loss, +discarded weight, and update runtime to seed bounded NNI, subtree, and +cross-cut leaf proposals for the next round. `choice["pilot"]["rounds"]` +contains every round and `report["edge_diagnostics"]` identifies hot tree +edges. `objective="full_tree"` combines all-scale static work/bond estimates +with this short state-aware replay. `install=True` remounts the product state +on the final plan; it remains rejected for an entangled state. + Both helpers are also available from the package-level API: ```python @@ -755,13 +786,19 @@ finder and optimizer expose diagnostics to choose it: The same diagnostics are available as a Cotengra-style tent plot. `TreeLayoutFinder.plot(plan)` is the default tent view (also available as `plot_tent(plan)`): it keeps the raw graph at the bottom and lifts -the selected hierarchy above its descendant sites: the raw lattice and gate -connectivity are gray, while circular nodes use stable scale colors. Hierarchy -edges use one uniform solid color by default. Pass `edge_color=None` to make -each incoming edge exactly match the node it terminates at; `node_cmap` then -controls both. Arrows are disabled by default, matching Cotengra's -structural tent view; pass `show_edge_arrows=True` only when parent-to-child -direction is needed. +the selected hierarchy above its descendant sites: the raw lattice and +optional gate connectivity are gray, while internal tree nodes use a stable order-based +`turbo` palette by default. Incoming edges match their child nodes by default; +pass an explicit `edge_color` for a uniform structural color. Arrows are +disabled by default, matching Cotengra's structural tent view; pass +`show_edge_arrows=True` only when parent-to-child direction is needed. +When the physical background already has its own markers, pass +`show_leaf_nodes=False` to hide the tree's physical leaf circles while keeping +internal tree nodes and hierarchy edges visible. This is useful for a gray +`+`-marked lattice backdrop. +When `site_coords` are supplied, the default tent presentation projects them +with `lattice_skew=0.30` and `lattice_rise=0.18`, and draws gray `+` markers at +the physical sites. Override those values to use a different base projection. Nearest-neighbor gate edges are not duplicated over the lattice. Supplying `site_coords={qubit: (x, y)}` places the physical sites on an existing lattice. It returns `(fig, ax)` and does not mutate the plan or live TTN: diff --git a/docs/api/tensors/symmetric.md b/docs/api/tensors/symmetric.md index 15a1788..9ff6233 100644 --- a/docs/api/tensors/symmetric.md +++ b/docs/api/tensors/symmetric.md @@ -827,6 +827,26 @@ Symmray-compatible), so it is deferred. For 2D today: order the lattice with ``to_mpo`` for the residual long-range terms, or evolve via a Symmray MPO-TDVP sweep (which preserves U1xU1 without gates). +When the circuit gate stream should choose the 1D path, use the MPS layout +finder as a mapping mode. This remains a coordinate/index operation only: it +does not allocate an MPS or perform replay, SVD, or truncation. + +```python +mapper = py.OneDMap( + Lx=6, + Ly=6, + mode="finder", + gate_stream=gates, + layout_kwargs={"objective": "compression", "order": "quality"}, +) +idx2coo, coo2idx = mapper.build() +``` + +The finder assumes the stream uses compact logical labels in +`range(Lx * Ly)`. Set `finder_base_mode="row-major"` (or another regular +mode) when those labels come from a different initial traversal. A previously +computed MPS layout plan may be passed as `finder=plan`. + For a flattened MPS path, feed the same canonical bundled stream to ``MpsOptimizer``: diff --git a/docs/development/modules/tensors.md b/docs/development/modules/tensors.md index 5c9934c..0a8ebf2 100644 --- a/docs/development/modules/tensors.md +++ b/docs/development/modules/tensors.md @@ -26,7 +26,25 @@ structure unless a change has a strong reason to split implementation. `OneDMap` maps regular 2D or 3D lattice coordinates onto a 1D path. Supported modes include `snake`, `snake-row-major`, `row-major`, `col-major`, `hilbert`, -`hilbert-row-major`, and `diag`. +`hilbert-row-major`, and `diag`. The additional `finder` mode composes an MPS +gate-stream layout permutation with a base lattice mode. It analyzes only gate +supports and does not construct, replay, or truncate an MPS: + +```python +mapper = py.OneDMap( + 6, + 6, + mode="finder", + gate_stream=gates, + layout_kwargs={"objective": "compression", "order": "quality"}, +) +idx2coo, coo2idx = mapper.build() +``` + +The gate stream's logical labels must be the compact integers `0..Lx*Ly-1`. +Use `finder_base_mode="row-major"` when those labels were originally assigned +by a different regular traversal. An existing MPS layout plan can be supplied +with `finder=plan` instead of `gate_stream=`. Constructors create common tensor-network states and operators: diff --git a/src/pepsy/optimizers/tree/layout.py b/src/pepsy/optimizers/tree/layout.py index 198fb38..bb563e1 100644 --- a/src/pepsy/optimizers/tree/layout.py +++ b/src/pepsy/optimizers/tree/layout.py @@ -167,10 +167,15 @@ def _normalize_layout_search(search): "simulated_annealing": "anneal", "subtree_anneal": "anneal", "annealing": "anneal", + "quality": "hybrid", + "combined": "hybrid", + "full_tree": "hybrid", } name = aliases.get(name, name) - if name not in {"nevergrad", "anneal"}: - raise ValueError("search must be None, 'nevergrad', or 'anneal'.") + if name not in {"nevergrad", "anneal", "hybrid"}: + raise ValueError( + "search must be None, 'nevergrad', 'anneal', or 'hybrid'." + ) return name @@ -199,6 +204,11 @@ def _nevergrad_available(): return True +def _quality_search_mode(): + """Return the complete quality search, with a clear fallback.""" + return "hybrid" if _nevergrad_available() else "anneal" + + def _validate_search_budget(value, name): """Validate a positive bounded layout-search evaluation budget.""" try: @@ -1292,8 +1302,9 @@ class TreeLayoutFinder: chi : int, optional Bond-dimension budget used to bias the default arity search toward plans that stay exact at ``chi`` (see :meth:`recommend_arities`). ``None`` - keeps the search purely objective-driven. :class:`TreeOptimizer` - forwards its own ``chi`` here automatically. + keeps the search purely objective-driven and is the static layout + default; it does not allocate tensors or perform truncations. + :class:`TreeOptimizer` forwards its own ``chi`` here automatically. community_frac : float Strong-edge fraction for ``structure="adaptive"`` (see :meth:`TreePlan.from_order`). @@ -1314,12 +1325,14 @@ class TreeLayoutFinder: then applies bounded leaf and binary-topology refinement by default. `"full_tree"` evaluates dynamic bond pressure, tensor width, estimated work, write volume, and route length across every tree scale. It is - the high-quality, Cotengra-inspired mode and is opt-in because its - bounded subtree search is more expensive. + the high-quality, Cotengra-inspired mode; ``order="quality"`` selects + it automatically and enables its bounded search stages. order : {None, "quality"}, optional - Optional high-quality offline mode. `"quality"` enables bounded - greedy refinement and opportunistic Nevergrad refinement; omitted - keeps the fast deterministic candidate selection. + Optional high-quality offline mode. `"quality"` means + `objective="full_tree"` and enables bounded greedy leaf refinement, + all-scale subtree topology refinement, and hybrid + Nevergrad/annealing search. Omitted keeps the fast deterministic + objective selected by `objective`. hybrid_weights : mapping or sequence of three floats, optional Weights for the hybrid path, maximum edge load, and total edge load. The default is ``(1.0, 1.0, 0.25)``. @@ -1340,14 +1353,18 @@ class TreeLayoutFinder: Maximum topology proposals per candidate plan. Defaults to at most 64 proposals when topology refinement is enabled. For ``"subtree"``, proposals are sampled across the available descendant scales. - search : {None, "nevergrad", "anneal"} + search : {None, "nevergrad", "anneal", "hybrid"} Optional offline derivative-free refinement. It is never run unless requested. `"nevergrad"` refines leaf order and requires the optional package; `"anneal"` performs bounded simulated annealing over subtree - reconfigurations and has no additional dependency. + reconfigurations and has no additional dependency. `"hybrid"` splits + the budget between subtree annealing and Nevergrad leaf refinement; + it falls back to annealing if Nevergrad is unavailable. search_budget : int Number of offline search evaluations per candidate plan. For - ``search="anneal"``, this is the number of subtree proposals. + ``search="anneal"``, this is the number of subtree proposals; for + ``search="hybrid"``, it is the shared total split between annealing + and Nevergrad. seed : int Reproducible seed used by the optional Nevergrad stage. nevergrad_optimizer : str @@ -1473,6 +1490,13 @@ def __init__(self, gates=None, n=None, *, supports=None, structure="quality", self.hybrid_weights = _normalize_hybrid_weights(hybrid_weights) self.weight_mode = _normalize_weight_mode(weight_mode) self.order = _normalize_layout_order(order) + if self.order == "quality": + # ``order="quality"`` is the explicit high-quality contract. It + # is intentionally stronger than merely enabling a leaf swap: the + # all-scale full-tree objective and its topology search are part + # of the mode. Callers can still opt out of individual stages in + # ``run`` with ``refine=None`` or ``search=None``. + self.objective = "full_tree" self.refine = _normalize_layout_refinement(refine) if refine_budget is not None: refine_budget = _validate_search_budget(refine_budget, "refine_budget") @@ -1515,6 +1539,7 @@ def __init__(self, gates=None, n=None, *, supports=None, structure="quality", self._edge_load_cache = {} self._rank_diagnostics_cache = {} self._full_tree_profile_cache = {} + self._full_tree_structure_cache = {} self._schmidt_rank_cache = {} self._similarity_cache = {} self._congestion_weights_cache = None @@ -2380,6 +2405,71 @@ def _anneal_plan_subtree( "scales_visited": tuple(sorted(visited_scales)), } + def _refine_plan_hybrid( + self, plan, *, chi, budget, seed, optimizer_name, progbar=False + ): + """Combine all-scale topology annealing with Nevergrad leaf search. + + The two search methods explore complementary spaces: annealing changes + descendant subtree topology while Nevergrad changes the labels assigned + to a fixed topology. ``budget`` is the total number of proposals and + is split between the two stages, so hybrid quality mode remains bounded + by the caller's requested offline budget. + """ + initial_key = self._selection_key(plan, chi) + if budget < 2: + anneal_budget = budget + nevergrad_budget = 0 + else: + anneal_budget = max(1, budget // 2) + nevergrad_budget = budget - anneal_budget + + current, anneal_info = self._anneal_plan_subtree( + plan, + chi=chi, + budget=anneal_budget, + seed=seed + 1, + progbar=progbar, + ) + nevergrad_info = None + if nevergrad_budget: + try: + current, nevergrad_info = self._refine_plan_nevergrad( + current, + chi=chi, + budget=nevergrad_budget, + seed=seed + 2, + optimizer_name=optimizer_name, + progbar=progbar, + ) + except ImportError: + # Hybrid mode is the automatic quality path. Keep its + # dependency-free topology result when the optional package + # is absent; explicit search="nevergrad" still raises. + nevergrad_info = { + "method": "nevergrad", + "optimizer": optimizer_name, + "budget": nevergrad_budget, + "evaluations": 0, + "seed": seed + 2, + "available": False, + "improved": False, + } + + return current, { + "method": "hybrid", + "search": "hybrid", + "budget": budget, + "evaluations": int(anneal_info["evaluations"]) + + int((nevergrad_info or {}).get("evaluations", 0)), + "accepted_moves": int(anneal_info["accepted_moves"]) + + int((nevergrad_info or {}).get("improved", False)), + "initial_key": initial_key, + "final_key": self._selection_key(current, chi), + "anneal": anneal_info, + "nevergrad": nevergrad_info, + } + def _refine_plan_nevergrad( self, plan, *, chi, budget, seed, optimizer_name, progbar=False ): @@ -2535,6 +2625,15 @@ def _improve_plan(self, plan, *, chi, settings, progbar=False): optimizer_name=settings["nevergrad_optimizer"], progbar=progbar, ) + elif settings["search"] == "hybrid": + plan, info["search"] = self._refine_plan_hybrid( + plan, + chi=chi, + budget=settings["search_budget"], + seed=settings["seed"], + optimizer_name=settings["nevergrad_optimizer"], + progbar=progbar, + ) info["final_order"] = self._leaf_order(plan) info["final_key"] = self._selection_key(plan, chi) return plan, info @@ -2787,10 +2886,11 @@ def recommend_layered( topology_budget : int, optional Maximum topology proposals per candidate. When omitted, an enabled search uses at most 64 proposals. - search : {None, "nevergrad", "anneal"}, optional + search : {None, "nevergrad", "anneal", "hybrid"}, optional Override the finder offline search setting. Nevergrad optimizes only the returned fixed plan; annealing explores subtree - replacements. Neither mutates a live TTN. + replacements. ``"hybrid"`` combines both with one shared budget. + Neither mutates a live TTN. search_budget, seed, nevergrad_optimizer Optional offline-search configuration for each candidate plan. progbar : bool, optional @@ -2948,11 +3048,11 @@ def run( explicitly supplied. It evaluates every tree scale using dynamic bond-pressure and tensor-work proxies. - ``order="quality"`` is a convenience mode matching the MPS layout - API: it enables bounded greedy refinement and opportunistic offline - refinement. If Nevergrad is unavailable, quality mode falls back to - greedy refinement. For ``objective="full_tree"``, quality mode uses - the dependency-free annealing stage. Pass + ``order="quality"`` is the high-quality mode matching the MPS layout + API: it upgrades the effective objective to ``"full_tree"``, enables + greedy leaf refinement, all-scale subtree topology refinement, and + hybrid topology annealing plus Nevergrad leaf search. When Nevergrad + is unavailable, it selects dependency-free simulated annealing. Pass ``search=None`` or ``refine=None`` explicitly to disable either stage. """ if order is _DEFAULT_ORDER: @@ -2960,18 +3060,26 @@ def run( else: order = _normalize_layout_order(order) if order == "quality": + if self.objective != "full_tree": + self.objective = "full_tree" + # Objective-dependent candidate caches may have been filled by + # an earlier fast/path run on this finder. Quality mode must + # not reuse those scores after upgrading to full-tree cost. + self._plan_cache.clear() + self._edge_load_cache.clear() + self._rank_diagnostics_cache.clear() + self._full_tree_profile_cache.clear() + self._full_tree_structure_cache.clear() + self._schmidt_rank_cache.clear() + self._similarity_cache.clear() + self._congestion_weights_cache = None + self._balanced_plan_cache = None if refine is _DEFAULT_SEARCH_OPTION: - refine = "greedy" + refine = self.refine or "greedy" if topology_refine is _DEFAULT_SEARCH_OPTION: - topology_refine = ( - "subtree" if self.objective == "full_tree" else "nni" - ) + topology_refine = self.topology_refine or "subtree" if search is _DEFAULT_SEARCH_OPTION: - search = ( - "anneal" - if self.objective == "full_tree" - else ("nevergrad" if _nevergrad_available() else None) - ) + search = self.search or _quality_search_mode() if chi is _DEFAULT_CHI: chi = self.chi else: @@ -3036,6 +3144,9 @@ def candidate_plans( include_quality=False, quality_refine_budget=None, quality_topology_budget=None, + quality_search=_DEFAULT_SEARCH_OPTION, + quality_search_budget=_DEFAULT_SEARCH_OPTION, + quality_seed=_DEFAULT_SEARCH_OPTION, ): """Return immutable candidate plans for optional pilot replay. @@ -3052,12 +3163,20 @@ def candidate_plans( include_quality : bool, optional Also add one ``"quality:arity=..."`` candidate per arity. These candidates start from the static objective candidates and apply - bounded greedy leaf and binary NNI topology refinement. This is - deliberately opt-in because it is more expensive than the static - candidate list. + bounded greedy leaf and topology refinement. For + ``objective="full_tree"`` this means all-scale subtree search and + the configured quality search; it is deliberately opt-in because + it is more expensive than the static candidate list. quality_refine_budget, quality_topology_budget : int, optional Bounds for the quality candidate's leaf-swap and NNI proposals. Each defaults to the normal bounded quality-mode budget. + quality_search : {None, "anneal", "nevergrad", "hybrid"}, optional + Optional second-stage search for quality candidates. The default + is no second stage for older objectives and ``"hybrid"`` (or + dependency-free ``"anneal"``) for ``objective="full_tree"``. + quality_search_budget, quality_seed : int, optional + Budget and seed for ``quality_search``. These are planning-only + controls; no tensor state is allocated or replayed here. """ if chi is _DEFAULT_CHI: chi = self.chi @@ -3071,12 +3190,35 @@ def candidate_plans( result = {} quality_settings = None if include_quality: + if quality_search is _DEFAULT_SEARCH_OPTION: + quality_search = ( + _quality_search_mode() + if self.objective == "full_tree" else None + ) + if quality_search_budget is _DEFAULT_SEARCH_OPTION: + quality_search_budget = self.search_budget + elif quality_search_budget is not None: + quality_search_budget = _validate_search_budget( + quality_search_budget, "quality_search_budget" + ) + if quality_seed is _DEFAULT_SEARCH_OPTION: + quality_seed = self.seed + else: + try: + quality_seed = int(quality_seed) + except (TypeError, ValueError) as exc: + raise ValueError("quality_seed must be an integer.") from exc + quality_topology_refine = ( + "subtree" if self.objective == "full_tree" else "nni" + ) quality_settings = self._resolve_search_settings( refine="greedy", refine_budget=quality_refine_budget, - topology_refine="nni", + topology_refine=quality_topology_refine, topology_budget=quality_topology_budget, - search=None, + search=quality_search, + search_budget=quality_search_budget, + seed=quality_seed, ) for arity in arities: plans = self._candidate_plans(arity) @@ -3118,6 +3260,146 @@ def candidate_plans( } return result + def targeted_candidates( + self, + plan, + edge_diagnostics, + *, + chi=_DEFAULT_CHI, + budget=32, + seed=0, + ): + """Propose static plans around replay-hot tree edges. + + This is the circuit-only feedback hook used by + :meth:`TreeOptimizer.optimize_layout`. ``edge_diagnostics`` is the + per-edge report produced by a short pilot replay. The method uses the + measured hot edges only to choose where to explore; every proposed + plan is still ranked by the configured static layout objective. It + never constructs a tensor network, applies a gate, or truncates a + state. + + The proposals include binary NNI moves where valid, local subtree + reconfigurations at the hot edge, and leaf exchanges across the hot + cut. Returned plans are immutable and deduplicated. + """ + if not isinstance(plan, TreePlan): + raise TypeError("plan must be a TreePlan.") + if chi is _DEFAULT_CHI: + chi = self.chi + else: + chi = _validate_chi(chi) + budget = _validate_search_budget(budget, "budget") + try: + seed = int(seed) + except (TypeError, ValueError) as exc: + raise ValueError("seed must be an integer.") from exc + if budget is None or budget < 1: + return [] + if not hasattr(edge_diagnostics, "items"): + raise TypeError("edge_diagnostics must be a mapping of edge metrics.") + + def hot_key(item): + edge, metrics = item + metrics = metrics if hasattr(metrics, "get") else {} + return ( + -float(metrics.get("discarded_fraction", 0.0) or 0.0), + -float(metrics.get("discarded_weight", 0.0) or 0.0), + -int(metrics.get("truncated", 0) or 0), + tuple(edge), + ) + + hot_edges = [] + for edge, metrics in sorted(edge_diagnostics.items(), key=hot_key): + try: + edge = tuple(int(x) for x in edge) + except (TypeError, ValueError): + continue + if len(edge) != 2: + continue + parent, child = edge + if plan.parent.get(child) != parent: + continue + hot_edges.append((edge, metrics)) + if not hot_edges: + return [] + + rng = np.random.default_rng(seed) + leaf_nodes = self._leaf_nodes(plan) + below = plan.subtree_qubit_masks() + proposals = [] + seen = set() + + def signature(candidate): + return ( + candidate.root, + candidate.root_qubit, + tuple( + sorted( + (node, tuple(children)) + for node, children in candidate.children.items() + ) + ), + tuple(sorted(candidate.qubit_of_leaf.items())), + ) + + original_signature = signature(plan) + + def add(candidate): + if candidate is None: + return + key = signature(candidate) + if key == original_signature or key in seen: + return + seen.add(key) + proposals.append(candidate) + + for (parent, child), _metrics in hot_edges: + if len(proposals) >= budget: + break + if (parent, child) in self._nni_edges(plan): + for variant in (0, 1): + if len(proposals) >= budget: + break + add(self._plan_with_nni(plan, parent, child, variant)) + + # Rebuild the smaller side first. A second draw at the parent + # lets the search change the attachment itself when the hot edge + # is a poor cut rather than merely a poor local ordering. + for node in (child, parent): + if len(proposals) >= budget: + break + if node not in plan.children: + continue + candidate = self._plan_with_subtree_reconfiguration( + plan, node, rng + ) + add(candidate) + + hot_qubits = [ + leaf for leaf in leaf_nodes + if below[child] & (1 << plan.qubit_of_leaf[leaf]) + ] + cold_qubits = [leaf for leaf in leaf_nodes if leaf not in hot_qubits] + if hot_qubits and cold_qubits: + rng.shuffle(hot_qubits) + rng.shuffle(cold_qubits) + for left_leaf, right_leaf in zip(hot_qubits, cold_qubits): + if len(proposals) >= budget: + break + add(self._plan_with_leaf_swap(plan, left_leaf, right_leaf)) + + # If the first cut generated fewer than the requested proposals, use + # deterministic neighbouring swaps so a pilot round still has a + # useful bounded exploration budget on shallow trees. + if len(proposals) < budget: + for left_leaf, right_leaf in zip(leaf_nodes, leaf_nodes[1:]): + if len(proposals) >= budget: + break + add(self._plan_with_leaf_swap(plan, left_leaf, right_leaf)) + proposals.sort(key=lambda candidate: self._selection_key(candidate, chi)) + return proposals[:budget] + def recommend_arities( self, max_arities=(2, 3, 4), @@ -3460,12 +3742,33 @@ def full_tree_profile(self, plan=None): return cached[1] below = plan.subtree_qubit_masks() - node_scales = _tree_node_scales(plan) - edges = tuple( - (parent, child) - for parent, children in plan.children.items() - for child in children + structure_key = ( + plan.root, + plan.root_qubit, + tuple( + sorted( + (node, tuple(children)) + for node, children in plan.children.items() + ) + ), + tuple(sorted(plan.qubit_of_node)), ) + cached_structure = self._full_tree_structure_cache.get(structure_key) + if cached_structure is None: + node_scales = _tree_node_scales(plan) + edges = tuple( + (parent, child) + for parent, children in plan.children.items() + for child in children + ) + self._full_tree_structure_cache[structure_key] = ( + node_scales, + edges, + ) + structure_reused = False + else: + node_scales, edges = cached_structure + structure_reused = True demand_log = {edge: 0.0 for edge in edges} bond_log = {edge: 0.0 for edge in edges} log_chi = ( @@ -3505,6 +3808,7 @@ def node_log_size(node): exact_events = 0 bounded_events = 0 bound_reasons = {} + support_span_cache = {} for payload, support, event_type, temporal_factor in zip( self.payloads, @@ -3521,9 +3825,12 @@ def node_log_size(node): } ): continue - support_mask, span_nodes, crossed_edges = self._support_span( - plan, support - ) + span_key = support + cached_span = support_span_cache.get(span_key) + if cached_span is None: + cached_span = self._support_span(plan, support) + support_span_cache[span_key] = cached_span + support_mask, span_nodes, crossed_edges = cached_span if not crossed_edges: continue event_count += 1 @@ -3611,6 +3918,10 @@ def node_log_size(node): "exact_events": int(exact_events), "bounded_events": int(bounded_events), "bound_reasons": bound_reasons, + "cache": { + "structure_reused": bool(structure_reused), + "unique_supports": len(support_span_cache), + }, "scales": scales, } self._full_tree_profile_cache[cache_key] = (plan, profile) @@ -4215,26 +4526,34 @@ def plot_tent( *, site_coords=None, ax=None, - figsize=(8, 7), + figsize=(8, 8), cmap="turbo", edge_cmap="GnBu", node_cmap="YlOrRd", - color_by="scale", - edge_color="#2f80a0", + color_by="order", + edge_color=None, show_edge_arrows=False, arrow_size=8.0, order=True, lattice=True, - show_gate_connectivity=True, + show_gate_connectivity=False, show_node_ids=False, show_site_labels=False, + show_leaf_nodes=False, + show_lattice_markers=True, + lattice_marker="+", + lattice_marker_size=100, + lattice_marker_color="#737e89", + lattice_marker_alpha=0.95, + lattice_skew=0.30, + lattice_rise=0.18, colorbar=False, show_axes=False, show_title=False, - node_size=38, + node_size=24, edge_linewidth=1.35, - edge_alpha=1.0, - vertical_spacing=None, + edge_alpha=0.8, + vertical_spacing=0.8, ): """Plot the hierarchy as a Cotengra-style tent over the raw graph. @@ -4252,6 +4571,15 @@ def plot_tent( hierarchy nodes by a deterministic post-order traversal, matching the ordering option in Cotengra's tent plots. Use ``color_by="order"`` if the same traversal should also control the colors. + Pass ``show_leaf_nodes=False`` when the physical lattice already has + its own site markers (for example gray ``+`` symbols) and only the + internal tree nodes should be drawn over that backdrop. + By default, supplied two-dimensional coordinates are projected into a + shallow tent base using ``x' = x + lattice_skew * y`` and + ``y' = lattice_rise * y``. Set ``lattice_skew=0`` and + ``lattice_rise=1`` to preserve the supplied coordinates. + The default order/turbo palette and matching hierarchy edges are + intended to give a compact Cotengra-style structural view. """ plt, colormaps, ScalarMappable, Normalize, _FancyArrowPatch = ( matplotlib_modules() @@ -4282,24 +4610,51 @@ def plot_tent( fig = ax.figure qubits = tuple(range(plan.n)) - coords = resolve_site_coords(qubits, site_coords) + raw_coords = resolve_site_coords(qubits, site_coords) + try: + lattice_skew = float(lattice_skew) + lattice_rise = float(lattice_rise) + except (TypeError, ValueError) as exc: + raise TypeError( + "lattice_skew and lattice_rise must be real numbers." + ) from exc + if not np.isfinite(lattice_skew) or not np.isfinite(lattice_rise): + raise ValueError("lattice_skew and lattice_rise must be finite.") + coords = { + qubit: ( + x + lattice_skew * y, + lattice_rise * y, + ) + for qubit, (x, y) in raw_coords.items() + } node_scales = _tree_node_scales(plan) n_scales = max(node_scales.values(), default=0) + 1 if lattice: - for left, right in coordinate_lattice_edges(coords): + for left, right in coordinate_lattice_edges(raw_coords): ax.plot( (coords[left][0], coords[right][0]), (coords[left][1], coords[right][1]), - color="#d5d9de", - linewidth=1.0, - alpha=0.78, + color="#b7c0c9", + linewidth=1.05, + alpha=0.82, zorder=1, ) + if show_lattice_markers: + ax.scatter( + [coords[qubit][0] for qubit in qubits], + [coords[qubit][1] for qubit in qubits], + marker=lattice_marker, + s=lattice_marker_size, + color=lattice_marker_color, + alpha=lattice_marker_alpha, + linewidths=1.25, + zorder=1.5, + ) if show_gate_connectivity: lattice_pairs = ( - coordinate_lattice_edge_keys(coords) + coordinate_lattice_edge_keys(raw_coords) if lattice else set() ) @@ -4453,6 +4808,8 @@ def hierarchy_edge_color(parent, child): ) for node in plan.nodes(): + if plan.is_leaf(node) and not show_leaf_nodes: + continue x, y = positions[node] ax.scatter( [x], diff --git a/src/pepsy/optimizers/tree/optimizer.py b/src/pepsy/optimizers/tree/optimizer.py index 86ee922..9fcf81b 100644 --- a/src/pepsy/optimizers/tree/optimizer.py +++ b/src/pepsy/optimizers/tree/optimizer.py @@ -1480,176 +1480,396 @@ def layout_report(self): } return self.layout_finder.report(self.plan) - def select_layout_for_compression( + def _layout_candidate_record(self, finder, plan, *, source=None): + """Build the common static record for a pilot candidate.""" + return { + "plan": plan, + "objective_key": finder._selection_key(plan, self.chi), + "path_score": finder.score(plan), + "tensor_cost": finder._tensor_cost_key(plan), + "edge_loads": finder.edge_loads(plan), + **({"source": source} if source is not None else {}), + } + + @staticmethod + def _pilot_edge_diagnostics(events): + """Aggregate replay losses by the immutable planned tree edge.""" + diagnostics = {} + for event in events: + edge = tuple(event["edge"]) + metric = diagnostics.setdefault(edge, { + "events": 0, + "truncated": 0, + "discarded_weight": 0.0, + "discarded_fraction": 0.0, + "tracked": False, + "max_before_bond": 0, + "max_after_bond": 0, + }) + metric["events"] += 1 + metric["truncated"] += int(bool(event.get("truncated", False))) + metric["max_before_bond"] = max( + metric["max_before_bond"], int(event.get("before_bond", 0)) + ) + metric["max_after_bond"] = max( + metric["max_after_bond"], int(event.get("after_bond", 0)) + ) + discarded_weight = event.get("discarded_weight") + discarded_fraction = event.get("discarded_fraction") + if discarded_weight is not None: + metric["tracked"] = True + metric["discarded_weight"] += float(discarded_weight) + if discarded_fraction is not None: + metric["tracked"] = True + metric["discarded_fraction"] = max( + metric["discarded_fraction"], float(discarded_fraction) + ) + return diagnostics + + def _pilot_layout_candidate( + self, plan, *, objective, pilot_steps=None, progbar=False + ): + """Replay one candidate and return state-aware diagnostics.""" + started = time.perf_counter() + trial = type(self)( + None, + n=self.n, + chi=self.chi, + cutoff=self.cutoff, + cutoff_mode=self.cutoff_mode, + mode=self.mode, + structure=self.structure, + max_arity=self.max_arity, + top_arity=self.top_arity, + community_frac=self.community_frac, + star_frac=self.star_frac, + layout_objective=objective, + tree=plan, + dtype=self.dtype, + threads=self.threads, + track_truncation=True, + track_infidelity=True, + max_intermediate_bond=self.max_intermediate_bond, + max_operator_qubits=self.max_operator_qubits, + max_subtree_nodes=self.max_subtree_nodes, + record_history=True, + run=False, + tn=self.tn, + ) + trial.G = list(self.G) + trial.where = list(self.where) + trial.event_types = list(self.event_types) + if pilot_steps is not None: + trial.G = trial.G[:pilot_steps] + trial.where = trial.where[:pilot_steps] + trial.event_types = trial.event_types[:pilot_steps] + try: + trial.run(progbar=progbar) + except Exception as exc: # pragma: no cover - backend-specific + return { + "status": "error", + "error": f"{type(exc).__name__}: {exc}", + "elapsed_seconds": float(time.perf_counter() - started), + "pilot_steps": len(trial.G), + } + + elapsed = float(time.perf_counter() - started) + events = deepcopy(trial.truncation_history) + edge_diagnostics = self._pilot_edge_diagnostics(events) + tracked = [ + event for event in events + if event.get("discarded_weight") is not None + ] + total_discarded = float( + sum(event["discarded_weight"] for event in tracked) + ) if tracked else 0.0 + max_fraction = max( + (float(event["discarded_fraction"]) for event in tracked), + default=0.0, + ) + update_runtime = float(sum( + update.get("elapsed_seconds", 0.0) + for update in trial.update_history + )) + return { + "status": "ok", + "elapsed_seconds": elapsed, + "update_runtime_seconds": update_runtime, + "infidelity": float(trial.infidelities[-1]), + "final_bond": int(trial.max_bond()), + "truncated_edges": int(sum( + event.get("truncated", False) for event in events + )), + "total_discarded_weight": total_discarded, + "max_discarded_fraction": float(max_fraction), + "pilot_steps": len(trial.G), + "edge_diagnostics": edge_diagnostics, + "updates": deepcopy(trial.update_history), + } + + def optimize_layout( self, *, + objective=None, pilot_candidates=4, + candidate_budget=None, pilot_steps=None, include_quality=True, + rounds=2, + topology_budget=None, + refine_budget=None, + search_budget=None, + seed=0, install=False, progbar=False, ): - """Select a tree layout using state-aware pilot replay. - - Static compression candidates are generated with - ``objective="compression"`` and then replayed on independent copies - of the current state. When ``include_quality=True`` (the default), one - bounded greedy/NNI quality candidate is reserved a pilot slot so it - cannot be excluded by static surrogate ranking. The pilot uses the - real tree update kernels, ``chi``, cutoff, backend, and queued gate - stream. The original state is unchanged. By default the selected plan - is returned for explicit hand-off; ``install=True`` is allowed only - for a product state and remounts that state exactly on the selected - geometry. Pass ``include_quality=False`` for the previous static-only - candidate set. + """Optimize a tree layout with bounded pilot-guided feedback. + + ``TreeLayoutFinder`` first generates static candidates, including the + all-scale ``objective="full_tree"`` search when requested. Each round + replays a short list on independent copies of the current *product* + state using the real tree kernels. The measured per-edge truncation + losses then seed targeted NNI, subtree, and cut-crossing leaf proposals + for the next round. The finder remains circuit-only: it does not + allocate tensors, replay gates, or perform truncations. + + The original optimizer is unchanged unless ``install=True`` is passed. + Installation is restricted to product states because an entangled TTN + cannot be relaid out exactly without an explicit state conversion. + ``objective="full_tree"`` is the recommended high-quality mode; its + static score covers routing demand, tensor width, work, and every tree + scale, while the pilot supplies the final state-aware choice. + ``candidate_budget`` is an alias for ``pilot_candidates`` for callers + that want to express the total candidate-pilot budget explicitly. """ try: + if candidate_budget is not None: + pilot_candidates = candidate_budget pilot_candidates = int(pilot_candidates) + rounds = int(rounds) except (TypeError, ValueError) as exc: - raise ValueError("pilot_candidates must be a positive integer.") from exc - if pilot_candidates < 1: - raise ValueError("pilot_candidates must be a positive integer.") + raise ValueError( + "pilot_candidates and rounds must be positive integers." + ) from exc + if pilot_candidates < 1 or rounds < 1: + raise ValueError("pilot_candidates and rounds must be positive integers.") if pilot_steps is not None: try: pilot_steps = int(pilot_steps) except (TypeError, ValueError) as exc: - raise ValueError("pilot_steps must be a positive integer or None.") from exc + raise ValueError( + "pilot_steps must be a positive integer or None." + ) from exc if pilot_steps < 1: raise ValueError("pilot_steps must be a positive integer or None.") - - finder = TreeLayoutFinder( - gates=self._layout_gate_stream(), - n=self.n, - structure=self.structure, - max_arity=self.max_arity, - community_frac=self.community_frac, - star_frac=self.star_frac, - objective="compression", - weight_mode=self.layout_weight_mode, - time_decay=self.layout_time_decay, - time_window=self.layout_time_window, - chi=self.chi, - max_operator_qubits=self.max_operator_qubits, - root_qubit=self.plan.root_qubit, - top_arity=self.top_arity, - ) - candidates = finder.candidate_plans( - chi=self.chi, - include_quality=bool(include_quality), - ) - ranked_static = sorted( - candidates, - key=lambda name: candidates[name]["objective_key"], - ) - if include_quality: - quality_names = [ - name for name in ranked_static if name.startswith("quality:") - ] - non_quality_names = [ - name for name in ranked_static if not name.startswith("quality:") - ] - reserved_quality = quality_names[:1] - ranked = ( - reserved_quality - + non_quality_names[: max(0, pilot_candidates - 1)] + if not _is_product_tensor_network(self.tn): + raise ValueError( + "Tree layout pilots require a product initial state when " + "comparing different tree geometries. Convert the entangled " + "state explicitly onto each candidate plan first." ) - else: - ranked = ranked_static[:pilot_candidates] - reports = {} - successful = [] - for name in ranked: - plan = candidates[name]["plan"] - if not _is_product_tensor_network(self.tn): - raise ValueError( - "Tree compression pilots require a product initial state " - "when comparing different tree geometries. Convert the " - "entangled state explicitly onto each candidate plan first." - ) - started = time.perf_counter() - trial = type(self)( - None, + try: + seed = int(seed) + except (TypeError, ValueError) as exc: + raise ValueError("seed must be an integer.") from exc + + objective = self.layout_objective if objective is None else objective + previous_plan = None + previous_edge_diagnostics = None + round_reports = [] + final_finder = None + final_candidates = None + final_ranked = None + final_selected_name = None + + for round_index in range(rounds): + finder = TreeLayoutFinder( + gates=self._layout_gate_stream(), n=self.n, - chi=self.chi, - cutoff=self.cutoff, - cutoff_mode=self.cutoff_mode, - mode=self.mode, structure=self.structure, max_arity=self.max_arity, - top_arity=self.top_arity, community_frac=self.community_frac, star_frac=self.star_frac, - tree=plan, - dtype=self.dtype, - threads=self.threads, - track_truncation=True, - track_infidelity=True, - max_intermediate_bond=self.max_intermediate_bond, + objective=objective, + weight_mode=self.layout_weight_mode, + time_decay=self.layout_time_decay, + time_window=self.layout_time_window, + chi=self.chi, max_operator_qubits=self.max_operator_qubits, - max_subtree_nodes=self.max_subtree_nodes, - record_history=self.record_history, - run=False, - tn=self.tn, + root_qubit=self.plan.root_qubit, + top_arity=self.top_arity, + seed=seed + round_index, ) - trial.G = list(self.G) - trial.where = list(self.where) - trial.event_types = list(self.event_types) - if pilot_steps is not None: - trial.G = trial.G[:pilot_steps] - trial.where = trial.where[:pilot_steps] - trial.event_types = trial.event_types[:pilot_steps] - try: - trial.run(progbar=progbar) - elapsed = time.perf_counter() - started - infidelity = float(trial.infidelities[-1]) - final_bond = int(trial.max_bond()) - truncated_edges = int(sum( - event.get("truncated", False) - for event in trial.truncation_history - )) - reports[name] = { - "status": "ok", - "elapsed_seconds": float(elapsed), - "infidelity": infidelity, - "final_bond": final_bond, - "truncated_edges": truncated_edges, - "pilot_steps": len(trial.G), - } - successful.append((infidelity, truncated_edges, final_bond, elapsed, name)) - except Exception as exc: # pragma: no cover - backend-specific - reports[name] = { - "status": "error", - "error": f"{type(exc).__name__}: {exc}", - "elapsed_seconds": float(time.perf_counter() - started), - "pilot_steps": len(trial.G), - } + quality_kwargs = { + "chi": self.chi, + "include_quality": bool(include_quality), + } + if topology_budget is not None: + quality_kwargs["quality_topology_budget"] = topology_budget + if refine_budget is not None: + quality_kwargs["quality_refine_budget"] = refine_budget + if search_budget is not None: + quality_kwargs["quality_search_budget"] = search_budget + quality_kwargs["quality_seed"] = seed + round_index + candidates = finder.candidate_plans(**quality_kwargs) + + if ( + previous_plan is not None + and previous_edge_diagnostics + and rounds > 1 + ): + targeted = finder.targeted_candidates( + previous_plan, + previous_edge_diagnostics, + chi=self.chi, + budget=max(2 * pilot_candidates, 8), + seed=seed + round_index, + ) + for proposal_index, plan in enumerate(targeted): + candidates[ + f"pilot:round={round_index}:proposal={proposal_index}" + ] = self._layout_candidate_record( + finder, + plan, + source="pilot_feedback", + ) - if not successful: - raise RuntimeError( - "All Tree compression layout pilot candidates failed. " - f"Diagnostics: {reports!r}" + ranked_static = sorted( + candidates, + key=lambda name: candidates[name]["objective_key"], + ) + quality_names = [ + name for name in ranked_static if name.startswith("quality:") + ] + feedback_names = [ + name for name in ranked_static + if name.startswith("pilot:") + ] + ordinary_names = [ + name for name in ranked_static + if not name.startswith(("quality:", "pilot:")) + ] + if include_quality: + ranked = quality_names[:1] + remaining = pilot_candidates - len(ranked) + ranked.extend(feedback_names[:remaining]) + remaining = pilot_candidates - len(ranked) + ranked.extend(ordinary_names[:remaining]) + remaining = pilot_candidates - len(ranked) + ranked.extend( + name for name in ranked_static + if name not in ranked + ) + ranked = ranked[:pilot_candidates] + else: + ranked = ranked_static[:pilot_candidates] + + reports = {} + successful = [] + for name in ranked: + report = self._pilot_layout_candidate( + candidates[name]["plan"], + objective=finder.objective, + pilot_steps=pilot_steps, + progbar=progbar, + ) + reports[name] = report + if report["status"] != "ok": + continue + successful.append(( + float(report["infidelity"]), + float(report["total_discarded_weight"]), + float(report["max_discarded_fraction"]), + int(report["truncated_edges"]), + float(report["elapsed_seconds"]), + int(report["final_bond"]), + name, + )) + if not successful: + raise RuntimeError( + "All Tree layout pilot candidates failed. " + f"Diagnostics: {reports!r}" + ) + selected_name = min(successful)[-1] + selected_plan = candidates[selected_name]["plan"] + selected_report = reports[selected_name] + round_reports.append({ + "round": round_index, + "objective": finder.objective, + "pilot_candidates": tuple(ranked), + "selected_candidate": selected_name, + "reports": reports, + }) + previous_plan = selected_plan + previous_edge_diagnostics = selected_report.get( + "edge_diagnostics", {} ) - selected_name = min(successful)[-1] - selected_plan = candidates[selected_name]["plan"] + final_finder = finder + final_candidates = candidates + final_ranked = ranked + final_selected_name = selected_name + + selected_plan = final_candidates[final_selected_name]["plan"] if install: self.plan = selected_plan self.tn = self._remount_product_state(self.tn) self.center = self.plan.root - self.layout_finder = finder - self.layout_objective = "compression" + self.layout_finder = final_finder + self.layout_objective = final_finder.objective + final_round = round_reports[-1] return { "plan": selected_plan, - "selected_candidate": selected_name, - "candidates": candidates, + "selected_candidate": final_selected_name, + "candidates": final_candidates, "pilot": { - "objective": "compression", + "objective": final_finder.objective, "include_quality": bool(include_quality), - "pilot_candidates": tuple(ranked), - "selected_candidate": selected_name, - "reports": reports, + "pilot_candidates": tuple(final_ranked), + "selected_candidate": final_selected_name, + "reports": final_round["reports"], + "rounds": round_reports, + "n_rounds": rounds, "installed": bool(install), }, } + def select_layout_for_compression( + self, + *, + pilot_candidates=4, + candidate_budget=None, + pilot_steps=None, + include_quality=True, + rounds=1, + topology_budget=None, + refine_budget=None, + search_budget=None, + seed=0, + install=False, + progbar=False, + ): + """Backward-compatible one-round compression layout selection. + + For iterative state-aware optimization, use + :meth:`optimize_layout`, for example with + ``objective="full_tree"`` and ``rounds=2``. This wrapper retains the + original compression objective and return shape. + """ + return self.optimize_layout( + objective="compression", + pilot_candidates=pilot_candidates, + candidate_budget=candidate_budget, + pilot_steps=pilot_steps, + include_quality=include_quality, + rounds=rounds, + topology_budget=topology_budget, + refine_budget=refine_budget, + search_budget=search_budget, + seed=seed, + install=install, + progbar=progbar, + ) + def plot_layout(self, plan=None, *, layout_kwargs=None, **plot_kwargs): """Plot the tree layout as a Cotengra-style tent. @@ -1816,6 +2036,7 @@ def _begin_update(self, kind, where): "kind": str(kind), "support": tuple(int(q) for q in where), "edge_start": len(self.truncation_history), + "started_at": time.perf_counter(), } return True @@ -1876,6 +2097,9 @@ def _finish_update(self): "update": len(self.update_history), "kind": active["kind"], "support": active["support"], + "elapsed_seconds": float( + time.perf_counter() - active["started_at"] + ), "edge_event_indices": list(range(start, len(self.truncation_history))), "edge_count": len(edge_events), "truncated_edges": sum(event["truncated"] for event in edge_events), diff --git a/src/pepsy/tensors/maps.py b/src/pepsy/tensors/maps.py index 9bac416..75095ec 100644 --- a/src/pepsy/tensors/maps.py +++ b/src/pepsy/tensors/maps.py @@ -2,6 +2,7 @@ from __future__ import annotations +from collections.abc import Mapping from numbers import Integral __all__ = ["OneDMap"] @@ -39,6 +40,7 @@ class OneDMap: "hilbert", "hilbert-row-major", "diag", + "finder", ) _MODE_ALIASES = { "snake": "snake", @@ -75,6 +77,8 @@ class OneDMap: "diag": "diag", "diagonal": "diag", "diag-snake": "diag", + "mps-finder": "finder", + "layout-finder": "finder", } @classmethod @@ -105,7 +109,22 @@ def _coalesce_dim_names( return Lx, Ly, Lz - def __init__(self, Lx=None, Ly=None, Lz=None, mode="snake", *, L_x=None, L_y=None, L_z=_ONE_D_MAP_UNSET): + def __init__( + self, + Lx=None, + Ly=None, + Lz=None, + mode="snake", + *, + L_x=None, + L_y=None, + L_z=_ONE_D_MAP_UNSET, + finder=None, + gate_stream=None, + gates=None, + layout_kwargs=None, + finder_base_mode="snake", + ): Lx, Ly, Lz = self._coalesce_dim_names( Lx=Lx, Ly=Ly, @@ -117,6 +136,24 @@ def __init__(self, Lx=None, Ly=None, Lz=None, mode="snake", *, L_x=None, L_y=Non self.Lx, self.Ly, self.Lz = self._normalize_dims(Lx, Ly, Lz=Lz) self.L_x, self.L_y, self.L_z = self.Lx, self.Ly, self.Lz self.mode = self._normalize_mode(mode) + if gate_stream is not None and gates is not None: + raise TypeError("pass only one of gate_stream= or gates=.") + if gates is not None: + gate_stream = gates + if finder is not None and gate_stream is not None: + raise TypeError( + "pass either finder= or gate_stream=/gates=, not both." + ) + if hasattr(gate_stream, "__next__"): + gate_stream = list(gate_stream) + self.finder = finder + self.gate_stream = gate_stream + self.layout_kwargs = ( + {} if layout_kwargs is None else dict(layout_kwargs) + ) + self.finder_base_mode = self._normalize_mode(finder_base_mode) + if self.finder_base_mode == "finder": + raise ValueError("finder_base_mode cannot itself be 'finder'.") def __repr__(self): shape = (self.Lx, self.Ly) if self.Lz is None else (self.Lx, self.Ly, self.Lz) @@ -343,21 +380,130 @@ def _coords_to_maps(coords): lattice_to_one_d = {coord: idx for idx, coord in one_d_to_lattice.items()} return one_d_to_lattice, lattice_to_one_d + @classmethod + def _coords_from_mps_finder( + cls, + Lx, + Ly, + Lz, + *, + finder=None, + gate_stream=None, + layout_kwargs=None, + finder_base_mode="snake", + ): + """Compose an MPS layout permutation with a regular lattice map. + + The MPS finder works on logical integer site labels and returns a + position-to-logical-site permutation. This method applies that + permutation to the coordinates of ``finder_base_mode``. It performs + no tensor-network construction or state replay. + """ + if finder is not None and gate_stream is not None: + raise TypeError( + "pass either finder= or gate_stream=/gates=, not both." + ) + if finder is None: + if gate_stream is None: + raise ValueError( + "mode='finder' requires gate_stream=, gates=, or finder=." + ) + from ..optimizers.mps.layout import MpsGateStreamLayoutFinder + + nsites = Lx * Ly if Lz is None else Lx * Ly * Lz + finder = MpsGateStreamLayoutFinder(gate_stream, L=nsites) + + layout_kwargs = {} if layout_kwargs is None else dict(layout_kwargs) + if isinstance(finder, Mapping): + plan = finder + else: + run = getattr(finder, "run", None) + if not callable(run): + raise TypeError( + "finder must be an MpsGateStreamLayoutFinder or a layout " + "plan mapping returned by its run() method." + ) + plan = run(**layout_kwargs) + + site_order = plan.get( + "site_order", + plan.get("qubit_inds", plan.get("order")), + ) + if site_order is None: + raise ValueError( + "MPS finder plan must contain site_order, qubit_inds, or order." + ) + try: + site_order = tuple(int(site) for site in site_order) + except (TypeError, ValueError) as exc: + raise ValueError( + "MPS finder site_order must contain integer lattice labels." + ) from exc + nsites = Lx * Ly if Lz is None else Lx * Ly * Lz + if len(site_order) != nsites or set(site_order) != set(range(nsites)): + raise ValueError( + "MPS finder site_order must be a permutation of the lattice " + f"labels 0..{nsites - 1}." + ) + + base_idx2coo, _ = cls.build( + Lx, + Ly, + Lz=Lz, + mode=finder_base_mode, + ) + return [base_idx2coo[site] for site in site_order] + @classmethod def _normalize_mode(cls, mode): mode_norm = str(mode).strip().lower().replace("_", "-") return cls._MODE_ALIASES.get(mode_norm, mode_norm) @_dualmethod - def build(target, Lx=None, Ly=None, Lz=_ONE_D_MAP_UNSET, mode=None, *, L_x=None, L_y=None, L_z=_ONE_D_MAP_UNSET): + def build( + target, + Lx=None, + Ly=None, + Lz=_ONE_D_MAP_UNSET, + mode=None, + *, + L_x=None, + L_y=None, + L_z=_ONE_D_MAP_UNSET, + finder=None, + gate_stream=None, + gates=None, + layout_kwargs=None, + finder_base_mode=None, + ): """Build ``(one_d_to_lattice, lattice_to_one_d)`` for a traversal mode. This can be called either as ``OneDMap.build(Lx, Ly, ...)`` or on an instance, e.g. ``OneDMap(Lx, Ly, mode="row-major").build()``. Instance calls can override options per use, for example ``mapper.build(mode="snake")``. + + ``mode="finder"`` composes an MPS gate-stream layout with the base + lattice traversal. Supply ``gate_stream=``/``gates=`` or a previously + constructed MPS layout ``finder=``. The finder only analyzes supports + and returns a site permutation; it never allocates or truncates an MPS. """ cls = target if isinstance(target, type) else type(target) + if gate_stream is not None and gates is not None: + raise TypeError("pass only one of gate_stream= or gates=.") + if gates is not None: + gate_stream = gates + if not isinstance(target, type): + if finder is None: + finder = target.finder + if gate_stream is None: + gate_stream = target.gate_stream + if layout_kwargs is None: + layout_kwargs = target.layout_kwargs + if finder_base_mode is None: + finder_base_mode = target.finder_base_mode + if finder_base_mode is None: + finder_base_mode = "snake" Lx, Ly, Lz, mode_norm = cls._resolve_call_params( target, Lx=Lx, @@ -369,6 +515,18 @@ def build(target, Lx=None, Ly=None, Lz=_ONE_D_MAP_UNSET, mode=None, *, L_x=None, L_z=L_z, ) + if mode_norm == "finder": + coords = cls._coords_from_mps_finder( + Lx, + Ly, + Lz, + finder=finder, + gate_stream=gate_stream, + layout_kwargs=layout_kwargs, + finder_base_mode=finder_base_mode, + ) + return cls._coords_to_maps(coords) + if mode_norm == "snake": coords = ( cls._coords_snake_2d(Lx, Ly, major="col") @@ -589,6 +747,11 @@ def show( L_x=None, L_y=None, L_z=_ONE_D_MAP_UNSET, + finder=None, + gate_stream=None, + gates=None, + layout_kwargs=None, + finder_base_mode=None, print_out=False, ax=None, title=None, @@ -614,7 +777,26 @@ def show( L_y=L_y, L_z=L_z, ) - one_d_to_lattice, lattice_to_one_d = cls.build(Lx, Ly, Lz=Lz, mode=mode_norm) + if not isinstance(target, type): + if finder is None: + finder = target.finder + if gate_stream is None: + gate_stream = target.gate_stream + if layout_kwargs is None: + layout_kwargs = target.layout_kwargs + if finder_base_mode is None: + finder_base_mode = target.finder_base_mode + one_d_to_lattice, lattice_to_one_d = cls.build( + Lx, + Ly, + Lz=Lz, + mode=mode_norm, + finder=finder, + gate_stream=gate_stream, + gates=gates, + layout_kwargs=layout_kwargs, + finder_base_mode=finder_base_mode, + ) if Lz is not None: raise NotImplementedError( "OneDMap.show() is currently only available for 2D lattices." diff --git a/tests/test_ham.py b/tests/test_ham.py index 2880947..0bbcdf7 100644 --- a/tests/test_ham.py +++ b/tests/test_ham.py @@ -619,6 +619,46 @@ def test_map_builder_instance_style_can_override_mode_per_call(): assert map_[2] == (1, 1) +def test_map_builder_finder_composes_mps_site_order_with_lattice_coords(): + """Finder mode should map optimized MPS positions back to coordinates.""" + base_idx2coo, _ = OneDMap(2, 2, mode="snake").build() + plan = {"site_order": (2, 0, 3, 1)} + mapper = OneDMap(2, 2, mode="finder", finder=plan) + + idx2coo, coo2idx = mapper.build() + + assert idx2coo == { + position: base_idx2coo[logical_site] + for position, logical_site in enumerate(plan["site_order"]) + } + assert coo2idx == {coord: position for position, coord in idx2coo.items()} + + +def test_map_builder_finder_mode_runs_mps_layout_finder(): + """Finder mode should accept a gate stream without touching an MPS state.""" + gates = [ + (quimb.CNOT(), (0, 3)), + (quimb.CNOT(), (3, 1)), + ] + mapper = OneDMap( + 2, + 2, + mode="finder", + gates=gates, + layout_kwargs={"order": "input"}, + ) + + idx2coo, _ = mapper.build() + + assert idx2coo == OneDMap(2, 2, mode="snake").build()[0] + + +def test_map_builder_finder_mode_requires_layout_source(): + """Finder mode should fail clearly when no stream or plan is supplied.""" + with pytest.raises(ValueError, match="mode='finder'"): + OneDMap(2, 2, mode="finder").build() + + def test_map_builder_instance_style_supports_3d_build(): """Instance-style build() should preserve the 3D mapping modes.""" mapper = OneDMap(2, 2, Lz=2, mode="col-major") diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index 2612ed8..c42829c 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -1058,6 +1058,63 @@ def test_tree_compression_layout_pilot_is_non_mutating(): assert opt.max_bond() == 1 +def test_tree_layout_pilot_feedback_is_iterative_and_edge_aware(): + """Full-tree pilots feed bounded hot-edge proposals into the next round.""" + h = np.array([[1.0, 1.0], [1.0, -1.0]], dtype=complex) / np.sqrt(2.0) + gates = [(h, 0), (pepsy.cnot(), (0, 3)), (pepsy.cnot(), (3, 1))] + opt = TreeOptimizer(gates, n=4, max_arity=2, chi=1, run=False) + original_plan = opt.plan + + selected = opt.optimize_layout( + objective="full_tree", + pilot_candidates=1, + pilot_steps=3, + rounds=2, + topology_budget=1, + refine_budget=1, + search_budget=2, + ) + + assert selected["pilot"]["objective"] == "full_tree" + assert selected["pilot"]["n_rounds"] == 2 + assert len(selected["pilot"]["rounds"]) == 2 + assert opt.plan is original_plan + report = selected["pilot"]["reports"][selected["selected_candidate"]] + assert report["status"] == "ok" + assert report["update_runtime_seconds"] >= 0.0 + assert isinstance(report["edge_diagnostics"], dict) + assert opt.max_bond() == 1 + + +def test_tree_layout_targeted_candidates_are_static_and_bounded(): + """Hot-edge proposal generation never allocates or mutates a TTN.""" + finder = TreeLayoutFinder( + [(pepsy.cnot(), (0, 3)), (pepsy.cnot(), (3, 1))], + n=4, + max_arity=2, + objective="full_tree", + chi=None, + ) + plan = TreePlan.from_order(range(4), structure="balanced", max_arity=2) + edge = next( + (parent, child) + for parent, children in plan.children.items() + for child in children + ) + proposals = finder.targeted_candidates( + plan, + {edge: {"truncated": 1, "discarded_fraction": 0.5}}, + budget=3, + seed=3, + ) + + assert len(proposals) <= 3 + assert all(candidate.is_binary() for candidate in proposals) + unchanged = TreePlan.from_order(range(4), structure="balanced", max_arity=2) + assert plan.children == unchanged.children + assert plan.qubit_of_leaf == unchanged.qubit_of_leaf + + def test_tree_candidate_plans_include_quality_for_state_aware_pilots(): """Quality refinement is exposed as an explicit pilot candidate.""" finder = TreeLayoutFinder( @@ -1149,6 +1206,35 @@ def test_full_tree_anneals_subtrees_without_changing_binary_contract(): assert candidate["full_tree_profile"]["scales"] +def test_full_tree_hybrid_quality_search_is_static_and_budgeted(): + """Full-tree quality combines topology and leaf search without a TTN.""" + pytest.importorskip("nevergrad") + finder = TreeLayoutFinder( + [(pepsy.cnot(), (0, 3)), (pepsy.cnot(), (3, 1)), + (pepsy.cnot(), (2, 5)), (pepsy.cnot(), (4, 5))], + n=6, + max_arity=(2,), + top_arity=3, + objective="full_tree", + seed=7, + ) + plan = finder.run( + order="quality", + topology_budget=2, + refine_budget=2, + search_budget=6, + ) + + planning = finder._last_arity_recommendation["candidates"][0]["planning"] + search = planning["search"] + assert finder.chi is None + assert plan.is_binary() + assert search["method"] == "hybrid" + assert search["anneal"]["search"] == "anneal" + assert search["nevergrad"]["method"] == "nevergrad" + assert search["evaluations"] <= 6 + + def test_tree_edge_loads_match_full_edge_reference(): """Steiner-only edge scanning preserves the full congestion calculation.""" rng = np.random.default_rng(109) @@ -1686,7 +1772,9 @@ def test_tree_layout_finder_plot_defaults_to_tent(): assert len(fig.axes) == 1 assert not ax.axison # schematic-style presentation by default assert not ax.texts - assert len(ax.collections) == len(plan.nodes()) + assert len(ax.collections) == 1 + sum( + not plan.is_leaf(node) for node in plan.nodes() + ) plt.close(fig) @@ -1739,7 +1827,35 @@ def test_tree_layout_finder_plot_tent_draws_hierarchy_over_raw_graph(): assert not ax.texts assert not ax.axison assert len(ax.lines) >= len(plan.nodes()) - 1 - assert len(ax.collections) == len(plan.nodes()) + assert len(ax.collections) == 1 + sum( + not plan.is_leaf(node) for node in plan.nodes() + ) + plt.close(fig) + + +def test_tree_layout_tent_can_hide_physical_leaf_nodes(): + """Physical plus-mark backdrops can replace tree leaf circles.""" + matplotlib = pytest.importorskip("matplotlib") + matplotlib.use("Agg", force=True) + import matplotlib.pyplot as plt + + finder = TreeLayoutFinder( + [(pepsy.cnot(), (0, 3)), (pepsy.cnot(), (3, 1))], + n=4, + max_arity=2, + ) + plan = finder.run() + fig, ax = finder.plot_tent( + plan, + lattice=False, + show_gate_connectivity=False, + show_leaf_nodes=False, + ) + + assert len(ax.collections) == sum( + not plan.is_leaf(node) for node in plan.nodes() + ) + assert len(ax.lines) == len(plan.nodes()) - 1 plt.close(fig) @@ -1769,17 +1885,9 @@ def test_tree_layout_scale_colors_do_not_depend_on_gate_stream_length(): edge_color=None, show_edge_arrows=False, ) - # With the default one-dimensional coordinates, the first lines are - # the lattice and only the non-lattice gate-connectivity background; - # nearest-neighbor gates are already represented by the lattice. - lattice_pairs = { - frozenset((site, site + 1)) - for site in range(plan.n - 1) - } - nonlattice_gates = sum( - frozenset(where) not in lattice_pairs for _, where in gates - ) - background_lines = len(plan.leaves()) - 1 + nonlattice_gates + # Gate-connectivity overlays are disabled by default, so only the + # one-dimensional physical lattice precedes the hierarchy edges. + background_lines = len(plan.leaves()) - 1 structural_colors.append( tuple( tuple(line.get_color()) @@ -1793,7 +1901,9 @@ def test_tree_layout_scale_colors_do_not_depend_on_gate_stream_length(): ) ) assert len(fig.axes) == 1 - assert len(ax.collections) == len(plan.nodes()) + assert len(ax.collections) == 1 + sum( + not plan.is_leaf(node) for node in plan.nodes() + ) plt.close(fig) assert structural_colors[0] == structural_colors[1] @@ -1802,8 +1912,8 @@ def test_tree_layout_scale_colors_do_not_depend_on_gate_stream_length(): assert len(set(scale_node_colors[0])) > 1 -def test_tree_layout_tent_edges_are_uniform_by_default(): - """Tent hierarchy edges use one solid color unless explicitly varied.""" +def test_tree_layout_tent_edges_match_order_colors_by_default(): + """Tent hierarchy edges follow the default order color palette.""" matplotlib = pytest.importorskip("matplotlib") matplotlib.use("Agg", force=True) import matplotlib.pyplot as plt @@ -1813,16 +1923,11 @@ def test_tree_layout_tent_edges_are_uniform_by_default(): plan = finder.run() fig, ax = finder.plot_tent(plan, color_by="scale") - lattice_pairs = { - frozenset((site, site + 1)) for site in range(plan.n - 1) - } - background_lines = len(plan.leaves()) - 1 + sum( - frozenset(where) not in lattice_pairs for _, where in gates - ) + background_lines = len(plan.leaves()) - 1 hierarchy_colors = { line.get_color() for line in ax.lines[background_lines:] } - assert hierarchy_colors == {"#2f80a0"} + assert len(hierarchy_colors) > 1 assert not ax.patches plt.close(fig) @@ -1846,24 +1951,20 @@ def test_tree_layout_tent_colored_edges_match_child_nodes(): show_edge_arrows=False, ) - lattice_pairs = { - frozenset((site, site + 1)) for site in range(plan.n - 1) - } - background_lines = plan.n - 1 + sum( - frozenset(where) not in lattice_pairs - for _, where in [(pepsy.cnot(), (0, 3)), (pepsy.cnot(), (1, 2))] - ) + background_lines = plan.n - 1 + internal_nodes = [node for node in plan.nodes() if not plan.is_leaf(node)] node_colors = { node: tuple(collection.get_facecolors()[0]) - for node, collection in zip(plan.nodes(), ax.collections) + for node, collection in zip(internal_nodes, ax.collections[1:]) } hierarchy_lines = ax.lines[background_lines:] line_index = 0 for parent, children in plan.children.items(): for child in children: - assert tuple(hierarchy_lines[line_index].get_color()) == pytest.approx( - node_colors[child] - ) + if not plan.is_leaf(child): + assert tuple( + hierarchy_lines[line_index].get_color() + ) == pytest.approx(node_colors[child]) line_index += 1 assert line_index == len(hierarchy_lines) plt.close(fig) @@ -1925,9 +2026,34 @@ def fake_improve(plan, *, chi, settings, progbar=False): plan = finder.run() assert plan.n == 4 + assert finder.objective == "full_tree" assert captured["refine"] == "greedy" - assert captured["topology_refine"] == "nni" - assert captured["search"] is None + assert captured["topology_refine"] == "subtree" + assert captured["search"] == "anneal" + assert captured["search_budget"] == finder.search_budget + + +def test_tree_layout_quality_run_upgrades_a_fast_finder(monkeypatch): + """The explicit quality run is the full-tree mode even after construction.""" + monkeypatch.setitem(sys.modules, "nevergrad", None) + finder = TreeLayoutFinder( + [(pepsy.cnot(), (0, 3)), (pepsy.cnot(), (1, 2))], + n=4, + max_arity=2, + ) + captured = {} + + def fake_improve(plan, *, chi, settings, progbar=False): + captured.update(settings) + return plan, {"method": "test"} + + monkeypatch.setattr(finder, "_improve_plan", fake_improve) + plan = finder.run(order="quality") + + assert plan.n == 4 + assert finder.objective == "full_tree" + assert captured["topology_refine"] == "subtree" + assert captured["search"] == "anneal" def test_tree_layout_nni_refinement_changes_binary_topology(): From aaf6363576d98c23d66456bcdb22a1975b849504 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Fri, 31 Jul 2026 12:23:11 -0600 Subject: [PATCH 41/70] Add bounded open-loop BP observables --- .github/skills/belief-propagation/SKILL.md | 29 + docs/api/bp.md | 50 + .../references/belief_propagation.md | 4 +- src/pepsy/bp/__init__.py | 10 + src/pepsy/bp/series.py | 973 ++++++++++++++++-- tests/test_bp_open_series.py | 180 +++- 6 files changed, 1165 insertions(+), 81 deletions(-) diff --git a/.github/skills/belief-propagation/SKILL.md b/.github/skills/belief-propagation/SKILL.md index 5e3626a..7dfe59a 100644 --- a/.github/skills/belief-propagation/SKILL.md +++ b/.github/skills/belief-propagation/SKILL.md @@ -56,6 +56,35 @@ over `quimb.tensor.belief_propagation`; the annotated paper trail lives in `compute_local_expectation_loop_cluster` accept Quimb-style ``{site_or_sites: operator}`` mappings and return scalar expectations. Share one D2BP solve across terms. Use the scalar APIs for fermionic observables. +- Long-range open-loop observables: + `partial_trace_open_loop_series_expand` retains open support-connecting + paths, closed loops, and path-plus-loop configurations in a rho; + `compute_local_expectation_open_loop_series` inserts the gate directly into + the graded ket/bra contraction. Use `diagnose_open_loop_series` first when + the workflow must preflight geometry, Cotengra FLOP/peak-memory costs, and + reusable contraction plans. `diagnose_open_rho_series` adds rho output shape + and output-memory estimates. `OpenLoopSeriesCache` caches geometry and + `OpenLoopSeriesDiagnosticCache` caches topology-keyed diagnostics/plans; + measurements may reuse them or enumerate lazily on demand. +- Open-loop route and budgets: use `mode="auto"` for distance awareness. It + keeps nearby supports on the exact route and switches sufficiently distant + dense/tree supports to a bounded weighted-shortest-path corridor. Set + `auto_corridor_distance` to change the threshold. `max_terms` bounds all + explicit configurations, while `max_loop_terms` separately bounds only + closed-loop and path-plus-loop corrections; `max_enumeration_time` and + `max_enumeration_memory` bound geometry discovery. For contraction costs, + use `max_flops_log10` and `max_peak_memory_log2`, then choose + `on_budget="raise"` for production. `return_result=True` returns an + auditable `OpenLoopMeasurementResult`; `measure_resources=True` records + observed Python/host/GPU resources. `adaptive_open_loop_series` compares a + corridor or cluster ladder, but is a numerical stabilization check rather + than a rigorous truncation bound. +- Native cyclic fermion route: detect the cyclic native Symmray case before + explicit edge enumeration and use the graded cluster-compatible route with + `cluster_size`. Do not force `edge_cutoff`, `corridor_width`, or a dense + `trace(rho @ gate)` sign oracle onto that route. For a 2-periodic PBC axis, + treat the virtual lattice as a multigraph so parallel seam bonds remain + distinct loop-series edges. - Explicit-edge local D2BP observables: `partial_trace_edge_loop_series_expand` and `compute_local_expectation_edge_loop_series` use the same canonical diff --git a/docs/api/bp.md b/docs/api/bp.md index 8fb4e56..a965bb8 100644 --- a/docs/api/bp.md +++ b/docs/api/bp.md @@ -186,6 +186,10 @@ limits raise `OpenLoopEnumerationLimitError` rather than returning a partial sum. `gloops` remains a compatibility alias, but new code should use the route-specific names. +Use `max_loop_terms` for a separate budget on closed-loop and path-plus-loop +corrections. This lets nearby supports retain a richer loop tail while keeping +long-range measurements focused on their shortest connecting paths. + For very distant supports, set `corridor_width` to use the bounded corridor route. It retains a small weighted-shortest-path beam, inflates those paths by the requested graph width, and adds connected loop decorations only near @@ -256,6 +260,52 @@ share numerical results. If no `diagnostic` is supplied, scalar measurement with `mode="auto"` performs this diagnostic pass internally before starting the numerical contractions. +For production measurements, use `on_budget="raise"` and request the +auditable result record: + +```python +from pepsy.bp import OpenLoopMeasurementResult + +result = compute_local_expectation_open_loop_series( + peps.tn, + {support: gate}, + mode="auto", + on_budget="raise", + measure_resources=True, + return_result=True, +) +assert isinstance(result, OpenLoopMeasurementResult) +assert result.complete and result.bp_converged +``` + +The default `on_budget="report"` preserves historical partial-sum behavior +but records `complete=False` and omitted terms in `info`. `on_budget="skip"` +is available for exploratory workflows. Cotengra estimates remain preflight +guards; `measure_resources=True` additionally records observed Python and +host-RSS high-water marks. The same flags are available on +`partial_trace_open_loop_series_expand`; its result record stores the rho in +`result.value` and its trace in `result.normalization`. + +For controlled convergence, use an adaptive corridor or cluster ladder: + +```python +from pepsy.bp import adaptive_open_loop_series + +ladder = adaptive_open_loop_series( + peps.tn, + {support: gate}, + corridor_widths=(0, 1, 2, 4), + on_budget="raise", +) +value = ladder.value +``` + +The ladder tests numerical stabilization, not a rigorous truncation bound. +For cyclic native fermions, pass `cluster_sizes=(...)` instead. The +`diagnose_open_rho_series` helper adds physical output shape and output-memory +estimates. Rectangular PBC corridor discovery treats the virtual graph as a +multigraph, so period-two seam bonds remain distinct paths and loop edges. + This path performs an explicit configuration sum and normalizes only after the sum; it does not apply the scalar disconnected-loop resummation used by `partial_trace_edge_loop_series_expand`. diff --git a/docs/development/references/belief_propagation.md b/docs/development/references/belief_propagation.md index 3178d58..f48ed41 100644 --- a/docs/development/references/belief_propagation.md +++ b/docs/development/references/belief_propagation.md @@ -23,8 +23,8 @@ side · **[roots]** foundational / prior art. | `loop_series_expand`, `LoopSeriesTerm`, `LoopSeriesCache` | edge-resolved `P + Q` loop series for D1BP and D2BP; retains excited-bond degree and distinct embeddings/chord subsets | Evenbly et al. 2409.03108 | | `partial_trace_loop_series_expand`, `compute_local_expectation_loop_series` | D2BP local reduced-density-matrix and scalar `P + Q` loop series; keeps physical output legs open and uses native Symmray virtual projectors | Evenbly et al. 2409.03108; quimb local loop-series API | | `partial_trace_edge_loop_series_expand`, `compute_local_expectation_edge_loop_series` | D2BP local RDM and graded scalar observable expansion over canonical explicit Q-edge terms; does not reinterpret Quimb's local-region cutoff | Evenbly et al. 2409.03108; Pepsy API | -| `partial_trace_open_loop_series_expand`, `partial_trace_open_loop_series_sweep` | Lazy/path-first D2BP rho configuration sum over open Q paths, closed loops, and attached or disconnected path-plus-loop terms; `edge_cutoff` counts excited Q edges, bounded discovery uses `max_terms` / enumeration time / approximate geometry memory limits, and regional contraction paths are reused; `corridor_width` activates bounded weighted-shortest-path discovery, sampled connected loop decorations, and optional compressed boundary contraction; cyclic native fermionic graphs select the graded cluster route before edge discovery and use `cluster_size` | Evenbly et al. 2409.03108; Pepsy API | -| `compute_local_expectation_open_loop_series`, `diagnose_open_loop_series` | Scalar companion for long-range gates plus a non-contracting diagnostic pass: accepts native Fermion operators, reuses open-loop geometry, selects exact/corridor/graded-cluster routes, reports path/loop families and Cotengra FLOP/peak-memory estimates, and can reuse cached diagnostics before inserting the gate through the graded open-bond projector route | Evenbly et al. 2409.03108; Pepsy API | +| `partial_trace_open_loop_series_expand`, `partial_trace_open_loop_series_sweep` | Lazy/path-first D2BP rho configuration sum over open Q paths, closed loops, and attached or disconnected path-plus-loop terms; `edge_cutoff` counts excited Q edges, `max_terms` bounds all configurations while `max_loop_terms` separately bounds loop corrections, and regional contraction paths are reused; rectangular PBC corridors preserve parallel seam bonds; `corridor_width` activates bounded weighted-shortest-path discovery, sampled connected loop decorations, and optional compressed boundary contraction; cyclic native fermionic graphs select the graded cluster route before edge discovery and use `cluster_size` | Evenbly et al. 2409.03108; Pepsy API | +| `compute_local_expectation_open_loop_series`, `diagnose_open_loop_series`, `diagnose_open_rho_series`, `adaptive_open_loop_series` | Scalar/rho companions for long-range gates plus non-contracting diagnostics and adaptive convergence ladders: accepts native Fermion operators, reuses open-loop geometry and compatible contraction plans, selects exact/corridor/graded-cluster routes, reports path/loop families, Cotengra FLOP/peak-memory estimates, output shape/memory, completeness, and optional observed resources | Evenbly et al. 2409.03108; Pepsy API | | `partial_trace_loop_cluster_expand`, `compute_local_expectation_loop_cluster` | D2BP local reduced-density-matrix and scalar generalized-loop cluster expansion; combines BP-closed regions with inclusion--exclusion counts | Gray et al. 2510.05647; quimb local cluster API | | `loop_expand` | explicit selector between the edge loop series and region loop-cluster expansion; preserves each method's cutoff and result metadata | Pepsy API | | `partitioned_expand`, `pne_expand`, `PNEExpansionResult` | linear and combinatorial partitioned network expansions for D1BP/D2BP, with optional residue, explicit projectors, open outputs, and fixed recursive schedules | Evenbly, Gray & Chan 2512.10910 | diff --git a/src/pepsy/bp/__init__.py b/src/pepsy/bp/__init__.py index 5532d93..66cd236 100644 --- a/src/pepsy/bp/__init__.py +++ b/src/pepsy/bp/__init__.py @@ -45,7 +45,10 @@ select_bp_candidate, ) from .series import ( + OpenLoopBudgetError, + OpenLoopAdaptiveResult, OpenLoopEnumerationLimitError, + OpenLoopMeasurementResult, OpenLoopObservableTerm, OpenLoopSeriesDiagnostic, OpenLoopSeriesDiagnosticCache, @@ -62,6 +65,8 @@ partial_trace_open_loop_series_expand, partial_trace_open_loop_series_sweep, diagnose_open_loop_series, + diagnose_open_rho_series, + adaptive_open_loop_series, partial_trace_loop_cluster_expand, partial_trace_loop_series_expand, loop_series_expand, @@ -150,12 +155,17 @@ "OpenLoopObservableTerm", "OpenLoopSeriesDiagnostic", "OpenLoopSeriesDiagnosticCache", + "OpenLoopBudgetError", + "OpenLoopAdaptiveResult", + "OpenLoopMeasurementResult", "OpenLoopSeriesSweepResult", "LoopSeriesResult", "LoopSeriesTerm", "compute_local_expectation_edge_loop_series", "compute_local_expectation_open_loop_series", "diagnose_open_loop_series", + "diagnose_open_rho_series", + "adaptive_open_loop_series", "compute_local_expectation_loop_cluster", "compute_local_expectation_loop_series", "partial_trace_edge_loop_series_expand", diff --git a/src/pepsy/bp/series.py b/src/pepsy/bp/series.py index d0c43ae..ff86921 100644 --- a/src/pepsy/bp/series.py +++ b/src/pepsy/bp/series.py @@ -30,6 +30,12 @@ import sys import time from typing import Any, ClassVar +import warnings + +try: + import resource +except ImportError: # pragma: no cover - Windows does not expose rusage. + resource = None import autoray as ar import numpy as np @@ -52,7 +58,10 @@ ) __all__ = [ + "OpenLoopBudgetError", "OpenLoopEnumerationLimitError", + "OpenLoopMeasurementResult", + "OpenLoopAdaptiveResult", "OpenLoopObservableTerm", "OpenLoopSeriesDiagnostic", "OpenLoopSeriesDiagnosticCache", @@ -64,6 +73,8 @@ "compute_local_expectation_edge_loop_series", "compute_local_expectation_open_loop_series", "diagnose_open_loop_series", + "diagnose_open_rho_series", + "adaptive_open_loop_series", "compute_local_expectation_loop_cluster", "partial_trace_loop_cluster_expand", "partial_trace_edge_loop_series_expand", @@ -94,6 +105,156 @@ def __init__(self, reason: str, limit: float, observed: float): ) +class OpenLoopBudgetError(RuntimeError): + """Raised when a requested open-series contraction exceeds its budget.""" + + +@dataclass +class OpenLoopMeasurementResult: + """Auditable result returned by the opt-in production API. + + The legacy measurement functions continue to return a scalar or rho by + default. ``return_result=True`` returns this record instead, keeping the + numerical value together with completeness, route, budget, BP, and + resource information. + """ + + value: Any + normalization: Any = None + info: dict[str, Any] = field(default_factory=dict) + diagnostic: "OpenLoopSeriesDiagnostic | None" = None + complete: bool = True + approximate: bool = False + route: str | None = None + omitted_terms: tuple[Any, ...] = () + resources: dict[str, Any] = field(default_factory=dict) + bp_converged: bool | None = None + bp_iterations: int | None = None + bp_max_mdiff: float | None = None + bp: Any = field(default=None, repr=False) + + +@dataclass +class OpenLoopAdaptiveResult: + """Convergence ladder for corridor or native-cluster measurements.""" + + values: tuple[Any, ...] + settings: tuple[dict[str, Any], ...] + differences: tuple[float | None, ...] + converged: bool + selected_index: int | None + diagnostics: tuple[OpenLoopSeriesDiagnostic | None, ...] = () + infos: tuple[dict[str, Any], ...] = () + bp: Any = field(default=None, repr=False) + + @property + def value(self): + """Return the last stable value, or the final value if unsettled.""" + if not self.values: + return None + index = self.selected_index + return self.values[-1 if index is None else index] + + +class _OpenLoopResourceMonitor: + """Small optional host-resource monitor for an open-series call. + + ``tracemalloc`` measures Python allocations, while ``ru_maxrss`` gives a + process-level high-water mark including NumPy allocations on Unix. They + are deliberately reported as observations, not exact tensor liveness + proofs. Cotengra's symbolic peak remains the preflight guard. + """ + + def __init__(self, enabled): + self.enabled = bool(enabled) + self.started = time.perf_counter() + self._tracemalloc = None + self._owns_tracemalloc = False + if not self.enabled: + return + import tracemalloc + + self._tracemalloc = tracemalloc + if not tracemalloc.is_tracing(): + tracemalloc.start() + self._owns_tracemalloc = True + self._start = tracemalloc.get_traced_memory() + self._rss_start = self._rss() + self._gpu_start = self._gpu_snapshot() + + @staticmethod + def _rss(): + if resource is None: + return None + try: + value = int(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss) + # Linux and the BSDs report KiB; macOS reports bytes. + return value if sys.platform == "darwin" else value * 1024 + except (AttributeError, OSError): + return None + + @staticmethod + def _gpu_snapshot(): + snapshots = [] + try: + import torch + + if torch.cuda.is_available(): + for device in range(torch.cuda.device_count()): + snapshots.append( + { + "backend": "torch", + "device": device, + "current_bytes": int( + torch.cuda.memory_allocated(device) + ), + "peak_bytes": int( + torch.cuda.max_memory_allocated(device) + ), + } + ) + except (ImportError, RuntimeError): + pass + try: + import cupy + + pool = cupy.get_default_memory_pool() + snapshots.append( + { + "backend": "cupy", + "device": int(cupy.cuda.Device().id), + "current_bytes": int(pool.used_bytes()), + "peak_bytes": None, + } + ) + except (ImportError, RuntimeError, AttributeError): + pass + return snapshots + + def finish(self): + result = { + "enabled": self.enabled, + "elapsed_seconds": time.perf_counter() - self.started, + } + if not self.enabled: + return result + current, peak = self._tracemalloc.get_traced_memory() + result.update( + { + "python_current_bytes": int(current), + "python_peak_bytes": int(peak), + "python_start_bytes": int(self._start[0]), + "host_rss_start_bytes": self._rss_start, + "host_rss_peak_bytes": self._rss(), + "gpu_start": self._gpu_start, + "gpu_end": self._gpu_snapshot(), + } + ) + if self._owns_tracemalloc: + self._tracemalloc.stop() + return result + + @dataclass(frozen=True) class OpenLoopObservableTerm: """An observable term with its physical support made explicit. @@ -127,6 +288,9 @@ class OpenLoopSeriesDiagnostic: total_flops_log10: float | None = None peak_memory_log2: float | None = None cache_hits: int = 0 + bp_converged: bool | None = None + bp_iterations: int | None = None + bp_max_mdiff: float | None = None @property def routes(self) -> dict[tuple[Any, ...], str]: @@ -155,6 +319,7 @@ class OpenLoopSeriesDiagnosticCache: the cache's ownership. """ + schema_version: ClassVar[str] = "open-loop-diagnostic-v2" diagnostics_by_key: dict[Any, OpenLoopSeriesDiagnostic] = field( default_factory=dict ) @@ -222,6 +387,7 @@ class _OpenEnumerationLimits: """Validated limits for lazy open-series geometry discovery.""" max_terms: int | None = None + max_loop_terms: int | None = None max_enumeration_time: float | None = None max_enumeration_memory: int | None = None @@ -230,6 +396,7 @@ def validate( cls, *, max_terms=None, + max_loop_terms=None, max_enumeration_time=None, max_enumeration_memory=None, ): @@ -237,6 +404,15 @@ def validate( if not isinstance(max_terms, (int, np.integer)) or max_terms < 0: raise ValueError("max_terms must be a non-negative integer or None") max_terms = int(max_terms) + if max_loop_terms is not None: + if ( + not isinstance(max_loop_terms, (int, np.integer)) + or max_loop_terms < 0 + ): + raise ValueError( + "max_loop_terms must be a non-negative integer or None" + ) + max_loop_terms = int(max_loop_terms) if max_enumeration_time is not None: if ( not isinstance( @@ -263,6 +439,7 @@ def validate( max_enumeration_memory = int(max_enumeration_memory) return cls( max_terms=max_terms, + max_loop_terms=max_loop_terms, max_enumeration_time=max_enumeration_time, max_enumeration_memory=max_enumeration_memory, ) @@ -275,6 +452,7 @@ def __init__(self, limits: _OpenEnumerationLimits): self.limits = limits self.started = time.perf_counter() self.emitted = 0 + self.loop_emitted = 0 self.estimated_memory = 0 @staticmethod @@ -297,7 +475,7 @@ def check(self): "max_enumeration_time", limit, elapsed ) - def accept(self, term: LoopSeriesTerm): + def accept(self, term: LoopSeriesTerm, *, loop=False): self.check() if ( self.limits.max_terms is not None @@ -306,6 +484,16 @@ def accept(self, term: LoopSeriesTerm): raise OpenLoopEnumerationLimitError( "max_terms", self.limits.max_terms, self.emitted + 1 ) + if ( + loop + and self.limits.max_loop_terms is not None + and self.loop_emitted >= self.limits.max_loop_terms + ): + raise OpenLoopEnumerationLimitError( + "max_loop_terms", + self.limits.max_loop_terms, + self.loop_emitted + 1, + ) term_memory = self._term_memory(term) if ( self.limits.max_enumeration_memory is not None @@ -318,11 +506,14 @@ def accept(self, term: LoopSeriesTerm): self.estimated_memory + term_memory, ) self.emitted += 1 + if loop: + self.loop_emitted += 1 self.estimated_memory += term_memory def diagnostics(self): return { "terms": self.emitted, + "loop_terms": self.loop_emitted, "elapsed_seconds": time.perf_counter() - self.started, "estimated_memory_bytes": self.estimated_memory, } @@ -349,11 +540,30 @@ class LoopSeriesCache: @staticmethod def _signature(tn): + def index_size(index): + try: + return int(tn.ind_size(index)) + except (AttributeError, TypeError, ValueError): + return None + + def tensor_shape(tensor): + try: + return tuple(tensor.shape) + except (AttributeError, TypeError): + return None + return ( frozenset(tn.tensor_map), frozenset( (index, frozenset(tids)) for index, tids in tn.ind_map.items() ), + frozenset( + (index, index_size(index)) for index in tn.ind_map + ), + frozenset( + (tid, tensor_shape(tensor)) + for tid, tensor in tn.tensor_map.items() + ), ) def _check_topology(self, tn) -> None: @@ -445,6 +655,7 @@ def iter_terms_for( excluded_edges=(), *, max_terms: int | None = None, + max_loop_terms: int | None = None, max_enumeration_time: float | None = None, max_enumeration_memory: int | None = None, ): @@ -461,6 +672,7 @@ def iter_terms_for( excluded_edges = frozenset(excluded_edges) limits = _OpenEnumerationLimits.validate( max_terms=max_terms, + max_loop_terms=max_loop_terms, max_enumeration_time=max_enumeration_time, max_enumeration_memory=max_enumeration_memory, ) @@ -485,7 +697,10 @@ def iter_terms_for( if cached is not None: guard = _OpenEnumerationGuard(limits) for term in cached: - guard.accept(term) + guard.accept( + term, + loop=_open_term_family(tn, term) != "open_path", + ) yield term return @@ -946,7 +1161,14 @@ def _weighted_shortest_distances(adjacency, target, guard=None): def _grid_corridor_context(tn): - """Return lazy coordinate-neighbor access for rectangular PEPS graphs.""" + """Return lazy coordinate-neighbor access for rectangular PEPS graphs. + + The returned graph is a *multigraph*. On a period-two PBC direction the + two seam bonds can connect the same pair of coordinates, but they remain + distinct tensor-network indices and therefore distinct loop-series edges. + Collapsing them to one coordinate neighbor loses valid paths and, for + native fermions, can also lose a seam sign route. + """ if not all(hasattr(tn, name) for name in ("Lx", "Ly", "has_site")): return None if not callable(getattr(tn, "site_tag", None)): @@ -968,6 +1190,20 @@ def tid_at(coo): tid_cache[coo] = tids[0] return tids[0] + def shared_edges(tid, neighbor_tid): + left_inds = set(tn.tensor_map[tid].inds) + right_inds = set(tn.tensor_map[neighbor_tid].inds) + return tuple( + sorted( + ( + index + for index in left_inds.intersection(right_inds) + if len(tn.ind_map[index]) == 2 + ), + key=repr, + ) + ) + def neighbors(coo): x, y = coo candidates = ((x - 1, y), (x + 1, y), (x, y - 1), (x, y + 1)) @@ -982,21 +1218,19 @@ def neighbors(coo): continue ny %= tn.Ly neighbor = (nx, ny) - if neighbor in seen or not tn.has_site(neighbor): + if not tn.has_site(neighbor): continue - seen.add(neighbor) neighbor_tid = tid_at(neighbor) if neighbor_tid is not None: - yield neighbor, neighbor_tid - - def edge_between(tid, neighbor_tid): - left_inds = set(tn.tensor_map[tid].inds) - right_inds = set(tn.tensor_map[neighbor_tid].inds) - shared = tuple(left_inds & right_inds) - for index in shared: - if len(tn.ind_map[index]) == 2: - return index - return None + for edge in shared_edges(tid_at(coo), neighbor_tid): + # The same coordinate can occur twice in the candidate + # list when a size-two axis wraps. Deduplicate only the + # exact multiedge record, never the coordinate itself. + record = (neighbor, edge) + if record in seen: + continue + seen.add(record) + yield neighbor, neighbor_tid, edge def axis_distance(left, right, size, cyclic): distance = abs(left - right) @@ -1010,7 +1244,7 @@ def distance(left, right): return { "tid_at": tid_at, "neighbors": neighbors, - "edge_between": edge_between, + "shared_edges": shared_edges, "distance": distance, } @@ -1082,13 +1316,11 @@ def _discover_grid_corridor_paths( continue if current_distance >= target_distance and current != source: continue - for neighbor, neighbor_tid in context["neighbors"](current): + for neighbor, neighbor_tid, edge in context["neighbors"]( + current + ): if neighbor in visited: continue - edge = context["edge_between"]( - context["tid_at"](current), - neighbor_tid, - ) if edge is None or edge in excluded_edges: continue if context["distance"](neighbor, target) != current_distance - 1: @@ -1122,7 +1354,7 @@ def _discover_grid_corridor_paths( coo, distance = pending.popleft() if distance >= corridor_width: continue - for neighbor, _ in context["neighbors"](coo): + for neighbor, _, _ in context["neighbors"](coo): if neighbor in corridor_coos: continue corridor_coos.add(neighbor) @@ -1133,13 +1365,9 @@ def _discover_grid_corridor_paths( context["tid_at"](coo) for coo in corridor_coos } for coo in corridor_coos: - for neighbor, neighbor_tid in context["neighbors"](coo): + for neighbor, neighbor_tid, edge in context["neighbors"](coo): if neighbor not in corridor_coos: continue - edge = context["edge_between"]( - context["tid_at"](coo), - neighbor_tid, - ) if edge is not None and edge not in excluded_edges: corridor_edges.add(edge) corridor_edges = frozenset(corridor_edges) @@ -1390,7 +1618,11 @@ def _iter_corridor_loop_clusters( guard.check() current, vertices, path_edges, visited = stack.pop() if ( - len(path_edges) >= 3 + # A pair of parallel PBC bonds is a legitimate + # two-edge cycle. This occurs routinely on period-2 + # PEPS axes and must not be discarded as a simple + # graph artifact. + len(path_edges) >= 2 and len(path_edges) <= max_size ): for neighbor, edge in by_vertex.get(current, ()): @@ -1481,7 +1713,7 @@ def _iter_corridor_open_terms( continue if loop.edges in seen: continue - guard.accept(loop) + guard.accept(loop, loop=True) seen.add(loop.edges) yield loop @@ -1501,7 +1733,7 @@ def _iter_corridor_open_terms( allowed_tids=allowed_tids, excluded_edges=excluded_edges, ) - guard.accept(term) + guard.accept(term, loop=True) seen.add(term.edges) yield term @@ -1577,7 +1809,7 @@ def visit(edge_pos, selected_count): frozenset(degrees), ) if term.edges not in path_terms: - guard.accept(term) + guard.accept(term, loop=True) yield term return @@ -2397,6 +2629,43 @@ def _validate_contraction_cost_limits( ) +def _validate_on_budget(on_budget): + """Validate the policy used when a term exceeds a contraction budget.""" + if on_budget not in {"report", "skip", "raise"}: + raise ValueError("on_budget must be 'report', 'skip', or 'raise'") + return on_budget + + +def _apply_budget_policy(skipped_terms, *, on_budget, info, kind): + """Record and optionally reject an incomplete contraction. + + ``report`` is the compatibility-safe default: it returns the partial + budgeted sum but marks it incomplete. ``raise`` is the recommended + production policy. ``skip`` is retained for exploratory workflows and + is intentionally equivalent numerically to the historical behavior. + """ + on_budget = _validate_on_budget(on_budget) + skipped_terms = dict(skipped_terms or {}) + complete = not skipped_terms + if info is not None: + info[f"{kind}_complete"] = complete + info[f"{kind}_budget_policy"] = on_budget + info[f"{kind}_omitted_terms"] = tuple(skipped_terms) + elif skipped_terms and on_budget == "report": + warnings.warn( + f"{kind} returned an incomplete budgeted sum; pass info=... or " + "use on_budget='raise' to make this failure explicit", + RuntimeWarning, + stacklevel=3, + ) + if skipped_terms and on_budget == "raise": + raise OpenLoopBudgetError( + f"{kind} skipped {len(skipped_terms)} term(s) because their " + "contraction estimate exceeded the configured budget" + ) + return complete, tuple(skipped_terms) + + def _contract_cost_record(tree, *, max_bond=None): """Extract the standard Cotengra log-cost diagnostics from a tree.""" if max_bond is None: @@ -3152,6 +3421,7 @@ def _open_edge_series_terms_for_support( *, cache, max_terms=None, + max_loop_terms=None, max_enumeration_time=None, max_enumeration_memory=None, corridor_width=None, @@ -3170,6 +3440,7 @@ def _open_edge_series_terms_for_support( allowed_tids = frozenset(tids) limits = _OpenEnumerationLimits.validate( max_terms=max_terms, + max_loop_terms=max_loop_terms, max_enumeration_time=max_enumeration_time, max_enumeration_memory=max_enumeration_memory, ) @@ -3227,6 +3498,7 @@ def corridor_terms(): "max_loop_clusters_per_segment" ], "max_edge_cutoff": total_edge_cutoff, + "max_loop_terms": limits.max_loop_terms, "approximation": "path_plus_connected_loop_decorations", } ) @@ -3279,6 +3551,7 @@ def corridor_terms(): allowed_tids, excluded_edges=inner_bonds, max_terms=limits.max_terms, + max_loop_terms=limits.max_loop_terms, max_enumeration_time=limits.max_enumeration_time, max_enumeration_memory=limits.max_enumeration_memory, ) @@ -3304,6 +3577,7 @@ def corridor_terms(): allowed_tids, excluded_edges=inner_bonds, max_terms=limits.max_terms, + max_loop_terms=limits.max_loop_terms, max_enumeration_time=limits.max_enumeration_time, max_enumeration_memory=limits.max_enumeration_memory, ) @@ -3317,12 +3591,33 @@ def limited_terms(): edge_cutoff, inner_bonds=inner_bonds, ): - guard.accept(term) + guard.accept( + term, + loop=_open_term_family(bp.tn, term) != "open_path", + ) yield term return limited_terms(), inner_bonds +def _limit_diagnostic_open_terms(tn, terms, *, max_loop_terms): + """Apply a loop-term budget even when geometry came from a diagnostic.""" + if max_loop_terms is None: + return iter(terms) + limits = _OpenEnumerationLimits.validate(max_loop_terms=max_loop_terms) + guard = _OpenEnumerationGuard(limits) + + def limited_terms(): + for term in terms: + guard.accept( + term, + loop=_open_term_family(tn, term) != "open_path", + ) + yield term + + return limited_terms() + + def _log10_sum_costs(costs): """Sum positive costs represented in base-10 logarithmic form.""" values = [ @@ -3356,6 +3651,8 @@ def _diagnose_network_cost( optimize, contract_opts, max_bond=None, + plan_cache=None, + plan_key=None, ): """Build a contraction tree and return costs without contracting data.""" if "get" in contract_opts: @@ -3364,7 +3661,10 @@ def _diagnose_network_cost( "diagnostics; diagnostics always build a contraction tree" ) tree = network.contract(get="tree", optimize=optimize, **contract_opts) - return _contract_cost_record(tree, max_bond=max_bond) + cost = _contract_cost_record(tree, max_bond=max_bond) + if plan_cache is not None and plan_key is not None: + plan_cache[plan_key] = (tree, cost) + return cost def _open_diagnostic_key( @@ -3376,11 +3676,14 @@ def _open_diagnostic_key( cluster_size, corridor_options, max_terms, + max_loop_terms, max_enumeration_time, max_enumeration_memory, max_flops_log10, max_peak_memory_log2, path_edge_weights, + optimize=None, + contract_opts=None, ): """Build a stable cache key for geometry and cost diagnostics.""" try: @@ -3392,17 +3695,26 @@ def _open_diagnostic_key( except Exception: dtype = repr(type(gate)) return ( + "open-loop-diagnostic-v2", tuple(sites), route, repr(edge_cutoff), repr(cluster_size), tuple(sorted((key, repr(value)) for key, value in corridor_options.items())), max_terms, + max_loop_terms, max_enumeration_time, max_enumeration_memory, max_flops_log10, max_peak_memory_log2, repr(path_edge_weights), + repr(optimize), + tuple( + sorted( + (str(key), repr(value)) + for key, value in (contract_opts or {}).items() + ) + ), shape, dtype, ) @@ -3424,6 +3736,7 @@ def _diagnose_open_scalar_support( max_flops_log10, max_peak_memory_log2, max_terms, + max_loop_terms, max_enumeration_time, max_enumeration_memory, corridor_options, @@ -3442,6 +3755,7 @@ def _diagnose_open_scalar_support( fermionic_q = _uses_symmray(bp.tn) and _gate_needs_fermionic_open_q(gate) route = route_selection["route"] total_costs = [] + plan_cache = {} if route == "graded_cluster_compatible": from quimb.tensor.belief_propagation import gen_region_counts @@ -3510,6 +3824,7 @@ def _diagnose_open_scalar_support( edge_cutoff, cache=cache, max_terms=max_terms, + max_loop_terms=max_loop_terms, max_enumeration_time=max_enumeration_time, max_enumeration_memory=max_enumeration_memory, corridor_width=( @@ -3537,6 +3852,14 @@ def _diagnose_open_scalar_support( compressed = compressed_max_bond is not None for term in requested_terms: region = frozenset((*tids, *term.tids)) + norm_path_key = ( + "norm", where_key, tuple(sorted(region, key=repr)), + tuple(sorted(inner_bonds, key=repr)), fermionic_q, + ) + gate_path_key = ( + "gate", where_key, tuple(sorted(region, key=repr)), + tuple(sorted(inner_bonds, key=repr)), fermionic_q, + ) norm_network = _get_d2_edge_partial_trace_excited( bp, region, @@ -3544,7 +3867,7 @@ def _diagnose_open_scalar_support( exclude=inner_bonds, projector_layout="open" if _uses_symmray(bp.tn) else "series", fermionic_q=fermionic_q, - index_namespace=("diagnostic", "norm", where_key, term.edges), + index_namespace=("open-scalar", *norm_path_key), ) gate_network = _get_d2_edge_partial_trace_excited( bp, @@ -3556,19 +3879,23 @@ def _diagnose_open_scalar_support( projector_layout="open" if _uses_symmray(bp.tn) else "series", gate_as_operator=True, fermionic_q=fermionic_q, - index_namespace=("diagnostic", "gate", where_key, term.edges), + index_namespace=("open-scalar", *gate_path_key), ) norm_cost = _diagnose_network_cost( norm_network, optimize=optimize, contract_opts=contract_opts, max_bond=compressed_max_bond if compressed else None, + plan_cache=plan_cache, + plan_key=norm_path_key, ) gate_cost = _diagnose_network_cost( gate_network, optimize=optimize, contract_opts=contract_opts, max_bond=compressed_max_bond if compressed else None, + plan_cache=plan_cache, + plan_key=gate_path_key, ) record = { "norm": norm_cost, @@ -3596,7 +3923,11 @@ def _diagnose_open_scalar_support( exclude=inner_bonds, projector_layout="open" if _uses_symmray(bp.tn) else "series", fermionic_q=fermionic_q, - index_namespace=("diagnostic", "base-norm", where_key), + index_namespace=( + "open-scalar", "base-norm", where_key, + tuple(sorted(tids, key=repr)), + tuple(sorted(inner_bonds, key=repr)), fermionic_q, + ), ) base_gate_network = _get_d2_edge_partial_trace_excited( bp, @@ -3607,19 +3938,33 @@ def _diagnose_open_scalar_support( projector_layout="open" if _uses_symmray(bp.tn) else "series", gate_as_operator=True, fermionic_q=fermionic_q, - index_namespace=("diagnostic", "base-gate", where_key), + index_namespace=( + "open-scalar", "base-gate", where_key, + tuple(sorted(tids, key=repr)), + tuple(sorted(inner_bonds, key=repr)), fermionic_q, + ), ) base_norm_cost = _diagnose_network_cost( base_norm_network, optimize=optimize, contract_opts=contract_opts, max_bond=compressed_max_bond if compressed else None, + plan_cache=plan_cache, + plan_key=( + "base-norm", where_key, tuple(sorted(tids, key=repr)), + tuple(sorted(inner_bonds, key=repr)), fermionic_q, + ), ) base_gate_cost = _diagnose_network_cost( base_gate_network, optimize=optimize, contract_opts=contract_opts, max_bond=compressed_max_bond if compressed else None, + plan_cache=plan_cache, + plan_key=( + "base-gate", where_key, tuple(sorted(tids, key=repr)), + tuple(sorted(inner_bonds, key=repr)), fermionic_q, + ), ) total_costs.extend((base_norm_cost, base_gate_cost)) base_cost = { @@ -3650,6 +3995,8 @@ def _diagnose_open_scalar_support( default=None, ), "inner_bonds": inner_bonds, + "contraction_plans": plan_cache, + "plan_cache_schema": "open-scalar-v1", } @@ -3790,6 +4137,7 @@ def _partial_trace_open_loop_series( max_flops_log10, max_peak_memory_log2, max_terms, + max_loop_terms, max_enumeration_time, max_enumeration_memory, corridor_width, @@ -3804,6 +4152,7 @@ def _partial_trace_open_loop_series( mode, auto_corridor_distance, diagnostic_support, + on_budget, ): """Contract the explicit open-edge rho loop-series expansion.""" if bp.__class__.__name__ != "D2BP": @@ -3920,19 +4269,19 @@ def _partial_trace_open_loop_series( max_flops_log10=max_flops_log10, max_peak_memory_log2=max_peak_memory_log2, ) + cluster_region_costs = { + (where_key, region): cost + for region, cost in cluster_info.get( + "cluster_rho_term_costs", {} + ).items() + } + cluster_region_skipped = { + (where_key, region): cost + for region, cost in cluster_info.get( + "cluster_rho_skipped_terms", {} + ).items() + } if info is not None: - cluster_region_costs = { - (where_key, region): cost - for region, cost in cluster_info.get( - "cluster_rho_term_costs", {} - ).items() - } - cluster_region_skipped = { - (where_key, region): cost - for region, cost in cluster_info.get( - "cluster_rho_skipped_terms", {} - ).items() - } info["open_rho_requested_terms"] = () info["open_rho_terms_list"] = () info["open_rho_term_costs"] = dict( @@ -3959,10 +4308,12 @@ def _partial_trace_open_loop_series( info["open_rho_support_tids"] = tids info["open_rho_excluded_edges"] = inner_bonds info["open_rho_native_route"] = "graded_cluster_compatible" + info["open_rho_route"] = route_selection["route"] info["open_rho_edge_cutoff"] = None info["open_rho_cluster_size"] = cluster_size info["open_rho_enumeration_limits"] = { "max_terms": max_terms, + "max_loop_terms": max_loop_terms, "max_enumeration_time": max_enumeration_time, "max_enumeration_memory": max_enumeration_memory, } @@ -3975,6 +4326,20 @@ def _partial_trace_open_loop_series( if diagnostic_support is None else dict(diagnostic_support) ) + complete, omitted = _apply_budget_policy( + cluster_region_skipped, + on_budget=on_budget, + info=info, + kind="open_rho", + ) + info["open_rho_omitted_terms"] = omitted + else: + _apply_budget_policy( + cluster_region_skipped, + on_budget=on_budget, + info=None, + kind="open_rho", + ) return rho corridor_info = ( @@ -3988,7 +4353,11 @@ def _partial_trace_open_loop_series( "the supplied open-series diagnostic does not match the " "route selected for this rho support" ) - terms = iter(diagnostic_support.get("terms", ())) + terms = _limit_diagnostic_open_terms( + bp.tn, + diagnostic_support.get("terms", ()), + max_loop_terms=max_loop_terms, + ) if corridor_info is not None: corridor_info.clear() corridor_info.update(diagnostic_support.get("corridor", {})) @@ -4004,6 +4373,7 @@ def _partial_trace_open_loop_series( edge_cutoff, cache=cache, max_terms=max_terms, + max_loop_terms=max_loop_terms, max_enumeration_time=max_enumeration_time, max_enumeration_memory=max_enumeration_memory, corridor_width=corridor_options["corridor_width"], @@ -4088,6 +4458,13 @@ def _partial_trace_open_loop_series( rho_terms[term.edges] = rho_e accepted_terms.append(term) + complete, omitted_terms = _apply_budget_policy( + skipped_terms, + on_budget=on_budget, + info=info, + kind="open_rho", + ) + base_cache = ( {} if info is None @@ -4198,11 +4575,13 @@ def _partial_trace_open_loop_series( info["open_rho_cluster_size"] = None info["open_rho_enumeration_limits"] = { "max_terms": max_terms, + "max_loop_terms": max_loop_terms, "max_enumeration_time": max_enumeration_time, "max_enumeration_memory": max_enumeration_memory, } info["open_rho_corridor_options"] = dict(corridor_options) info["open_rho_mode"] = mode + info["open_rho_route"] = route_selection["route"] info["open_rho_support_distance"] = route_selection[ "support_distance" ] @@ -4211,6 +4590,7 @@ def _partial_trace_open_loop_series( if diagnostic_support is None else dict(diagnostic_support) ) + info["open_rho_omitted_terms"] = omitted_terms return rho @@ -4354,6 +4734,7 @@ def _local_expectation_open_loop_series( max_flops_log10, max_peak_memory_log2, max_terms, + max_loop_terms, max_enumeration_time, max_enumeration_memory, corridor_width, @@ -4368,6 +4749,7 @@ def _local_expectation_open_loop_series( mode, auto_corridor_distance, diagnostic_support, + on_budget, ): """Contract a gate through the explicit open-edge loop series.""" if normalized == "prod": @@ -4476,6 +4858,12 @@ def _local_expectation_open_loop_series( max_flops_log10=max_flops_log10, max_peak_memory_log2=max_peak_memory_log2, ) + cluster_region_skipped = { + (where_key, region): cost + for region, cost in cluster_info.get( + "cluster_scalar_skipped_terms", {} + ).items() + } if info is not None: cluster_region_costs = { (where_key, region): cost @@ -4483,12 +4871,6 @@ def _local_expectation_open_loop_series( "cluster_scalar_term_costs", {} ).items() } - cluster_region_skipped = { - (where_key, region): cost - for region, cost in cluster_info.get( - "cluster_scalar_skipped_terms", {} - ).items() - } info["open_scalar_requested_terms"] = () info["open_scalar_terms"] = () info["open_scalar_skipped_terms"] = dict( @@ -4509,6 +4891,7 @@ def _local_expectation_open_loop_series( info["open_scalar_native_route"] = ( "graded_cluster_compatible" ) + info["open_scalar_route"] = route_selection["route"] info["open_scalar_fermionic_q_phase"] = fermionic_q info["open_scalar_cost_limits"] = { "max_flops_log10": max_flops_log10, @@ -4524,6 +4907,7 @@ def _local_expectation_open_loop_series( info["open_scalar_cluster_size"] = cluster_size info["open_scalar_enumeration_limits"] = { "max_terms": max_terms, + "max_loop_terms": max_loop_terms, "max_enumeration_time": max_enumeration_time, "max_enumeration_memory": max_enumeration_memory, } @@ -4536,6 +4920,20 @@ def _local_expectation_open_loop_series( if diagnostic_support is None else dict(diagnostic_support) ) + complete, omitted = _apply_budget_policy( + cluster_region_skipped, + on_budget=on_budget, + info=info, + kind="open_scalar", + ) + info["open_scalar_omitted_terms"] = omitted + else: + _apply_budget_policy( + cluster_region_skipped, + on_budget=on_budget, + info=None, + kind="open_scalar", + ) return value, normalization corridor_info = ( @@ -4544,7 +4942,11 @@ def _local_expectation_open_loop_series( else info.setdefault("open_scalar_corridor", {}) ) if diagnostic_support is not None: - terms = iter(diagnostic_support.get("terms", ())) + terms = _limit_diagnostic_open_terms( + bp.tn, + diagnostic_support.get("terms", ()), + max_loop_terms=max_loop_terms, + ) if corridor_info is not None: corridor_info.clear() corridor_info.update(diagnostic_support.get("corridor", {})) @@ -4560,6 +4962,7 @@ def _local_expectation_open_loop_series( edge_cutoff, cache=cache, max_terms=max_terms, + max_loop_terms=max_loop_terms, max_enumeration_time=max_enumeration_time, max_enumeration_memory=max_enumeration_memory, corridor_width=corridor_options["corridor_width"], @@ -4618,6 +5021,11 @@ def _local_expectation_open_loop_series( if info is None else info.setdefault("open_scalar_region_path_cache", {}) ) + if diagnostic_support is not None: + for plan_key, plan in diagnostic_support.get( + "contraction_plans", {} + ).items(): + path_cache.setdefault(plan_key, plan) for term in terms: requested_terms.append(term) if (where_key, term.edges) in skipped_terms: @@ -4729,6 +5137,13 @@ def _local_expectation_open_loop_series( norm_terms[term.edges] = norm_e gate_terms[term.edges] = gate_e + complete, omitted_terms = _apply_budget_policy( + skipped_terms, + on_budget=on_budget, + info=info, + kind="open_scalar", + ) + base_key = (where_key, tuple(kix), tids, inner_bonds, fermionic_q) base_norm_path_key = ( "base-norm", @@ -4866,11 +5281,13 @@ def _local_expectation_open_loop_series( if _uses_symmray(bp.tn) else "dense_open_projectors" ) + info["open_scalar_route"] = route_selection["route"] info["open_scalar_fermionic_q_phase"] = fermionic_q info["open_scalar_edge_cutoff"] = edge_cutoff info["open_scalar_cluster_size"] = None info["open_scalar_enumeration_limits"] = { "max_terms": max_terms, + "max_loop_terms": max_loop_terms, "max_enumeration_time": max_enumeration_time, "max_enumeration_memory": max_enumeration_memory, } @@ -4884,6 +5301,7 @@ def _local_expectation_open_loop_series( if diagnostic_support is None else dict(diagnostic_support) ) + info["open_scalar_omitted_terms"] = omitted_terms return value, norm @@ -5626,6 +6044,7 @@ def partial_trace_open_loop_series_sweep( edge_cutoffs=None, cluster_sizes=None, max_terms: int | None = None, + max_loop_terms: int | None = None, max_enumeration_time: float | None = None, max_enumeration_memory: int | None = None, mode: str = "exact", @@ -5639,6 +6058,7 @@ def partial_trace_open_loop_series_sweep( max_corridor_edges: int | None = 100_000, path_edge_weights=None, corridor_max_bond: int | None = None, + on_budget: str = "report", contract_opts: dict[str, Any] | None = None, **bp_opts, ) -> OpenLoopSeriesSweepResult: @@ -5714,6 +6134,12 @@ def partial_trace_open_loop_series_sweep( raise ValueError("each support must contain at least one site") contract_opts = {} if contract_opts is None else dict(contract_opts) + _OpenEnumerationLimits.validate( + max_terms=max_terms, + max_loop_terms=max_loop_terms, + max_enumeration_time=max_enumeration_time, + max_enumeration_memory=max_enumeration_memory, + ) cache = cache or OpenLoopSeriesCache() bp, bp_info = _build_bp( tn, @@ -5766,6 +6192,7 @@ def partial_trace_open_loop_series_sweep( cache=cache, info=support_info, max_terms=max_terms, + max_loop_terms=max_loop_terms, max_enumeration_time=max_enumeration_time, max_enumeration_memory=max_enumeration_memory, corridor_width=corridor_width, @@ -5780,6 +6207,7 @@ def partial_trace_open_loop_series_sweep( mode=mode, auto_corridor_distance=auto_corridor_distance, diagnostic_support=None, + on_budget=on_budget, ) support_rhos[cutoff] = rho support_diagnostics[cutoff] = { @@ -5863,6 +6291,7 @@ def partial_trace_open_loop_series_expand( edge_cutoff=None, cluster_size=None, max_terms: int | None = None, + max_loop_terms: int | None = None, max_enumeration_time: float | None = None, max_enumeration_memory: int | None = None, mode: str = "exact", @@ -5876,6 +6305,9 @@ def partial_trace_open_loop_series_expand( max_corridor_edges: int | None = 100_000, path_edge_weights=None, corridor_max_bond: int | None = None, + on_budget: str = "report", + measure_resources: bool = False, + return_result: bool = False, diagnostic: OpenLoopSeriesDiagnostic | None = None, info: dict[str, Any] | None = None, contract_opts: dict[str, Any] | None = None, @@ -5939,6 +6371,10 @@ def partial_trace_open_loop_series_expand( Hard limit on discovered explicit edge terms. Exceeding it raises :class:`OpenLoopEnumerationLimitError`; partial sums are never returned silently. + max_loop_terms : int, optional + Separate hard limit on loop-containing configurations (closed loops + and path-plus-loop terms). Open support paths do not consume this + budget. max_enumeration_time : float, optional Maximum edge-geometry discovery time in seconds. max_enumeration_memory : int, optional @@ -5978,6 +6414,15 @@ def partial_trace_open_loop_series_expand( Optional Cotengra tree-cost limits. Terms over either limit are skipped and recorded in ``info``; the unexcited base configuration must still fit both limits. + on_budget : {"report", "skip", "raise"}, optional + Policy for terms rejected by the FLOP or peak-memory limits. + ``"report"`` preserves the historical partial result and marks it + incomplete; ``"raise"`` is recommended for production runs. + measure_resources : bool, optional + Record observed Python, host-RSS, and available GPU resource data. + return_result : bool, optional + Return :class:`OpenLoopMeasurementResult` with the value, route, + completeness, BP metadata, diagnostics, and resource observations. cache : OpenLoopSeriesCache, optional Reusable open-term geometry cache for the same topology and support. info : dict, optional @@ -5990,6 +6435,17 @@ def partial_trace_open_loop_series_expand( older ``open_rho_term_costs`` fields are route-specific aliases. """ contract_opts = {} if contract_opts is None else dict(contract_opts) + on_budget = _validate_on_budget(on_budget) + _OpenEnumerationLimits.validate( + max_terms=max_terms, + max_loop_terms=max_loop_terms, + max_enumeration_time=max_enumeration_time, + max_enumeration_memory=max_enumeration_memory, + ) + monitor = _OpenLoopResourceMonitor(measure_resources) + internal_info = info + if internal_info is None and (return_result or measure_resources): + internal_info = {} where = tuple(where) if not where: raise ValueError("where must contain at least one site") @@ -6027,13 +6483,21 @@ def partial_trace_open_loop_series_expand( "messages; pass require_fixed_point=False for an exploratory " "estimate" ) + if internal_info is not None: + internal_info.update( + { + "bp_converged": bp_info.get("converged"), + "bp_iterations": bp_info.get("iterations"), + "bp_max_mdiff": bp_info.get("max_mdiff"), + } + ) diagnostic_support = ( None if diagnostic is None else diagnostic.supports.get(tuple(where)) ) - return _partial_trace_open_loop_series( + value = _partial_trace_open_loop_series( bp, where, gloops, @@ -6043,10 +6507,11 @@ def partial_trace_open_loop_series_expand( max_peak_memory_log2=max_peak_memory_log2, contract_opts=contract_opts, cache=cache or OpenLoopSeriesCache(), - info=info, + info=internal_info, edge_cutoff=edge_cutoff, cluster_size=cluster_size, max_terms=max_terms, + max_loop_terms=max_loop_terms, max_enumeration_time=max_enumeration_time, max_enumeration_memory=max_enumeration_memory, corridor_width=corridor_width, @@ -6061,6 +6526,40 @@ def partial_trace_open_loop_series_expand( mode=mode, auto_corridor_distance=auto_corridor_distance, diagnostic_support=diagnostic_support, + on_budget=on_budget, + ) + resources = monitor.finish() + if internal_info is not None: + internal_info["open_rho_resources"] = resources + if not return_result: + return value + support_info = internal_info or {} + route = support_info.get("open_rho_route") + if route is None: + route = ( + "corridor" + if support_info.get("open_rho_mode") == "corridor" + else support_info.get("open_rho_native_route") + ) + skipped = tuple(support_info.get("open_rho_omitted_terms", ())) + return OpenLoopMeasurementResult( + value=value, + normalization=_rho_trace(value), + info=support_info, + diagnostic=diagnostic, + complete=support_info.get("open_rho_complete", not skipped), + approximate=( + support_info.get("open_rho_mode") == "corridor" + or support_info.get("open_rho_native_route") + == "graded_cluster_compatible" + ), + route=route, + omitted_terms=skipped, + resources=resources, + bp_converged=bp_info.get("converged"), + bp_iterations=bp_info.get("iterations"), + bp_max_mdiff=bp_info.get("max_mdiff"), + bp=bp, ) @@ -6511,6 +7010,7 @@ def diagnose_open_loop_series( edge_cutoff=None, cluster_size=None, max_terms: int | None = None, + max_loop_terms: int | None = None, max_enumeration_time: float | None = None, max_enumeration_memory: int | None = None, mode: str = "auto", @@ -6524,6 +7024,7 @@ def diagnose_open_loop_series( max_corridor_edges: int | None = 100_000, path_edge_weights=None, corridor_max_bond: int | None = None, + on_budget: str = "report", contract_opts: dict[str, Any] | None = None, **bp_opts, ): @@ -6542,6 +7043,13 @@ def diagnose_open_loop_series( corridor route, with ``auto_corridor_distance`` controlling that switch. """ records = _normalize_open_observable_terms(terms) + on_budget = _validate_on_budget(on_budget) + _OpenEnumerationLimits.validate( + max_terms=max_terms, + max_loop_terms=max_loop_terms, + max_enumeration_time=max_enumeration_time, + max_enumeration_memory=max_enumeration_memory, + ) contract_opts = {} if contract_opts is None else dict(contract_opts) if require_fixed_point and run_bp and ( not isinstance(max_iterations, (int, np.integer)) or max_iterations < 1 @@ -6614,11 +7122,14 @@ def diagnose_open_loop_series( cluster_size=cluster_value, corridor_options=corridor_options, max_terms=max_terms, + max_loop_terms=max_loop_terms, max_enumeration_time=max_enumeration_time, max_enumeration_memory=max_enumeration_memory, max_flops_log10=max_flops_log10, max_peak_memory_log2=max_peak_memory_log2, path_edge_weights=path_edge_weights, + optimize=optimize, + contract_opts=contract_opts, ) report = diagnostic_cache.get(bp.tn, key) if report is not None: @@ -6640,6 +7151,7 @@ def diagnose_open_loop_series( max_flops_log10=max_flops_log10, max_peak_memory_log2=max_peak_memory_log2, max_terms=max_terms, + max_loop_terms=max_loop_terms, max_enumeration_time=max_enumeration_time, max_enumeration_memory=max_enumeration_memory, corridor_options=corridor_options, @@ -6650,10 +7162,22 @@ def diagnose_open_loop_series( "mode": mode, "edge_cutoff": edge_value, "cluster_size": cluster_value, + "max_loop_terms": max_loop_terms, "corridor_options": corridor_options, "cache_key": key, } ) + skipped = report_data.get("skipped_terms", {}) + skipped = skipped or report_data.get("cluster_region_skipped_terms", {}) + report_data["complete"] = not bool(skipped) + report_data["omitted_terms"] = tuple(skipped) + if skipped and on_budget == "raise": + _apply_budget_policy( + skipped, + on_budget=on_budget, + info=None, + kind="open_loop_diagnostic", + ) support_report = OpenLoopSeriesDiagnostic({tuple(sites): report_data}) diagnostic_cache.put(bp.tn, key, support_report) support_reports[tuple(sites)] = report_data @@ -6676,6 +7200,206 @@ def diagnose_open_loop_series( total_flops_log10=total_flops, peak_memory_log2=max(peaks, default=None), cache_hits=cache_hits, + bp_converged=bp_info.get("converged"), + bp_iterations=bp_info.get("iterations"), + bp_max_mdiff=bp_info.get("max_mdiff"), + ) + + +def diagnose_open_rho_series( + tn, + supports, + *, + gloops=None, + edge_cutoff=None, + cluster_size=None, + mode="auto", + max_rho_identity_dimension=2048, + **kwargs, +): + """Preflight open-rho geometry and output-resource costs. + + The loop geometry is shared with scalar observables, while the report also + accounts for the physical output tensor. The identity gate is used only + to exercise the same support/route diagnostics; no numerical contraction + is performed. Large retained regions should be diagnosed with a custom + application-level output budget rather than materializing a huge identity. + """ + supports = tuple(tuple(support) for support in supports) + if not supports: + raise ValueError("supports must contain at least one support") + if ( + not isinstance(max_rho_identity_dimension, (int, np.integer)) + or max_rho_identity_dimension < 1 + ): + raise ValueError("max_rho_identity_dimension must be a positive integer") + max_rho_identity_dimension = int(max_rho_identity_dimension) + terms = {} + for support in supports: + dims = tuple(int(tn.ind_size(tn.site_ind(site))) for site in support) + dimension = int(np.prod(dims, dtype=np.int64)) + if dimension > max_rho_identity_dimension: + raise ValueError( + "diagnose_open_rho_series would materialize an identity of " + f"dimension {dimension}; lower the support or increase " + "max_rho_identity_dimension" + ) + terms[support] = np.eye(dimension) + diagnostic = diagnose_open_loop_series( + tn, + terms, + gloops=gloops, + edge_cutoff=edge_cutoff, + cluster_size=cluster_size, + mode=mode, + **kwargs, + ) + itemsize = np.dtype(np.complex128).itemsize + for support in supports: + record = diagnostic.supports[support] + dims = tuple(int(tn.ind_size(tn.site_ind(site))) for site in support) + dimension = int(np.prod(dims, dtype=np.int64)) + # The public rho API returns a flattened ket-by-bra matrix. Keep the + # logical tensor-index shape alongside it for callers that need to + # reshape the result into one axis per retained site. + output_shape = (dimension, dimension) + logical_output_shape = dims + dims + output_elements = int(np.prod(output_shape, dtype=np.int64)) + output_bytes = output_elements * itemsize + record.update( + { + "observable_kind": "rho", + "output_shape": output_shape, + "logical_output_shape": logical_output_shape, + "output_elements": output_elements, + "output_memory_bytes": output_bytes, + "rho_peak_memory_log2": float( + max( + record.get("peak_memory_log2", 0.0), + np.log2(max(1, output_elements)), + ) + ), + } + ) + record["peak_memory_log2"] = record["rho_peak_memory_log2"] + diagnostic.peak_memory_log2 = max( + (record["peak_memory_log2"] for record in diagnostic.supports.values()), + default=None, + ) + return diagnostic + + +def adaptive_open_loop_series( + tn, + terms, + *, + corridor_widths=(0, 1, 2), + cluster_sizes=None, + atol=1e-8, + rtol=1e-6, + min_stable=2, + diagnostic_cache=None, + **kwargs, +): + """Measure an observable along an adaptive corridor/cluster ladder. + + Each level reuses the first level's BP messages and the diagnostic cache. + For native cyclic fermions pass ``cluster_sizes``; for dense/tree + networks pass ``corridor_widths``. The routine reports convergence of the + numerical value, not a rigorous truncation bound. + """ + raw_levels = cluster_sizes if cluster_sizes is not None else corridor_widths + if isinstance(raw_levels, (int, np.integer)): + raw_levels = (raw_levels,) + try: + levels = tuple(_validate_nonnegative_degree(level) for level in raw_levels) + except TypeError as exc: + raise TypeError("adaptive levels must be an integer or iterable of integers") from exc + if len(levels) == 0: + raise ValueError("the adaptive ladder must contain at least one level") + if not isinstance(min_stable, (int, np.integer)) or min_stable < 1: + raise ValueError("min_stable must be a positive integer") + for name, value in (("atol", atol), ("rtol", rtol)): + if not isinstance(value, (int, float, np.integer, np.floating)): + raise TypeError(f"{name} must be a real number") + if not np.isfinite(value) or value < 0: + raise ValueError(f"{name} must be finite and non-negative") + min_stable = int(min_stable) + atol = float(atol) + rtol = float(rtol) + diagnostic_cache = diagnostic_cache or OpenLoopSeriesDiagnosticCache() + values = [] + settings = [] + differences = [] + diagnostics = [] + infos = [] + shared_messages = kwargs.pop("messages", None) + run_bp = kwargs.pop("run_bp", True) + bp = None + stable_count = 0 + selected_index = None + for level in levels: + level_kwargs = dict(kwargs) + level_kwargs["diagnostic_cache"] = diagnostic_cache + level_kwargs["return_result"] = True + level_info = {} + level_kwargs["info"] = level_info + if shared_messages is not None: + level_kwargs.update(messages=shared_messages, run_bp=False) + else: + level_kwargs["run_bp"] = run_bp + if cluster_sizes is not None: + level_kwargs.update(cluster_size=int(level), mode="auto") + setting = {"cluster_size": int(level), "mode": "auto"} + else: + # ``mode="auto"`` retains the explicit width while enabling the + # diagnostic prepass and cache reuse for every ladder level. + level_kwargs.update(corridor_width=int(level), mode="auto") + setting = {"corridor_width": int(level), "mode": "auto"} + result = compute_local_expectation_open_loop_series( + tn, terms, **level_kwargs + ) + if not isinstance(result, OpenLoopMeasurementResult): + raise RuntimeError("adaptive measurement requires return_result support") + if bp is None: + bp = result.bp + if bp is not None: + shared_messages = bp.messages + value = result.value + values.append(value) + settings.append(setting) + diagnostics.append(result.diagnostic) + infos.append(result.info) + if len(values) == 1: + differences.append(None) + continue + try: + difference = float(np.max(np.abs(np.asarray(value) - np.asarray(values[-2])))) + scale = max(1.0, float(np.max(np.abs(np.asarray(value))))) + except (TypeError, ValueError): + difference = None + scale = 1.0 + differences.append(difference) + if ( + result.complete + and difference is not None + and difference <= atol + rtol * scale + ): + stable_count += 1 + else: + stable_count = 0 + if stable_count >= min_stable: + selected_index = len(values) - 1 + break + return OpenLoopAdaptiveResult( + values=tuple(values), + settings=tuple(settings), + differences=tuple(differences), + converged=selected_index is not None, + selected_index=selected_index, + diagnostics=tuple(diagnostics), + infos=tuple(infos), + bp=bp, ) @@ -6705,6 +7429,7 @@ def compute_local_expectation_open_loop_series( edge_cutoff=None, cluster_size=None, max_terms: int | None = None, + max_loop_terms: int | None = None, max_enumeration_time: float | None = None, max_enumeration_memory: int | None = None, mode: str = "exact", @@ -6718,6 +7443,9 @@ def compute_local_expectation_open_loop_series( max_corridor_edges: int | None = 100_000, path_edge_weights=None, corridor_max_bond: int | None = None, + on_budget: str = "report", + measure_resources: bool = False, + return_result: bool = False, diagnostic_cache: OpenLoopSeriesDiagnosticCache | None = None, diagnostic: OpenLoopSeriesDiagnostic | None = None, info: dict[str, Any] | None = None, @@ -6771,6 +7499,10 @@ def compute_local_expectation_open_loop_series( bound violation raises :class:`OpenLoopEnumerationLimitError` rather than returning a silently incomplete observable. + ``max_loop_terms`` is an additional loop-only bound. It is useful when + nearby supports can afford many loop corrections while long-range + supports should retain only a small loop tail. + ``corridor_width`` activates the bounded long-separation route: weighted shortest paths are retained in a small beam, the paths are inflated into a corridor, and connected loop decorations are sampled near its segments. @@ -6787,6 +7519,17 @@ def compute_local_expectation_open_loop_series( "normalized must be one of True, False, 'prod', or 'separate'" ) contract_opts = {} if contract_opts is None else dict(contract_opts) + on_budget = _validate_on_budget(on_budget) + _OpenEnumerationLimits.validate( + max_terms=max_terms, + max_loop_terms=max_loop_terms, + max_enumeration_time=max_enumeration_time, + max_enumeration_memory=max_enumeration_memory, + ) + monitor = _OpenLoopResourceMonitor(measure_resources) + internal_info = info + if internal_info is None and (return_result or measure_resources): + internal_info = {} bp, bp_info = _build_bp( tn, norm="2norm", @@ -6812,6 +7555,14 @@ def compute_local_expectation_open_loop_series( "BP messages; pass require_fixed_point=False for an exploratory " "estimate" ) + if internal_info is not None: + internal_info.update( + { + "bp_converged": bp_info.get("converged"), + "bp_iterations": bp_info.get("iterations"), + "bp_max_mdiff": bp_info.get("max_mdiff"), + } + ) cache = cache or OpenLoopSeriesCache() cached_diagnostic_supports = {} @@ -6853,11 +7604,14 @@ def compute_local_expectation_open_loop_series( cluster_size=cluster_value, corridor_options=corridor_options, max_terms=max_terms, + max_loop_terms=max_loop_terms, max_enumeration_time=max_enumeration_time, max_enumeration_memory=max_enumeration_memory, max_flops_log10=max_flops_log10, max_peak_memory_log2=max_peak_memory_log2, path_edge_weights=path_edge_weights, + optimize=optimize, + contract_opts=contract_opts, ) cached = ( None @@ -6884,6 +7638,7 @@ def compute_local_expectation_open_loop_series( max_flops_log10=diagnostic_flops, max_peak_memory_log2=diagnostic_peak, max_terms=max_terms, + max_loop_terms=max_loop_terms, max_enumeration_time=max_enumeration_time, max_enumeration_memory=max_enumeration_memory, corridor_options=corridor_options, @@ -6894,6 +7649,7 @@ def compute_local_expectation_open_loop_series( "mode": mode, "edge_cutoff": edge_value, "cluster_size": cluster_value, + "max_loop_terms": max_loop_terms, "corridor_options": corridor_options, "cache_key": key, } @@ -6907,13 +7663,13 @@ def compute_local_expectation_open_loop_series( ) term_info = ( {} - if info is None - else info.setdefault("open_scalar_normalization_by_term", {}) + if internal_info is None + else internal_info.setdefault("open_scalar_normalization_by_term", {}) ) support_info = ( {} - if info is None - else info.setdefault("open_scalar_supports", {}) + if internal_info is None + else internal_info.setdefault("open_scalar_supports", {}) ) expecs = {} for where, gate in records: @@ -6935,6 +7691,7 @@ def compute_local_expectation_open_loop_series( edge_cutoff=edge_cutoff, cluster_size=cluster_size, max_terms=max_terms, + max_loop_terms=max_loop_terms, max_enumeration_time=max_enumeration_time, max_enumeration_memory=max_enumeration_memory, corridor_width=corridor_width, @@ -6951,7 +7708,8 @@ def compute_local_expectation_open_loop_series( diagnostic_support=diagnostic_support, contract_opts=contract_opts, cache=cache, - info=info, + info=internal_info, + on_budget=on_budget, ) result_key = where try: @@ -6960,51 +7718,51 @@ def compute_local_expectation_open_loop_series( result_key = tuple(sites) term_info[result_key] = normalization expecs[result_key] = value - if info is not None: + if internal_info is not None: support_key = tuple(sites) support_edge_costs = { key: value - for key, value in info["open_scalar_edge_term_costs"].items() + for key, value in internal_info["open_scalar_edge_term_costs"].items() if key[0] == support_key } support_edge_skipped = { key: value - for key, value in info[ + for key, value in internal_info[ "open_scalar_edge_skipped_terms" ].items() if key[0] == support_key } support_cluster_costs = { key: value - for key, value in info[ + for key, value in internal_info[ "open_scalar_cluster_region_costs" ].items() if key[0] == support_key } support_cluster_skipped = { key: value - for key, value in info[ + for key, value in internal_info[ "open_scalar_cluster_region_skipped_terms" ].items() if key[0] == support_key } support_info[support_key] = { - "terms": tuple(info["open_scalar_terms"]), + "terms": tuple(internal_info["open_scalar_terms"]), "requested_terms": tuple( - info["open_scalar_requested_terms"] + internal_info["open_scalar_requested_terms"] ), "skipped_terms": dict( - info["open_scalar_skipped_terms"] + internal_info["open_scalar_skipped_terms"] ), - "term_costs": dict(info["open_scalar_term_costs"]), + "term_costs": dict(internal_info["open_scalar_term_costs"]), "edge_skipped_terms": support_edge_skipped, "edge_term_costs": support_edge_costs, "cluster_region_skipped_terms": support_cluster_skipped, "cluster_region_costs": support_cluster_costs, - "family_counts": dict(info["open_scalar_family_counts"]), - "family_weights": dict(info["open_scalar_family_weights"]), + "family_counts": dict(internal_info["open_scalar_family_counts"]), + "family_weights": dict(internal_info["open_scalar_family_weights"]), "corridor": dict( - info.get("open_scalar_corridor", {}) + internal_info.get("open_scalar_corridor", {}) ), "diagnostic": ( None @@ -7013,5 +7771,64 @@ def compute_local_expectation_open_loop_series( ), } if return_all: - return expecs - return functools.reduce(operator.add, expecs.values()) + if not return_result: + if internal_info is not None: + internal_info["open_scalar_resources"] = monitor.finish() + return expecs + value = expecs + else: + value = functools.reduce(operator.add, expecs.values()) + resources = monitor.finish() + if internal_info is not None: + internal_info["open_scalar_resources"] = resources + effective_diagnostic = diagnostic + if effective_diagnostic is None and cached_diagnostic_supports: + effective_diagnostic = OpenLoopSeriesDiagnostic( + supports=dict(cached_diagnostic_supports), + bp_converged=bp_info.get("converged"), + bp_iterations=bp_info.get("iterations"), + bp_max_mdiff=bp_info.get("max_mdiff"), + ) + route_values = { + record.get("route") + for record in ( + effective_diagnostic.supports.values() + if effective_diagnostic + else () + ) + if record.get("route") is not None + } + if not route_values and internal_info is not None: + route_values = { + internal_info.get("open_scalar_route") + } - {None} + omitted = () if internal_info is None else tuple( + internal_info.get("open_scalar_omitted_terms", ()) + ) + complete = True if internal_info is None else all( + bool(value) + for key, value in internal_info.items() + if key.startswith("open_scalar_") and key.endswith("_complete") + ) + return OpenLoopMeasurementResult( + value=value, + normalization=( + None + if internal_info is None + else internal_info.get("open_scalar_denominator") + ), + info=internal_info or {}, + diagnostic=effective_diagnostic, + complete=complete, + approximate=( + bool(internal_info and internal_info.get("open_scalar_mode") == "corridor") + or bool(route_values.intersection({"graded_cluster_compatible"})) + ), + route=(next(iter(route_values)) if len(route_values) == 1 else "mixed"), + omitted_terms=omitted, + resources=resources, + bp_converged=bp_info.get("converged"), + bp_iterations=bp_info.get("iterations"), + bp_max_mdiff=bp_info.get("max_mdiff"), + bp=bp, + ) if return_result else value diff --git a/tests/test_bp_open_series.py b/tests/test_bp_open_series.py index 9f6a19b..925e6f1 100644 --- a/tests/test_bp_open_series.py +++ b/tests/test_bp_open_series.py @@ -9,17 +9,25 @@ from pepsy.bp import ( OpenLoopEnumerationLimitError, + OpenLoopBudgetError, + OpenLoopMeasurementResult, OpenLoopObservableTerm, OpenLoopSeriesCache, OpenLoopSeriesDiagnosticCache, OpenLoopSeriesSweepResult, compute_local_expectation_open_loop_series, + adaptive_open_loop_series, + diagnose_open_rho_series, diagnose_open_loop_series, partial_trace_open_loop_series_expand, partial_trace_open_loop_series_sweep, two_norm_bp, ) -from pepsy.bp.series import _open_term_family +from pepsy.bp.series import ( + _discover_grid_corridor_paths, + _grid_corridor_context, + _open_term_family, +) def _edge_degrees(tn, edges): @@ -235,6 +243,28 @@ def test_corridor_mode_limits_geometry_before_contraction(): ) +def test_loop_term_budget_is_separate_from_total_term_budget(): + state = qtn.PEPS.rand( + 3, + 3, + bond_dim=2, + phys_dim=2, + seed=1938, + dtype="complex128", + ) + with pytest.raises(OpenLoopEnumerationLimitError, match="max_loop_terms"): + partial_trace_open_loop_series_expand( + state, + ((0, 0), (0, 2)), + edge_cutoff=6, + max_terms=100, + max_loop_terms=0, + max_iterations=200, + tol=1e-10, + diis=False, + ) + + def test_corridor_mode_can_use_compressed_boundary_contraction(): state = qtn.PEPS.rand( 1, @@ -317,6 +347,7 @@ def test_open_measurement_diagnostic_selects_auto_route_and_reuses_terms(): assert np.isfinite(value) assert info["open_scalar_mode"] == "auto" assert tuple(info["open_scalar_requested_terms"]) == support["terms"] + assert info["open_scalar_region_path_cache"] cached = diagnose_open_loop_series( state, @@ -605,6 +636,153 @@ def test_open_series_reuses_contraction_paths_for_shared_regions(): ) +def test_pbc_corridor_keeps_parallel_period_two_bonds_as_distinct_paths(): + """A 3x2 torus is a multigraph, not a simple coordinate graph.""" + state = qtn.PEPS.rand( + 3, + 2, + bond_dim=2, + phys_dim=2, + cyclic=(True, True), + seed=1948, + ) + context = _grid_corridor_context(state) + neighbors = list(context["neighbors"]((0, 0))) + seam_edges = { + edge for neighbor, _, edge in neighbors if neighbor == (0, 1) + } + assert len(seam_edges) == 2 + + paths, corridor_edges, diagnostics = _discover_grid_corridor_paths( + state, + ((0, 0), (2, 1)), + corridor_width=0, + max_path_candidates=20, + ) + assert diagnostics["path_count"] == 4 + assert len({path.edges for path in paths}) == 4 + assert len(corridor_edges) == len(set(corridor_edges)) + assert all(len(path.edges) == 2 for path in paths) + + +def test_open_series_production_result_reports_budget_and_resources(): + state = qtn.PEPS.rand( + 1, + 4, + bond_dim=2, + phys_dim=2, + seed=1949, + dtype="complex128", + ) + result = compute_local_expectation_open_loop_series( + state, + {((0, 0), (0, 3)): np.eye(4)}, + edge_cutoff=3, + max_flops_log10=2.0, + max_iterations=200, + tol=1e-10, + diis=False, + return_result=True, + measure_resources=True, + ) + assert isinstance(result, OpenLoopMeasurementResult) + assert not result.complete + assert result.omitted_terms + assert result.resources["enabled"] + assert result.info["open_scalar_complete"] is False + with pytest.raises(OpenLoopBudgetError): + compute_local_expectation_open_loop_series( + state, + {((0, 0), (0, 3)): np.eye(4)}, + edge_cutoff=3, + max_flops_log10=2.0, + max_iterations=200, + tol=1e-10, + diis=False, + on_budget="raise", + ) + + +def test_rho_diagnostic_and_adaptive_corridor_ladder(): + state = qtn.PEPS.rand( + 1, + 4, + bond_dim=2, + phys_dim=2, + seed=1950, + dtype="complex128", + ) + support = ((0, 0), (0, 3)) + diagnostic = diagnose_open_rho_series( + state, + (support,), + edge_cutoff=3, + max_iterations=200, + tol=1e-10, + diis=False, + ) + record = diagnostic.for_support(support) + assert record["observable_kind"] == "rho" + assert record["output_shape"] == (4, 4) + assert record["logical_output_shape"] == (2, 2, 2, 2) + assert record["output_memory_bytes"] > 0 + + rho_result = partial_trace_open_loop_series_expand( + state, + support, + edge_cutoff=1, + max_iterations=200, + tol=1e-10, + diis=False, + measure_resources=True, + return_result=True, + ) + assert isinstance(rho_result, OpenLoopMeasurementResult) + assert rho_result.value.shape == (4, 4) + assert rho_result.normalization is not None + assert rho_result.resources["enabled"] + + adaptive = adaptive_open_loop_series( + state, + {support: np.eye(4)}, + corridor_widths=(0, 1), + edge_cutoff=3, + max_iterations=200, + tol=1e-10, + diis=False, + min_stable=1, + ) + assert adaptive.values + assert adaptive.settings[0]["corridor_width"] == 0 + assert adaptive.diagnostics[0] is not None + assert adaptive.bp is not None + + +def test_open_series_public_controls_are_validated_and_cache_is_positional_safe(): + cache = OpenLoopSeriesDiagnosticCache({}) + assert cache.diagnostics_by_key == {} + + with pytest.raises(ValueError, match="positive integer"): + diagnose_open_rho_series( + None, + (((0, 0),),), + max_rho_identity_dimension=0, + ) + with pytest.raises(ValueError, match="non-negative"): + adaptive_open_loop_series( + None, + {}, + corridor_widths=(-1,), + ) + with pytest.raises(ValueError, match="positive integer"): + adaptive_open_loop_series( + None, + {}, + corridor_widths=(0,), + min_stable=0, + ) + + def test_open_rho_series_reuses_one_d2bp_message_set(): state = qtn.PEPS.rand( 1, From 051c600f88619a26030499579358e2f8675dd1b9 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Fri, 31 Jul 2026 11:41:00 -0700 Subject: [PATCH 42/70] make ternary virtual root the tree default --- docs/api/optimizers/tree.md | 43 ++++---- src/pepsy/optimizers/planning.py | 2 +- src/pepsy/optimizers/tree/layout.py | 101 +++++++++++------- src/pepsy/optimizers/tree/optimizer.py | 80 ++++++++------ src/pepsy/optimizers/tree/ttn.py | 19 ++-- .../optimizers/tree_stabilizer/optimizer.py | 17 ++- src/pepsy/tensors/constructors.py | 41 +++++-- tests/test_optimize_tree.py | 74 +++++++++---- tests/test_optimize_tree_stabilizer.py | 12 ++- 9 files changed, 255 insertions(+), 134 deletions(-) diff --git a/docs/api/optimizers/tree.md b/docs/api/optimizers/tree.md index 340d36f..d2f72db 100644 --- a/docs/api/optimizers/tree.md +++ b/docs/api/optimizers/tree.md @@ -8,10 +8,9 @@ Huang, Mendl; Quantum 7, 964, 2023; [arXiv:2206.01000](https://arxiv.org/abs/220 By default the state is stored with one leaf tensor per qubit. A plan may instead designate one `root_qubit`, placing that physical index directly on the top tensor while every other qubit remains a leaf. Internal nodes may have -**any arity** -- by default the layout finder *searches* a small set of -candidate arities `(2, 3, 4)` and keeps the objective-best plan, but a fixed -binary tree, flatter `k`-ary trees, or gate-connectivity-driven communities -(see *Tree structure*) all work through the same machinery. +**any arity** -- the default is a binary tree below a three-virtual-leg root, +but a fixed binary root, flatter `k`-ary trees, or gate-connectivity-driven +communities (see *Tree structure*) all work through the same machinery. For example, this constructs a binary detector tree whose logical qubit is the open top index. The root tensor has two virtual child bonds and physical index @@ -47,9 +46,11 @@ structured sub-MPOs may include it in their support. `TreeLayoutFinder` keeps the site fixed at the root while its path, Steiner, congestion, greedy, and Nevergrad objectives permute only the remaining leaf sites. -For the conventional binary TTN with a three-leg top tensor, pass -`max_arity=2, top_arity=3` to `TreePlan.from_order`, `TreeLayoutFinder`, or -`TreeTensorNetwork.from_order`. The structural root then has three **virtual** +The conventional binary TTN with a three-leg top tensor is the default when +there are at least three leaves and no `root_qubit`. Pass +`max_arity=2, top_arity=3` explicitly to `TreePlan.from_order`, +`TreeLayoutFinder`, or `TreeTensorNetwork.from_order` for the same geometry. +The structural root then has three **virtual** child bonds and no parent bond; every non-root internal tensor has two child bonds and one parent bond. Thus the root is still in the rank-three binary class, rather than being a genuinely wider tensor. `top_arity=3` is not @@ -316,8 +317,8 @@ Because the geometry (`plan`) and naming live in `_EXTRA_PROPS`, they survive `.copy()` and every Quimb view, exactly like `site_ind_id` does for an MPS. Build one with `TreeTensorNetwork.from_plan(plan)` (product `|0...0>`), `TreeTensorNetwork.from_order(order, structure=...)` (build the plan and the -product state in one step; its `top_arity=3` option exposes the ternary virtual -root), or `TreeTensorNetwork.rand(plan, D=..., seed=...)` +product state in one step; its default exposes the ternary virtual root), or +`TreeTensorNetwork.rand(plan, D=..., seed=...)` (a random state, canonicalised around the root by default). `TreeOptimizer` builds and evolves its state on this class, delegating all node/qubit naming and geometry queries to it. @@ -481,8 +482,8 @@ any arity, controlled by two knobs on `TreeLayoutFinder` / `TreePlan.from_order` single fixed tree: `2` reproduces the strictly-binary tree exactly, larger values give flatter `k`-ary trees with shorter geodesics, `None` leaves the arity unbounded) or an iterable of candidate arities to **search**. The - default `(2, 3, 4)` searches those three and keeps the objective-best plan; - pass `max_arity=2` to force a fixed binary tree. + default `2` selects the fixed binary tree; pass an iterable such as + `(2, 3, 4)` to search candidate arities explicitly. - `structure="adaptive"` reads the gate-stream interaction graph and lets each level branch into as many children as it has strongly coupled communities (edges above `community_frac` times the level's strongest edge). A densely @@ -503,11 +504,10 @@ three-virtual-bond root convention described above. `TreePlan.max_arity()` and `TreePlan.is_binary()` report the shape; `TreePlan.is_strictly_binary()` is the strict two-child-at-every-internal-node predicate. -For an automatic arity choice, call -`finder.recommend_arities((2, 3, 4))`. This is also what the finder and -`TreeOptimizer` do **by default** (their `max_arity` defaults to `(2, 3, 4)`), -so `TreeOptimizer(gate_stream, n=n, chi=chi)` already searches these arities -- -and does so `chi`-aware, since the optimizer forwards its own `chi` (see below). +For an automatic arity search, call `finder.recommend_arities((2, 3, 4))` +explicitly. The default `TreeOptimizer(gate_stream, n=n, chi=chi)` uses the +fixed binary/ternary-root geometry; it does not allocate tensors or perform +truncations while finding the layout. The result contains the recommended `TreePlan` plus per-candidate path, edge-load, peak-bond-growth, and local virtual-degree summaries. An explicit handoff looks like: @@ -641,12 +641,11 @@ studies rather than routine short simulations. Candidate records expose their initial/final leaf order and the greedy/Nevergrad diagnostics under `candidate["planning"]`. -The default search is made `chi`-aware automatically when a `chi` is available: -`TreeLayoutFinder(gate_stream, n=n, chi=chi)` biases its default `(2, 3, 4)` -search toward `chi`-exact structures, and `TreeOptimizer(gate_stream, n=n, -chi=chi)` forwards its own `chi` into the finder it builds -- so the everyday -`TreeOptimizer(gate_stream, n=n, chi=chi)` already prefers a tree that stays -exact at `chi`. A bare finder with no `chi` searches `chi`-blind. +When an explicit iterable of arity candidates is supplied, a `chi` biases the +search toward `chi`-exact structures. The default fixed binary/ternary-root +geometry is independent of `chi`; layout finding still does not allocate +tensors or perform truncations. A bare finder with no `chi` is likewise +static unless candidate search is explicitly requested. Set `max_operator_qubits` to bound dense rank diagnostics and operator allocation; wider native MPO events can still replay without dense materialization. `TreeLayoutFinder(..., max_operator_qubits=...)` uses a diff --git a/src/pepsy/optimizers/planning.py b/src/pepsy/optimizers/planning.py index deb95aa..ced285b 100644 --- a/src/pepsy/optimizers/planning.py +++ b/src/pepsy/optimizers/planning.py @@ -413,7 +413,7 @@ def _mps_layout(self, records): def _tree_layout(self, records): kwargs = { "structure": "quality", - "max_arity": (2, 3, 4), + "max_arity": 2, "objective": "path", **self.tree_layout_kwargs, } diff --git a/src/pepsy/optimizers/tree/layout.py b/src/pepsy/optimizers/tree/layout.py index bb563e1..a6b178f 100644 --- a/src/pepsy/optimizers/tree/layout.py +++ b/src/pepsy/optimizers/tree/layout.py @@ -17,9 +17,9 @@ have any arity: ``max_arity`` gives flatter ``k``-ary trees (shallower geodesics), while ``structure="adaptive"`` reads the gate-stream interaction graph and lets each level branch into as many children as it has strongly -coupled communities. By default the finder *searches* a small set of -candidate arities (``max_arity=(2, 3, 4)``) and keeps the objective-best plan; -pass a scalar ``max_arity=2`` to opt back into a single fixed binary tree. +coupled communities. The default is a binary tree below a three-virtual-leg +root, which keeps every tensor at rank three. Pass an explicit ``top_arity`` +or an iterable ``max_arity`` to request another geometry. """ from __future__ import annotations @@ -56,6 +56,7 @@ __all__ = ["TreePlan", "TreeLayoutFinder"] _DEFAULT_MAX_ARITY = object() +_DEFAULT_TOP_ARITY = object() _DEFAULT_CHI = object() _DEFAULT_ORDER = object() _DEFAULT_SEARCH_OPTION = object() @@ -402,8 +403,8 @@ def _normalize_arity_candidates(max_arity): """Return ``(representative_arity, candidates)`` from a ``max_arity`` arg. ``max_arity`` may be a single int (a fixed arity), ``None`` (unbounded), or - an iterable of candidate arities to *search* (the finder default - ``(2, 3, 4)``). ``candidates`` is ``None`` unless a search set was given; + an iterable of candidate arities to *search*. ``candidates`` is ``None`` + unless a search set was given; the representative single arity is what the legacy single-plan builders use and is the first concrete candidate. """ @@ -471,8 +472,8 @@ class TreePlan: Nodes are integer ids. Leaves map one-to-one to qubits. Optionally, one additional qubit can be carried by the structural root via ``root_qubit``; this gives a binary top tensor two child bonds plus one open physical leg. - Other internal nodes carry no physical qubit. A strictly-binary tree (every - internal node with two children) is the common default, but the structure + Other internal nodes carry no physical qubit. The common default is binary + below a ternary virtual root, but the structure supports arbitrary arity so a level can branch into as many subtrees as the gate stream suggests. The plan is a pure structure description: it carries no tensor data and is consumed by @@ -510,7 +511,8 @@ def __init__( @classmethod def from_order(cls, order, *, weights=None, structure="quality", max_arity=2, community_frac=0.35, star_frac=0.75, - dense_max=512, root_qubit=None, top_arity=None): + dense_max=512, root_qubit=None, + top_arity=_DEFAULT_TOP_ARITY): """Build a rooted tree by recursive partition of ``order``. Parameters @@ -549,14 +551,13 @@ def from_order(cls, order, *, weights=None, structure="quality", Maximum subsystem size for dense spectral reordering. root_qubit : int, optional Qubit label carried by the top tensor rather than a leaf. - top_arity : int, optional - Number of virtual child bonds on the structural root. Set - ``top_arity=3`` with ``max_arity=2`` for the conventional binary - TTN with a ternary top tensor: the root has three virtual legs and - every non-root internal tensor has two child legs plus one parent - leg. This keeps every tensor rank at most three. It is incompatible - with ``root_qubit`` when greater than two because that would make a - rank-four root tensor. + top_arity : int or None, optional + Number of virtual child bonds on the structural root. By default, + ``max_arity=2`` uses ``top_arity=3`` when there are at least three + leaf qubits and no ``root_qubit``. Set ``top_arity=None`` or + ``top_arity=2`` to use the ordinary binary root. A value greater + than two is incompatible with ``root_qubit`` because that would + make a rank-four root tensor. """ order = list(order) if not order and root_qubit is None: @@ -570,6 +571,16 @@ def from_order(cls, order, *, weights=None, structure="quality", root_qubit = int(root_qubit) except (TypeError, ValueError) as exc: raise ValueError("root_qubit must be an integer or None.") from exc + if max_arity is not None: + max_arity = int(max_arity) + if max_arity < 2: + raise ValueError("max_arity must be >= 2 (or None).") + if top_arity is _DEFAULT_TOP_ARITY: + top_arity = ( + 3 + if root_qubit is None and max_arity == 2 and len(order) >= 3 + else None + ) if top_arity is not None: try: top_arity = int(top_arity) @@ -598,10 +609,6 @@ def from_order(cls, order, *, weights=None, structure="quality", raise ValueError( "structure must be 'quality', 'balanced', or 'adaptive'." ) - if max_arity is not None: - max_arity = int(max_arity) - if max_arity < 2: - raise ValueError("max_arity must be >= 2 (or None).") counter = [0] children = {} parent = {} @@ -1283,27 +1290,29 @@ class TreeLayoutFinder: Explicit interaction supports, used instead of extracting them from ``gates``. structure : {"quality", "balanced", "adaptive"} - Partition strategy passed to :meth:`TreePlan.from_order`. ``"quality"`` - and ``"balanced"`` build strictly-binary trees when ``max_arity=2``; + Partition strategy passed to :meth:`TreePlan.from_order`. ``"quality"`` + and ``"balanced"`` build binary trees below the optional ternary root + when ``max_arity=2``; ``"adaptive"`` lets each level branch into its strongly coupled communities so the arity follows the gate connectivity. max_arity : int, None, or iterable of ints Maximum children per internal node. A scalar builds one fixed tree (``2`` gives the binary tree; larger values or ``None`` give flatter / wider trees). An iterable of candidate arities makes :meth:`run` *search* - them and keep the objective-best plan; this is the default - ``(2, 3, 4)``. Pass a scalar to opt back into a single fixed tree. - top_arity : int, optional + them and keep the objective-best plan. The default is the scalar + ``2``. + top_arity : int or None, optional Override the structural root's number of virtual child bonds. With - ``max_arity=2, top_arity=3`` the finder builds the conventional binary - TTN whose top tensor has three virtual legs while all non-root internal - tensors remain two-in/one-out. This keeps the maximum tensor rank at - three. It cannot be combined with ``root_qubit`` when greater than two. + the default ``max_arity=2``, omitted ``top_arity`` selects ``3`` when + possible, so the top tensor has three virtual legs while all non-root + internal tensors remain two-in/one-out. Set ``top_arity=None`` or + ``top_arity=2`` to opt out. It cannot be greater than two with + ``root_qubit``. chi : int, optional - Bond-dimension budget used to bias the default arity search toward plans - that stay exact at ``chi`` (see :meth:`recommend_arities`). ``None`` - keeps the search purely objective-driven and is the static layout - default; it does not allocate tensors or perform truncations. + Bond-dimension budget used when an explicit iterable of arities is + searched to prefer plans that stay exact at ``chi`` (see + :meth:`recommend_arities`). The fixed default geometry does not + allocate tensors or perform truncations. :class:`TreeOptimizer` forwards its own ``chi`` here automatically. community_frac : float Strong-edge fraction for ``structure="adaptive"`` (see @@ -1381,7 +1390,7 @@ class TreeLayoutFinder: """ def __init__(self, gates=None, n=None, *, supports=None, structure="quality", - max_arity=(2, 3, 4), top_arity=None, + max_arity=2, top_arity=_DEFAULT_TOP_ARITY, community_frac=0.35, star_frac=0.75, dense_max=512, objective="path", weight_mode="count", chi=None, max_operator_qubits=8, hybrid_weights=None, refine=None, @@ -1458,6 +1467,19 @@ def __init__(self, gates=None, n=None, *, supports=None, structure="quality", self.leaf_qubits = tuple( q for q in range(self.n) if q != self.root_qubit ) + self.max_arity, self.arity_candidates = _normalize_arity_candidates( + max_arity + ) + if top_arity is _DEFAULT_TOP_ARITY: + top_arity = ( + 3 + if ( + root_qubit is None + and self.max_arity == 2 + and len(self.leaf_qubits) >= 3 + ) + else None + ) if top_arity is not None: try: top_arity = int(top_arity) @@ -1479,9 +1501,6 @@ def __init__(self, gates=None, n=None, *, supports=None, structure="quality", self.top_arity = top_arity self.supports = tuple(normalized_supports) self.structure = structure - self.max_arity, self.arity_candidates = _normalize_arity_candidates( - max_arity - ) self.chi = _validate_chi(chi) self.community_frac = float(community_frac) self.star_frac = float(star_frac) @@ -3028,11 +3047,11 @@ def run( ): """Return a TreePlan for the selected layout objective. - When the finder was built with a set of candidate arities (the default - ``max_arity=(2, 3, 4)``), this searches them with + A scalar ``max_arity`` (the default ``2``) builds one fixed binary + plan with the default ternary virtual root. When the finder is built + with an iterable of candidate arities, this searches them with :meth:`recommend_arities` -- ``chi``-aware when the finder carries a - ``chi`` -- and returns the objective-best plan. A scalar ``max_arity`` - builds one fixed plan. + ``chi`` -- and returns the objective-best plan. ``chi`` and the fixed-plan ``refine`` / ``search`` controls can be overridden for this call. Pass ``progbar=True`` to display greedy and diff --git a/src/pepsy/optimizers/tree/optimizer.py b/src/pepsy/optimizers/tree/optimizer.py index 9fcf81b..52858e5 100644 --- a/src/pepsy/optimizers/tree/optimizer.py +++ b/src/pepsy/optimizers/tree/optimizer.py @@ -7,7 +7,7 @@ A quantum state is stored as a rooted tree tensor network whose leaves carry physical qubit indices. One optional physical qubit may instead live on the root tensor. Internal nodes may have any arity: the default structure is a -strictly-binary tree, but flatter ``k``-ary trees +binary tree below a ternary virtual root, but flatter ``k``-ary trees (``max_arity``) or gate-connectivity-driven communities (``structure="adaptive"``) are supported unchanged. A bundled gate stream ``[(gate, where), ...]`` is replayed: @@ -61,6 +61,7 @@ submpo_event_parts, ) from .layout import ( + _DEFAULT_TOP_ARITY, TreeLayoutFinder, TreePlan, _normalize_time_decay, @@ -320,17 +321,16 @@ class TreeOptimizer: max_arity : int, None, or iterable of ints Maximum children per internal node for the auto-built structure. A scalar builds one fixed tree (``2`` = binary; larger values or ``None`` - = flatter / wider). An iterable of candidate arities makes the finder - *search* them and keep the objective-best plan; this is the default - ``(2, 3, 4)`` and is made ``chi``-aware automatically using this - optimizer's ``chi`` (structures that stay exact at ``chi`` are - preferred). Pass ``max_arity=2`` to force a fixed binary tree. Ignored - when an explicit ``tree`` is supplied. - top_arity : int, optional + = flatter / wider). The default is the fixed binary tree and its + ternary virtual root; an iterable can still be supplied explicitly to + search candidate arities. Ignored when an explicit ``tree`` is + supplied. + top_arity : int or None, optional Number of virtual child bonds on the structural root when the layout - is built automatically. Set ``top_arity=3`` with ``max_arity=2`` for - the conventional binary TTN with a three-leg top tensor. This keeps - every tensor rank at most three. + is built automatically. Omitted with ``max_arity=2`` selects + ``top_arity=3`` when possible, giving the conventional binary TTN + with a three-leg top tensor. Pass ``top_arity=None`` or ``2`` to use a + binary root. This keeps every tensor rank at most three. layout_objective : {"path", "congestion", "compression", "hypergraph", "full_tree", "hybrid"} Objective used when building an automatic tree. ``"path"`` is the backward-compatible interaction-path heuristic; ``"congestion"`` @@ -449,7 +449,8 @@ def __init__(self, gates=None, n=None, *, chi=64, cutoff=_DEFAULT_CUTOFF, cutoff_mode=_DEFAULT_CUTOFF_MODE, mode="auto", two_site_mode=None, - structure="quality", max_arity=(2, 3, 4), top_arity=None, + structure="quality", max_arity=2, + top_arity=_DEFAULT_TOP_ARITY, community_frac=0.35, star_frac=0.75, layout_objective="path", layout_weight_mode="count", layout_time_decay=None, @@ -485,8 +486,6 @@ def __init__(self, gates=None, n=None, *, chi=64, "layout must be a TreeLayoutFinder or TreePlan; " "pass an entangled TreeTensorNetwork as state= or tn=." ) - if top_arity is None and layout_top_arity is not None: - top_arity = layout_top_arity if state is not None: if tn is not None: raise ValueError("pass either state= or tn=, not both.") @@ -503,20 +502,6 @@ def __init__(self, gates=None, n=None, *, chi=64, raise ValueError( "root_qubit must be an integer or None." ) from exc - if top_arity is not None: - try: - top_arity = int(top_arity) - except (TypeError, ValueError) as exc: - raise ValueError( - "top_arity must be an integer >= 2 or None." - ) from exc - if top_arity < 2: - raise ValueError("top_arity must be >= 2 or None.") - if root_qubit is not None and top_arity != 2: - raise ValueError( - "top_arity > 2 cannot be combined with root_qubit: " - "the root would have a rank-four tensor." - ) self.G, self.where, self.event_types = self._normalize_gate_queue(gates) self.layout_finder = layout if isinstance(layout, TreeLayoutFinder) else None @@ -587,6 +572,37 @@ def __init__(self, gates=None, n=None, *, chi=64, raise ValueError( f"root_qubit {root_qubit!r} is outside 0..{self.n - 1}." ) + if top_arity is _DEFAULT_TOP_ARITY: + if layout_top_arity is not None: + top_arity = layout_top_arity if layout_top_arity >= 2 else None + elif isinstance(tree, TreePlan): + tree_top_arity = tree.top_arity + top_arity = tree_top_arity if tree_top_arity >= 2 else None + elif ( + root_qubit is None + and isinstance(max_arity, Integral) + and int(max_arity) == 2 + and self.n >= 3 + ): + top_arity = 3 + else: + top_arity = None + elif top_arity is None and layout_top_arity is not None: + top_arity = layout_top_arity + if top_arity is not None: + try: + top_arity = int(top_arity) + except (TypeError, ValueError) as exc: + raise ValueError( + "top_arity must be an integer >= 2 or None." + ) from exc + if top_arity < 2: + raise ValueError("top_arity must be >= 2 or None.") + if root_qubit is not None and top_arity != 2: + raise ValueError( + "top_arity > 2 cannot be combined with root_qubit: " + "the root would have a rank-four tensor." + ) # The TTN itself always uses compact physical positions. This facade # optionally preserves caller-facing logical labels across a cap while # keeping Quimb's internal site/index space contiguous. @@ -4914,11 +4930,11 @@ def get_projection_diagnostics(self): @classmethod def find_tree_layout(cls, gates, n=None, *, structure="quality", - max_arity=(2, 3, 4), community_frac=0.35, + max_arity=2, community_frac=0.35, star_frac=0.75, layout_objective="path", layout_weight_mode="count", layout_time_decay=None, layout_time_window=None, - root_qubit=None, top_arity=None, + root_qubit=None, top_arity=_DEFAULT_TOP_ARITY, max_operator_qubits=_DEFAULT_MAX_OPERATOR_QUBITS): """Return the :class:`TreePlan` a :class:`TreeLayoutFinder` would use.""" return TreeLayoutFinder( @@ -4935,9 +4951,9 @@ def find_tree_layout(cls, gates, n=None, *, structure="quality", @classmethod def convergence_sweep(cls, gates, n=None, chi_values=(2, 4, 8, 16, 32), *, - ops=None, structure="quality", max_arity=(2, 3, 4), + ops=None, structure="quality", max_arity=2, community_frac=0.35, star_frac=0.75, tree=None, - root_qubit=None, top_arity=None, + root_qubit=None, top_arity=_DEFAULT_TOP_ARITY, dense_cap=1 << 14): """Replay ``gates`` at several ``chi`` and report convergence. diff --git a/src/pepsy/optimizers/tree/ttn.py b/src/pepsy/optimizers/tree/ttn.py index ef2c34a..3e796d4 100644 --- a/src/pepsy/optimizers/tree/ttn.py +++ b/src/pepsy/optimizers/tree/ttn.py @@ -19,9 +19,10 @@ ``site_tag_id.format(q)`` (default ``"I{}"``) and the physical index ``site_ind_id.format(q)`` (default ``"k{}"``) for qubit ``q``; these are structural leaves by default; -* a plan may use ``top_arity=3`` for the conventional binary TTN with three - virtual bonds entering the top tensor. Other internal nodes then have two - child bonds plus one parent bond, so every tensor remains rank three; +* the default auto-built plan uses ``top_arity=3`` for the conventional binary + TTN with three virtual bonds entering the top tensor. Other internal nodes + then have two child bonds plus one parent bond, so every tensor remains rank + three; explicit plans can use another root arity; * a plan may designate one additional ``root_qubit`` carried by the top tensor. A binary root then has exactly two child bonds plus this physical leg. Other internal nodes remain ancillary bond carriers. This class supplies the @@ -45,7 +46,7 @@ from quimb.tensor.tensor_core import TensorNetwork from numbers import Integral -from .layout import TreePlan +from .layout import TreePlan, _DEFAULT_TOP_ARITY __all__ = ["TreeTensorNetwork"] @@ -1940,16 +1941,16 @@ def zero_charge(value): def from_order(cls, order, *, weights=None, structure="quality", max_arity=2, community_frac=0.35, star_frac=0.75, dtype=complex, site_tag_id="I{}", site_ind_id="k{}", - node_tag_id="N{}", root_qubit=None, top_arity=None): + node_tag_id="N{}", root_qubit=None, + top_arity=_DEFAULT_TOP_ARITY): """Build a product state on a tree partitioned from ``order``. Convenience wrapper that first builds a :class:`TreePlan` with :meth:`TreePlan.from_order` and then :meth:`from_plan`. ``max_arity`` and ``structure`` control the tree shape (see - :meth:`TreePlan.from_order`); the defaults reproduce the binary tree. - Set ``top_arity=3`` with ``max_arity=2`` for a binary TTN with a - three-virtual-leg top tensor; all lower internal tensors remain rank - three (two child bonds plus one parent bond). + :meth:`TreePlan.from_order`). The default is a binary tree below a + three-virtual-leg top tensor when there are at least three leaves; + pass ``top_arity=None`` or ``top_arity=2`` to use a binary root. """ plan = TreePlan.from_order( order, weights=weights, structure=structure, diff --git a/src/pepsy/optimizers/tree_stabilizer/optimizer.py b/src/pepsy/optimizers/tree_stabilizer/optimizer.py index 2fa6b6c..8e254a3 100644 --- a/src/pepsy/optimizers/tree_stabilizer/optimizer.py +++ b/src/pepsy/optimizers/tree_stabilizer/optimizer.py @@ -39,7 +39,7 @@ from ..stabilizer_tn.settings import DEFAULT_MAX_PAULI_DECOMPOSITION_QUBITS from ..stabilizer_tn.stn_state import _CLIFFORD_GATES, _validate_bits from ..mps.optimizer import conditional_event_parts, submpo_event_parts -from ..tree.layout import TreeLayoutFinder, TreePlan +from ..tree.layout import TreeLayoutFinder, TreePlan, _DEFAULT_TOP_ARITY from ..tree.optimizer import ( TreeOptimizer, _DEFAULT_CUTOFF, @@ -630,7 +630,8 @@ def __init__( tree=None, layout=None, structure="quality", - max_arity=(2, 3, 4), + max_arity=2, + top_arity=_DEFAULT_TOP_ARITY, layout_objective="path", layout_weight_mode="count", mode="auto", @@ -792,6 +793,7 @@ def __init__( cutoff=cutoff, structure=structure, max_arity=max_arity, + top_arity=top_arity, layout_objective=layout_objective, layout_weight_mode=layout_weight_mode, max_operator_qubits=max_operator_qubits, @@ -808,6 +810,7 @@ def __init__( finder_kwargs = dict(layout_kwargs) finder_kwargs.setdefault("structure", structure) finder_kwargs.setdefault("max_arity", max_arity) + finder_kwargs.setdefault("top_arity", top_arity) finder_kwargs.setdefault("objective", layout_objective) finder_kwargs.setdefault("weight_mode", layout_weight_mode) finder_kwargs.setdefault("chi", chi) @@ -828,6 +831,7 @@ def __init__( mode=mode, structure=structure, max_arity=max_arity, + top_arity=top_arity, layout_objective=layout_objective, layout_weight_mode=layout_weight_mode, tree=tree, @@ -1128,6 +1132,7 @@ def _build_frame_layout( cutoff, structure, max_arity, + top_arity, layout_objective, layout_weight_mode, max_operator_qubits, @@ -1161,6 +1166,7 @@ def _build_frame_layout( weight_mode = options.pop("weight_mode", layout_weight_mode) structure = options.pop("structure", structure) max_arity = options.pop("max_arity", max_arity) + top_arity = options.pop("top_arity", top_arity) objective = options.pop( "objective", options.pop("layout_objective", layout_objective) ) @@ -1196,6 +1202,7 @@ def _build_frame_layout( cutoff=cutoff, structure=structure, max_arity=max_arity, + top_arity=top_arity, layout_objective=objective, layout_weight_mode=weight_mode, max_operator_qubits=max_operator_qubits, @@ -1211,6 +1218,7 @@ def _build_frame_layout( n=int(n), structure=structure, max_arity=max_arity, + top_arity=top_arity, objective=objective, weight_mode=weight_mode, chi=chi, @@ -1417,6 +1425,7 @@ def current_frame_layout(self, *, weight_mode="count", **kwargs): finder_kwargs = { "structure": self._tree.structure, "max_arity": self._tree.max_arity, + "top_arity": self._tree.top_arity, "objective": self._tree.layout_objective, "weight_mode": weight_mode, "chi": self._tree.chi, @@ -1467,6 +1476,7 @@ def apply_frame_layout(self, plan="auto", *, layout_kwargs=None): mode=self._tree.mode, structure=self._tree.structure, max_arity=self._tree.max_arity, + top_arity=selected.top_arity, layout_objective=self._tree.layout_objective, layout_weight_mode=self._tree.layout_weight_mode, tree=selected, @@ -2691,7 +2701,8 @@ def _magic_tree_plan(cls, entries, n, ancillas, kwargs): layout_options.setdefault("weight_mode", "auto") finder_options = { "structure": kwargs.get("structure", "quality"), - "max_arity": kwargs.get("max_arity", (2, 3, 4)), + "max_arity": kwargs.get("max_arity", 2), + "top_arity": kwargs.get("top_arity", _DEFAULT_TOP_ARITY), "chi": kwargs.get("chi", 64), "max_operator_qubits": kwargs.get( "max_operator_qubits", DEFAULT_MAX_PAULI_DECOMPOSITION_QUBITS diff --git a/src/pepsy/tensors/constructors.py b/src/pepsy/tensors/constructors.py index 489b68f..7cf2a76 100644 --- a/src/pepsy/tensors/constructors.py +++ b/src/pepsy/tensors/constructors.py @@ -21,6 +21,24 @@ "hrs_to_peps", "hrs_to_mps", "hrps_to_peps", "hrps_to_mps", "hrps_to_ttn", ] + +_DEFAULT_TREE_TOP_ARITY = object() + + +def _resolve_tree_top_arity(top_arity, *, max_arity, n, root_qubit): + """Resolve the constructor default without hiding explicit opt-outs.""" + if top_arity is not _DEFAULT_TREE_TOP_ARITY: + return top_arity + if ( + root_qubit is None + and isinstance(max_arity, Integral) + and int(max_arity) == 2 + and int(n) >= 3 + ): + return 3 + return None + + def add_cycle(peps, bond_dim, cylinder=False): """Add periodic bonds to a PEPS network in x (and optional y) directions.""" Ly = peps.Ly @@ -619,7 +637,7 @@ def ps_to_ttn( root_qubit=None, structure="balanced", max_arity=2, - top_arity=None, + top_arity=_DEFAULT_TREE_TOP_ARITY, community_frac=0.35, star_frac=0.75, chi: int = 1, @@ -659,7 +677,9 @@ def ps_to_ttn( Qubit carried by the top tensor rather than a structural leaf. When an explicit ``tree`` is supplied, this must match its root site. structure, max_arity, top_arity, community_frac, star_frac - Forwarded to :meth:`TreePlan.from_order`. + Forwarded to :meth:`TreePlan.from_order`. The default is a binary + tree below a three-virtual-leg root when possible; pass + ``top_arity=None`` or ``top_arity=2`` to use a binary root. chi : int, optional If greater than one, expand every virtual bond to at least ``chi``. rand_strength : float, optional @@ -709,6 +729,9 @@ def ps_to_ttn( if tree is None: if root_qubit is not None: root_qubit = int(root_qubit) + top_arity = _resolve_tree_top_arity( + top_arity, max_arity=max_arity, n=n, root_qubit=root_qubit + ) if order is None: order = ( range(n) @@ -847,7 +870,7 @@ def hrs_to_ttn( root_qubit=None, structure="balanced", max_arity=2, - top_arity=None, + top_arity=_DEFAULT_TREE_TOP_ARITY, community_frac=0.35, star_frac=0.75, seed=None, @@ -865,10 +888,11 @@ def hrs_to_ttn( With ``fermion=`` the physical sites receive the model's charge sectors, while virtual-only internal nodes are neutral and every virtual tree edge is a conjugate pair of Symmray charge-sector indices. ``root_qubit`` places - one physical site on the top tensor. ``top_arity=3`` with ``max_arity=2`` - gives the conventional three-virtual-leg binary root. ``chi`` is the - requested total virtual-bond dimension. All block-sparse and fermionic - operations are delegated to Symmray/Quimb. + one physical site on the top tensor. The default gives a conventional + three-virtual-leg binary root when possible; ``top_arity=None`` or + ``top_arity=2`` selects a binary root. ``chi`` is the requested total + virtual-bond dimension. All block-sparse and fermionic operations are + delegated to Symmray/Quimb. """ from ..optimizers.tree import TreePlan, TreeTensorNetwork @@ -889,6 +913,9 @@ def hrs_to_ttn( if tree is None: if root_qubit is not None: root_qubit = int(root_qubit) + top_arity = _resolve_tree_top_arity( + top_arity, max_arity=max_arity, n=n, root_qubit=root_qubit + ) if order is None: order = ( range(n) diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index c42829c..7bdc146 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -910,12 +910,12 @@ def test_explicit_tree_rejects_mismatched_n(): TreeOptimizer(None, n=4, layout=plan, run=False) -def test_layout_finder_builds_valid_tree(): - """With max_arity=2 the finder returns a rooted binary tree over all qubits.""" +def test_layout_finder_builds_strict_binary_tree_when_requested(): + """Explicit top_arity=2 opts out of the ternary virtual root.""" rng = np.random.default_rng(8) n = 8 stream = _random_stream(n, 60, rng) - plan = TreeLayoutFinder(stream, n=n, max_arity=2).run() + plan = TreeLayoutFinder(stream, n=n, max_arity=2, top_arity=2).run() assert plan.n == n assert set(plan.leaf_of_qubit) == set(range(n)) # every internal node has exactly two children @@ -1538,19 +1538,21 @@ def test_recommend_layered_rejects_bad_chi(): finder.recommend_layered(block_sizes=(2,), chi=0) -def test_layout_finder_searches_arities_by_default(): - """The finder default searches (2, 3, 4) and stays chi-blind without chi.""" +def test_layout_finder_uses_binary_ternary_root_by_default(): + """The finder default is fixed binary below a ternary virtual root.""" rng = np.random.default_rng(213) stream = _random_stream(16, 60, rng, two_qubit_frac=0.7) finder = TreeLayoutFinder(stream, n=16, weight_mode="operator_schmidt") - assert finder.arity_candidates == (2, 3, 4) + assert finder.arity_candidates is None assert finder.chi is None - # run() returns the objective-best arity from the chi-blind recommendation. searched = finder.run() + assert searched.top_arity == 3 + assert searched.is_binary() + + # Candidate arity search remains available explicitly. blind = finder.recommend_arities((2, 3, 4)) assert blind["chi"] is None - assert searched.children == blind["plan"].children # A scalar max_arity opts back into a single fixed binary tree. fixed = TreeLayoutFinder(stream, n=16, max_arity=2, @@ -1598,8 +1600,8 @@ def capture(max_arities, **kwargs): assert finder._last_arity_recommendation["chi"] is None -def test_layout_finder_default_search_is_chi_aware_with_chi(): - """A finder built with ``chi`` makes its default arity search chi-aware.""" +def test_layout_finder_explicit_arity_search_is_chi_aware(): + """An explicit candidate search remains ``chi``-aware.""" rng = np.random.default_rng(214) stream = _random_stream(16, 60, rng, two_qubit_frac=0.7) finder = TreeLayoutFinder(stream, n=16, chi=64, @@ -1608,24 +1610,25 @@ def test_layout_finder_default_search_is_chi_aware_with_chi(): assert finder.chi == 64 searched = finder.run() aware = finder.recommend_arities((2, 3, 4), chi=64) - assert searched.children == aware["plan"].children + assert searched.top_arity == 3 + assert searched.is_binary() # The chi-aware search never overflows chi by more than the binary tree. by_arity = {c["max_arity"]: c for c in aware["candidates"]} chosen = by_arity[aware["recommended_max_arity"]] assert chosen["chi_overflow"] <= by_arity[2]["chi_overflow"] -def test_optimizer_searches_arities_by_default_chi_aware(): - """TreeOptimizer defaults to a chi-aware arity search using its own chi.""" +def test_optimizer_uses_binary_ternary_root_by_default(): + """TreeOptimizer shares the fixed binary/ternary-root default.""" rng = np.random.default_rng(215) stream = _random_stream(16, 60, rng, two_qubit_frac=0.7) opt = TreeOptimizer(stream, n=16, chi=64, layout_weight_mode="operator_schmidt", run=False) - # The optimizer forwards its chi into the finder's default arity search. - finder = TreeLayoutFinder(stream, n=16, max_arity=(2, 3, 4), chi=64, + finder = TreeLayoutFinder(stream, n=16, max_arity=2, top_arity=3, chi=64, weight_mode="operator_schmidt") assert opt.plan.children == finder.run().children + assert opt.plan.top_arity == 3 # A scalar max_arity=2 forces a fixed binary tree through the optimizer. fixed = TreeOptimizer(stream, n=16, chi=64, max_arity=2, @@ -1757,7 +1760,7 @@ def test_tree_layout_finder_plot_defaults_to_tent(): (pepsy.cnot(), (0, 3)), (pepsy.cnot(), (3, 1)), ] - finder = TreeLayoutFinder(gates, n=4, max_arity=2) + finder = TreeLayoutFinder(gates, n=4, max_arity=2, top_arity=2) plan = finder.run() assert plan.is_binary() assert len(plan.children[plan.root]) == 2 @@ -2063,9 +2066,12 @@ def test_tree_layout_nni_refinement_changes_binary_topology(): [(cnot, (0, 2)), (cnot, (0, 3))], n=4, max_arity=2, + top_arity=2, objective="path", ) - initial = TreePlan.from_order(range(4), structure="balanced", max_arity=2) + initial = TreePlan.from_order( + range(4), structure="balanced", max_arity=2, top_arity=2, + ) refined, planning = finder._refine_plan_topology( initial, @@ -2261,7 +2267,7 @@ def test_truncation_report_tracks_per_edge_discarded_weight(): [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]], dtype=complex, ) - plan = TreePlan.from_order(range(4), structure="balanced") + plan = TreePlan.from_order(range(4), structure="balanced", top_arity=2) opt = TreeOptimizer( [(h, 0), (cnot, (0, 3))], n=4, @@ -3237,6 +3243,38 @@ def test_binary_tree_supports_a_three_virtual_leg_top_tensor(): assert random.max_tensor_rank == 3 +def test_binary_tree_with_ternary_root_is_the_shared_default(): + """All high-level tree builders share the rank-three root convention.""" + plan = TreePlan.from_order(range(9), structure="balanced") + assert plan.top_arity == 3 + assert plan.is_binary() + assert not plan.is_strictly_binary() + + ordered = TreeTensorNetwork.from_order(range(9)) + assert ordered.top_arity == 3 + assert len(ordered.node_tensor(ordered.plan.root).inds) == 3 + + finder = TreeLayoutFinder([], n=9) + found = finder.run() + assert found.top_arity == 3 + assert found.is_binary() + + optimizer = TreeOptimizer([], n=9, run=False) + assert optimizer.plan.top_arity == 3 + assert optimizer.tn.max_tensor_rank == 3 + + product = pepsy.ps_to_ttn(9) + random = pepsy.hrs_to_ttn(9, seed=11) + assert product.top_arity == random.top_arity == 3 + assert product.max_tensor_rank == random.max_tensor_rank == 3 + + # A physical root cannot also use three incoming virtual bonds, and small + # systems naturally fall back to the ordinary binary root. + rooted = TreePlan.from_order(range(8), root_qubit=8) + assert rooted.top_arity == 2 + assert TreePlan.from_order(range(2)).top_arity == 2 + + def test_ps_to_ttn_matches_product_state_constructor_api(): """The high-level TTN constructor mirrors ``ps_to_mps`` amplitudes.""" theta = 0.31 diff --git a/tests/test_optimize_tree_stabilizer.py b/tests/test_optimize_tree_stabilizer.py index eccb6fa..9e41472 100644 --- a/tests/test_optimize_tree_stabilizer.py +++ b/tests/test_optimize_tree_stabilizer.py @@ -26,7 +26,17 @@ def _rzz(theta): np.exp(0.5j * theta), np.exp(-0.5j * theta), ] - ).astype(complex) +).astype(complex) + + +def test_tree_stab_uses_binary_tree_with_ternary_root_by_default(): + """The stabilizer facade propagates the shared tree geometry default.""" + optimizer = pepsy.TreeStabOptimizer(5) + assert optimizer.tree_optimizer.plan.top_arity == 3 + assert optimizer.tree_optimizer.plan.is_binary() + + binary_root = pepsy.TreeStabOptimizer(5, top_arity=2) + assert binary_root.tree_optimizer.plan.top_arity == 2 def _rz(theta): From 05d18008ab1ee48008940e6a5c85f48e6b5fa257 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Fri, 31 Jul 2026 15:33:04 -0600 Subject: [PATCH 43/70] Fix native fermion PEPO simple update --- src/pepsy/operators/gates.py | 158 ++++++++++++++++++++----- tests/test_symmetric_tensors.py | 200 ++++++++++++++++++++++++++++++++ 2 files changed, 327 insertions(+), 31 deletions(-) diff --git a/src/pepsy/operators/gates.py b/src/pepsy/operators/gates.py index 5697316..3625936 100644 --- a/src/pepsy/operators/gates.py +++ b/src/pepsy/operators/gates.py @@ -996,6 +996,39 @@ def _symmray_sample_data_from_tn(tn): return None +def _is_native_fermionic_array(value): + """Return whether ``value`` is a native Symmray fermionic array.""" + return _is_symmray_array(value) and type(value).__name__.endswith( + "FermionicArray" + ) + + +def _is_operator_tensor_network(tn): + """Return whether ``tn`` exposes distinct upper and lower site legs.""" + return ( + callable(getattr(tn, "upper_ind", None)) + and callable(getattr(tn, "lower_ind", None)) + ) + + +def _operator_ind_id(tn, which, arity): + """Return an operator's live upper/lower physical-index format.""" + attr = "upper_ind_id" if which == "upper" else "lower_ind_id" + return getattr(tn, attr, _ind_id_from_which(which, arity)) + + +def _operator_which_from_ind_id(tn, ind_id): + """Map an explicit operator physical-index format to its side.""" + if ind_id == getattr(tn, "upper_ind_id", None): + return "upper" + if ind_id == getattr(tn, "lower_ind_id", None): + return "lower" + raise ValueError( + "gate_simple() ind_id for an operator tensor network must match its " + "upper_ind_id or lower_ind_id. Prefer which='upper' or 'lower'." + ) + + def _symmray_index_map_for_tn_ind(tn, ind): """Return the Symmray charge map for a live tensor-network index.""" tensor_ids = getattr(tn, "ind_map", {}).get(ind, ()) @@ -1729,7 +1762,11 @@ def gate_simple( * Automatic long-range SWAP routing when the two sites are not adjacent (works for 1D / 2D / 3D ``where`` coordinates). * ``which``/``ind_id`` selection for vector-like networks whose physical - site-index family is not the default ``k...`` family. + site-index family is not the default ``k...`` family, and true one-sided + ``which='upper'`` / ``which='lower'`` updates for MPO/PEPO operators. + * A graded-safe native-fermion operator sandwich. With ``which=None``, a + native fermionic MPO/PEPO applies the upper gate and its conjugate lower + gate sequentially, avoiding Quimb's dense-style eager sandwich. * Dimension-aware, backend-aligned internal SWAP tensors for long-range routing through mixed physical dimensions. * Optional out-of-place semantics via ``inplace=False``. @@ -1751,8 +1788,11 @@ def gate_simple( gauges : dict Simple-update gauge dictionary keyed by bond index (mutated in place). which : {"upper", "lower"} | None, optional - Convenience selector for the physical index family. ``"upper"`` maps - to ``k...`` indices and ``"lower"`` maps to ``b...`` indices. + Convenience selector for the physical index family. On an MPO/PEPO it + selects a true one-sided operator update. With ``None``, operator + networks use sandwich semantics; native fermionic operators perform + this as two graded one-sided updates. On vector-like networks, + ``"upper"`` maps to ``k...`` and ``"lower"`` maps to ``b...``. ind_id : str | None, optional Explicit physical index format, e.g. ``"k{}"``, ``"b{},{}"``. renorm : bool, optional @@ -1836,34 +1876,68 @@ def gate_simple( else: raise ValueError("Could not infer gate dimensionality from where.") + which_local = ( + which_payload if which_payload is not None else which_default + ) + is_operator = _is_operator_tensor_network(tn_work) + operator_which = None ind_id_local = ind_id - if which_payload is not None: - ind_id_local = _ind_id_from_which(which_payload, arity) - elif which_default is not None: - ind_id_local = _ind_id_from_which(which_default, arity) + if is_operator: + if which_local is not None: + operator_which = which_local + ind_id_local = _operator_ind_id( + tn_work, operator_which, arity + ) + elif ind_id_local is not None: + operator_which = _operator_which_from_ind_id( + tn_work, ind_id_local + ) + elif which_local is not None: + ind_id_local = _ind_id_from_which(which_local, arity) if ind_id_local is not None: _validate_gate_target_inds_exist(tn_work, where_norm, ind_id_local) - _gate_simple_one( - tn_work, - gate_payload, - where_norm, - gauges, - renorm=renorm, - smudge=smudge, - gate_opts=gate_opts, - ind_id=ind_id_local, - sequence=sequence, - path_canonize=path_canonize, - path_canonize_distance=path_canonize_distance, - path_canonize_opts=path_canonize_opts, - path_compress=path_compress, - path_compress_max_bond=path_compress_max_bond, - path_compress_cutoff=path_compress_cutoff, - path_compress_canonize_distance=path_compress_canonize_distance, - path_compress_opts=path_compress_opts, - ) + gate_ind_id = None if is_operator else ind_id_local + gate_calls = ((gate_payload, operator_which, gate_ind_id),) + sample = _symmray_sample_data_from_tn(tn_work) + if ( + is_operator + and operator_which is None + and ind_id_local is None + and _is_native_fermionic_array(sample) + ): + # Quimb's eager two-site sandwich forms a dense-style product of + # the upper gate and ``conj(gate)`` before splitting. That loses + # the graded ordering for native FermionicArray data. Its + # one-sided operator paths preserve the grading, so realize + # ``G @ O @ G.H`` as upper ``G`` followed by lower ``conj(G)``. + gate_calls = ( + (gate_payload, "upper", None), + (ar.do("conj", gate_payload), "lower", None), + ) + + for gate_one, operator_which_one, ind_id_one in gate_calls: + _gate_simple_one( + tn_work, + gate_one, + where_norm, + gauges, + renorm=renorm, + smudge=smudge, + gate_opts=gate_opts, + ind_id=ind_id_one, + operator_which=operator_which_one, + sequence=sequence, + path_canonize=path_canonize, + path_canonize_distance=path_canonize_distance, + path_canonize_opts=path_canonize_opts, + path_compress=path_compress, + path_compress_max_bond=path_compress_max_bond, + path_compress_cutoff=path_compress_cutoff, + path_compress_canonize_distance=path_compress_canonize_distance, + path_compress_opts=path_compress_opts, + ) return tn_work @@ -2056,6 +2130,7 @@ def _gate_simple_one( smudge, gate_opts, ind_id=None, + operator_which=None, sequence=None, path_canonize=False, path_canonize_distance=1, @@ -2074,6 +2149,10 @@ def _gate_simple_one( "for MPO/PEPO operator layers." ) + gate_opts = dict(gate_opts) + if operator_which is not None: + gate_opts["which"] = operator_which + has_site_ind_id = hasattr(tn_work, "site_ind_id") old_site_ind_id = getattr(tn_work, "site_ind_id", None) if ind_id is not None: @@ -2127,13 +2206,26 @@ def _gate_simple_one_with_current_site_ind_id( path_compress_canonize_distance=0, path_compress_opts=None, ): - """Apply a single gate assuming ``site_ind_id`` has already been selected.""" + """Apply a single gate after selecting its site-index family or side.""" # One-site gate — no gauge update needed. if len(where) == 1: - tn_work.gate_simple_( - G, where=where, gauges=gauges, - renorm=False, smudge=smudge, inplace=True, - ) + operator_which = gate_opts.get("which") + if operator_which is None: + tn_work.gate_simple_( + G, where=where, gauges=gauges, + renorm=False, smudge=smudge, inplace=True, + ) + else: + # Quimb's one-tensor gate_simple shortcut does not forward + # gate_opts (including ``which``). There is no gauge to update, + # so use its direct one-sided operator gate instead. + tn_work.gate_( + G, + where=where, + which=operator_which, + contract=True, + inplace=True, + ) return tn_work # Two-site gate — check if the sites share a bond. @@ -2183,6 +2275,10 @@ def _gate_simple_one_with_current_site_ind_id( cast_complex_to_real=True, ) swap_ind_id = getattr(tn_work, "site_ind_id", None) + operator_which = gate_opts.get("which") + if swap_ind_id is None and operator_which is not None: + ndim = len(site_a) if isinstance(site_a, (tuple, list)) else 1 + swap_ind_id = _operator_ind_id(tn_work, operator_which, ndim) ndim = len(site_a) if isinstance(site_a, (tuple, list)) else 1 if ndim == 1: diff --git a/tests/test_symmetric_tensors.py b/tests/test_symmetric_tensors.py index f2823e4..4726f27 100644 --- a/tests/test_symmetric_tensors.py +++ b/tests/test_symmetric_tensors.py @@ -1273,6 +1273,206 @@ def test_fermion_to_pepo_native_result_supports_reverse_simple_update(): assert len(gauges) > 0 +def test_native_fermion_pepo_reverse_simple_update_matches_state_evolution(): + """Graded operator SU must realize G.H @ O @ G, not a dense sandwich.""" + fermion = Fermion(spinful=True, symmetry="U1") + left = (0, 0) + right = (1, 0) + where = (left, right) + mapper = OneDMap(2, 2, mode="snake") + + state = pepsy.ps_to_peps( + (2, 2), + fermion=fermion, + occupations={(x, y): 1 for x in range(2) for y in range(2)}, + seed=3, + dtype="complex128", + cyclic=False, + ) + for x in range(2): + for y in range(2): + sign = -1 if (x + y) % 2 == 0 else 1 + tensor = state[x, y] + (sector, block), = tensor.data.blocks.items() + tensor.data.blocks[sector] = ( + np.asarray([1.0, sign], dtype=np.complex128) + .reshape(block.shape) + / np.sqrt(2.0) + ) + + operator = fermion.to_pepo( + {where: fermion.eta_pair_operator()}, + Lx=2, + Ly=2, + mapper=mapper, + max_bond=64, + cutoff=0.0, + compress=False, + cyclic=False, + ) + forward_gate = fermion.hopping_gate(0.15, t=1.0) + evolved_state = gate( + state, + forward_gate, + where=where, + contract="split", + max_bond=64, + cutoff=0.0, + inplace=False, + ) + + def expectation(psi, pepo): + applied = pepo.apply(psi, contract=True, compress=False) + return complex(np.asarray((psi.H & applied).contract(all)).item()) + + reference = expectation(evolved_state, operator) + assert abs(reference) > 1.0e-3 + + explicit_operator = operator.copy() + explicit_gauges = {} + explicit_operator.gauge_all_simple_( + gauges=explicit_gauges, progbar=False + ) + explicit_operator = gate_simple( + explicit_operator, + forward_gate.H, + where=where, + which="upper", + gauges=explicit_gauges, + renorm=False, + smudge=1.0e-12, + max_bond=64, + cutoff=0.0, + contract="split", + inplace=False, + ) + gate_simple( + explicit_operator, + forward_gate.T, + where=where, + which="lower", + gauges=explicit_gauges, + renorm=False, + smudge=1.0e-12, + max_bond=64, + cutoff=0.0, + contract="split", + inplace=True, + ) + explicit_operator.gauge_simple_insert(explicit_gauges) + + gauges = {} + operator.gauge_all_simple_(gauges=gauges, progbar=False) + reverse_evolved = gate_simple( + operator, + forward_gate.H, + where=where, + gauges=gauges, + renorm=False, + smudge=1.0e-12, + max_bond=64, + cutoff=0.0, + contract="split", + inplace=False, + ) + reverse_evolved.gauge_simple_insert(gauges) + + np.testing.assert_allclose( + expectation(state, reverse_evolved), + reference, + atol=1.0e-8, + rtol=1.0e-8, + ) + np.testing.assert_allclose( + expectation(state, explicit_operator), + reference, + atol=1.0e-8, + rtol=1.0e-8, + ) + + +@pytest.mark.parametrize( + ("symmetry", "gate_name", "operator_name"), + [ + ("U1U1", "hopping", "eta"), + ("U1U1", "heisenberg", "number_up"), + ("U1", "hopping", "eta"), + ("U1", "heisenberg", "number_up"), + ("U1", "sxx", "number_up"), + ("Z2", "hopping", "eta"), + ("Z2", "heisenberg", "number_up"), + ("Z2", "sxx", "number_up"), + ], +) +def test_native_fermion_pepo_reverse_simple_update_unitary_families( + symmetry, gate_name, operator_name +): + """Several native unitaries must match the exact graded local sandwich.""" + fermion = Fermion(spinful=True, symmetry=symmetry) + where = ((0, 0), (1, 0)) + mapper = OneDMap(2, 1, mode="snake") + if operator_name == "eta": + operator = fermion.eta_pair_operator() + else: + operator = fermion.operator_term( + [(1.0, ((where[0], "number_up"),))], + sites=where, + ) + + if gate_name == "hopping": + forward_gate = fermion.hopping_gate(0.15, t=1.0) + elif gate_name == "heisenberg": + forward_gate = fermion.heisenberg_gate(0.11) + else: + forward_gate = fermion.sxx_gate(0.13) + + gate_matrix = forward_gate.fuse((0, 1), (2, 3)) + operator_matrix = operator.fuse((0, 1), (2, 3)) + exact_local = ( + gate_matrix.H @ operator_matrix @ gate_matrix + ).reshape((4, 4, 4, 4)) + pepo_opts = { + "Lx": 2, + "Ly": 1, + "mapper": mapper, + "max_bond": 256, + "cutoff": 0.0, + "compress": False, + "cyclic": False, + } + evolved = fermion.to_pepo({where: operator}, **pepo_opts) + exact = fermion.to_pepo({where: exact_local}, **pepo_opts) + + gauges = {} + evolved.gauge_all_simple_(gauges=gauges, progbar=False) + gate_simple( + evolved, + forward_gate.H, + where=where, + gauges=gauges, + renorm=False, + smudge=1.0e-12, + max_bond=256, + cutoff=0.0, + contract="split", + inplace=True, + ) + evolved.gauge_simple_insert(gauges) + + def hilbert_schmidt(left, right): + return complex( + np.asarray((left.H & right).contract(all)).item() + ) + + evolved_norm = hilbert_schmidt(evolved, evolved) + exact_norm = hilbert_schmidt(exact, exact) + overlap = hilbert_schmidt(evolved, exact) + relative_distance = abs( + evolved_norm + exact_norm - 2.0 * overlap.real + ) / max(abs(evolved_norm), abs(exact_norm)) + assert relative_distance < 1.0e-10 + + @pytest.mark.parametrize("symmetry", ["U1U1", "U1", "Z2"]) def test_fermion_to_pepo_supports_charged_odd_native_terms(symmetry): """Charged odd terms retain their native charge and dummy mode.""" From ed94c672d327b93f211d06bec184645a45dcd6c7 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Fri, 31 Jul 2026 21:25:01 -0600 Subject: [PATCH 44/70] Fix native fermionic PEPO PBC bonds --- src/pepsy/operators/gates.py | 193 ++++++++++++-- src/pepsy/operators/hamiltonians.py | 23 +- src/pepsy/tensors/bonds.py | 79 ++++++ src/pepsy/tensors/constructors.py | 7 +- src/pepsy/tensors/symmetric.py | 356 +++++++++++++++++++++++--- tests/test_native_fermion_pepo_2x3.py | 284 ++++++++++++++++++++ tests/test_symmetric_tensors.py | 37 +++ 7 files changed, 923 insertions(+), 56 deletions(-) create mode 100644 src/pepsy/tensors/bonds.py create mode 100644 tests/test_native_fermion_pepo_2x3.py diff --git a/src/pepsy/operators/gates.py b/src/pepsy/operators/gates.py index 3625936..b6b60e4 100644 --- a/src/pepsy/operators/gates.py +++ b/src/pepsy/operators/gates.py @@ -2,6 +2,8 @@ from __future__ import annotations +from contextlib import contextmanager +from contextvars import ContextVar from heapq import heappop, heappush from numbers import Integral import random @@ -14,6 +16,150 @@ import numpy as np import quimb.tensor as qtn + +# Symmray's cutoff-based block-SVD keeps every singular value tied with the +# global threshold. This is symmetry-respecting, but means ``max_bond`` is a +# soft cap whenever a cutoff is also supplied. Operator simple-update needs a +# predictable memory cap, so gate_simple exposes an opt-in hard total cap. +_STRICT_MAX_BOND_ACTIVE = ContextVar("pepsy_strict_max_bond_active", default=False) + + +def _hard_cap_blocksparse_svd(U, s, VH, max_bond): + """Keep the globally largest ``max_bond`` retained block-SVD values. + + The input has already been filtered by Symmray's requested cutoff. Ties + crossing the total cap are resolved deterministically by the stable block + and in-block ordering. The result retains the symmetry sectors selected by + that global spectrum and never has a larger total internal dimension than + ``max_bond``. + """ + sectors = tuple(U.sectors) + blocks = tuple(s.get_all_blocks()) + sizes = tuple(int(ar.size(block)) for block in blocks) + if sum(sizes) <= max_bond: + return U, s, VH + + ranked = [] + for sector_index, block in enumerate(blocks): + values = np.asarray(ar.to_numpy(block)).reshape(-1) + ranked.extend( + (-float(value), sector_index, value_index) + for value_index, value in enumerate(values) + ) + ranked.sort() + + kept_per_sector = [0] * len(sectors) + for _, sector_index, _ in ranked[:max_bond]: + kept_per_sector[sector_index] += 1 + + chargemap = {} + for (c0, c1), n_keep in zip(sectors, kept_per_sector): + if n_keep == 0: + U.del_block((c0, c1)) + s.del_block(c1) + VH.del_block((c1, c1)) + continue + U.set_block((c0, c1), U.get_block((c0, c1))[:, :n_keep]) + s.set_block(c1, s.get_block(c1)[:n_keep]) + VH.set_block((c1, c1), VH.get_block((c1, c1))[:n_keep, :]) + chargemap[c1] = n_keep + + chargemap = dict(sorted(chargemap.items())) + U.modify( + indices=( + U.indices[0], + U.indices[1].copy_with(chargemap=chargemap, subinfo=None), + ) + ) + VH.modify( + indices=( + VH.indices[0].copy_with(chargemap=chargemap, subinfo=None), + VH.indices[1], + ) + ) + return U, s, VH + + +def _install_strict_blocksparse_truncation(): + """Install a context-controlled hard-cap layer over Symmray's SVD. + + This stays dormant for every existing caller. It is activated only by the + ``strict_max_bond`` context in :func:`gate_simple` below. + """ + try: + from symmray.sparse import sparse_array_common as sparse_common + except ImportError: + return + + if getattr(sparse_common, "_pepsy_strict_max_bond_patch", False): + return + + original = sparse_common.truncate_svd_result_blocksparse + + def truncate_with_optional_hard_cap( + U, + s, + VH, + cutoff, + cutoff_mode, + max_bond, + absorb, + renorm, + backend=None, + use_abs=False, + ): + if not ( + _STRICT_MAX_BOND_ACTIVE.get() + and cutoff > 0.0 + and max_bond is not None + and max_bond > 0 + ): + return original( + U, + s, + VH, + cutoff=cutoff, + cutoff_mode=cutoff_mode, + max_bond=max_bond, + absorb=absorb, + renorm=renorm, + backend=backend, + use_abs=use_abs, + ) + + # First apply Symmray's numerical cutoff without an artificial cap; + # then impose an exact global dimension budget before absorption. + U, s, VH = original( + U, + s, + VH, + cutoff=cutoff, + cutoff_mode=cutoff_mode, + max_bond=-1, + absorb=None, + renorm=renorm, + backend=backend, + use_abs=use_abs, + ) + U, s, VH = _hard_cap_blocksparse_svd(U, s, VH, int(max_bond)) + return sparse_common.absorb_svd_result(U, s, VH, absorb) + + sparse_common.truncate_svd_result_blocksparse = truncate_with_optional_hard_cap + sparse_common._pepsy_strict_max_bond_patch = True + + +@contextmanager +def _strict_max_bond_context(enabled): + if not enabled: + yield + return + _install_strict_blocksparse_truncation() + token = _STRICT_MAX_BOND_ACTIVE.set(True) + try: + yield + finally: + _STRICT_MAX_BOND_ACTIVE.reset(token) + from ..backends.convert import ( infer_backend_converter_from_sample, resolve_backend_sample_data_from_tn, @@ -1741,6 +1887,7 @@ def gate_simple( max_bond=None, cutoff=1e-12, cutoff_mode="rsum2", + strict_max_bond: bool = False, contract=None, sequence="auto", path_canonize=False, @@ -1808,6 +1955,11 @@ def gate_simple( cutoff_mode : str, optional Cutoff mode passed to ``gate_simple_`` (e.g. ``'rsum2'``, ``'rel'``). Default ``'rsum2'``. + strict_max_bond : bool, optional + Enforce ``max_bond`` as a hard *total* block-sparse bond limit even + when a nonzero cutoff has degenerate singular values at its boundary. + This is disabled by default to preserve Symmray's standard + degeneracy-preserving truncation policy. contract : {None, "auto", "split", "reduce-split"}, optional Two-site split strategy passed to Quimb. The default ``None``/``"auto"`` selects ``"split"`` as a conservative fallback for block-sparse @@ -1918,26 +2070,27 @@ def gate_simple( ) for gate_one, operator_which_one, ind_id_one in gate_calls: - _gate_simple_one( - tn_work, - gate_one, - where_norm, - gauges, - renorm=renorm, - smudge=smudge, - gate_opts=gate_opts, - ind_id=ind_id_one, - operator_which=operator_which_one, - sequence=sequence, - path_canonize=path_canonize, - path_canonize_distance=path_canonize_distance, - path_canonize_opts=path_canonize_opts, - path_compress=path_compress, - path_compress_max_bond=path_compress_max_bond, - path_compress_cutoff=path_compress_cutoff, - path_compress_canonize_distance=path_compress_canonize_distance, - path_compress_opts=path_compress_opts, - ) + with _strict_max_bond_context(strict_max_bond): + _gate_simple_one( + tn_work, + gate_one, + where_norm, + gauges, + renorm=renorm, + smudge=smudge, + gate_opts=gate_opts, + ind_id=ind_id_one, + operator_which=operator_which_one, + sequence=sequence, + path_canonize=path_canonize, + path_canonize_distance=path_canonize_distance, + path_canonize_opts=path_canonize_opts, + path_compress=path_compress, + path_compress_max_bond=path_compress_max_bond, + path_compress_cutoff=path_compress_cutoff, + path_compress_canonize_distance=path_compress_canonize_distance, + path_compress_opts=path_compress_opts, + ) return tn_work diff --git a/src/pepsy/operators/hamiltonians.py b/src/pepsy/operators/hamiltonians.py index 7262dea..1b5e287 100644 --- a/src/pepsy/operators/hamiltonians.py +++ b/src/pepsy/operators/hamiltonians.py @@ -19,6 +19,7 @@ resolve_color_mode, ) from ..tensors.core import OneDMap +from ..tensors.bonds import new_native_bond __all__ = [ "ham_tn", @@ -626,11 +627,19 @@ def _add_missing_lattice_bonds_(self, pepo): if x + 1 < self.L_x: edge = frozenset(((x, y), (x + 1, y))) if edge not in chain_edges: - pepo[f"I{x},{y}"].new_bond(pepo[f"I{x + 1},{y}"], size=1) + new_native_bond( + pepo[f"I{x},{y}"], + pepo[f"I{x + 1},{y}"], + size=1, + ) if y + 1 < self.L_y: edge = frozenset(((x, y), (x, y + 1))) if edge not in chain_edges: - pepo[f"I{x},{y}"].new_bond(pepo[f"I{x},{y + 1}"], size=1) + new_native_bond( + pepo[f"I{x},{y}"], + pepo[f"I{x},{y + 1}"], + size=1, + ) return pepo def _add_cycle_bonds_(self, pepo, *, bond_dim=1): @@ -641,11 +650,17 @@ def _add_cycle_bonds_(self, pepo, *, bond_dim=1): if self.L_x > 1: for y in range(self.L_y): - pepo[f"I{self.L_x - 1},{y}"].new_bond(pepo[f"I0,{y}"], size=int(bond_dim)) + left = pepo[f"I{self.L_x - 1},{y}"] + right = pepo[f"I0,{y}"] + if not qtn.bonds(left, right): + new_native_bond(left, right, size=int(bond_dim)) if self.L_y > 1: for x in range(self.L_x): - pepo[f"I{x},{self.L_y - 1}"].new_bond(pepo[f"I{x},0"], size=int(bond_dim)) + top = pepo[f"I{x},{self.L_y - 1}"] + bottom = pepo[f"I{x},0"] + if not qtn.bonds(top, bottom): + new_native_bond(top, bottom, size=int(bond_dim)) return pepo diff --git a/src/pepsy/tensors/bonds.py b/src/pepsy/tensors/bonds.py new file mode 100644 index 0000000..5dd0766 --- /dev/null +++ b/src/pepsy/tensors/bonds.py @@ -0,0 +1,79 @@ +"""Bond construction helpers for native tensor data.""" + +from __future__ import annotations + +import quimb.tensor as qtn + +__all__ = ["new_native_bond"] + + +def _repair_fermionic_duals(tensor_a, tensor_b, bond): + """Make a newly-created native fermionic bond contractible. + + Quimb's ``new_bond`` only knows the dense index size. For a native + Symmray tensor it consequently creates the new index with the same + dualness on both tensors. That is invisible for a dimension-one dense + bond, but it changes the graded sign when the bond is crossed by an + operator gate. Native contractions require opposite dual orientations. + """ + data_a = getattr(tensor_a, "data", None) + data_b = getattr(tensor_b, "data", None) + if not ( + bool(getattr(data_a, "fermionic", False)) + and bool(getattr(data_b, "fermionic", False)) + ): + return False + + try: + axis_a = tensor_a.inds.index(bond) + axis_b = tensor_b.inds.index(bond) + index_a = data_a.indices[axis_a] + index_b = data_b.indices[axis_b] + dual_a = bool(index_a.dual) + dual_b = bool(index_b.dual) + except (AttributeError, IndexError, ValueError): + return False + + if dual_a != dual_b: + return False + + # Flip only the second endpoint. The data shape and sector blocks are + # unchanged; only the Symmray index orientation is corrected. + indices_b = list(data_b.indices) + indices_b[axis_b] = index_b.conj() + data_b.modify(indices=tuple(indices_b)) + return True + + +def new_native_bond( + tensor_a, + tensor_b, + *, + size=1, + name=None, + axis1=0, + axis2=0, +): + """Add a bond and repair its dual orientation for native fermions. + + Dense and ordinary Abelian tensor networks follow the same path as + ``quimb.tensor.new_bond``. For native Symmray fermionic arrays, the + newly-created shared index is checked and one endpoint is dual-flipped + when both endpoints were initialized with the same orientation. + """ + before = set(tensor_a.inds).intersection(tensor_b.inds) + qtn.new_bond( + tensor_a, + tensor_b, + size=size, + name=name, + axis1=axis1, + axis2=axis2, + ) + after = set(tensor_a.inds).intersection(tensor_b.inds) + new_bonds = after.difference(before) + if len(new_bonds) == 1: + bond = next(iter(new_bonds)) + _repair_fermionic_duals(tensor_a, tensor_b, bond) + return bond + return next(iter(after.difference(before)), None) diff --git a/src/pepsy/tensors/constructors.py b/src/pepsy/tensors/constructors.py index 7cf2a76..e5ece60 100644 --- a/src/pepsy/tensors/constructors.py +++ b/src/pepsy/tensors/constructors.py @@ -12,6 +12,7 @@ import quimb.tensor as qtn from .contractions import build_optimizer, tn_norm +from .bonds import new_native_bond from .validation import validate_tensor_network_tags __all__ = [ @@ -46,13 +47,15 @@ def add_cycle(peps, bond_dim, cylinder=False): for j in range(Ly): T1 = peps[f"I{Lx-1},{j}"] T2 = peps[f"I{0},{j}"] - qtn.new_bond(T1, T2, size=bond_dim, name=None, axis1=0, axis2=0) + if not qtn.bonds(T1, T2): + new_native_bond(T1, T2, size=bond_dim, axis1=0, axis2=0) if not cylinder: for i in range(Lx): T1 = peps[f"I{i},{Ly-1}"] T2 = peps[f"I{i},{0}"] - qtn.new_bond(T1, T2, size=bond_dim, name=None, axis1=0, axis2=0) + if not qtn.bonds(T1, T2): + new_native_bond(T1, T2, size=bond_dim, axis1=0, axis2=0) return peps diff --git a/src/pepsy/tensors/symmetric.py b/src/pepsy/tensors/symmetric.py index bb16450..9f8a040 100644 --- a/src/pepsy/tensors/symmetric.py +++ b/src/pepsy/tensors/symmetric.py @@ -5919,6 +5919,232 @@ def _add_native_term_to_mpo( return list(physical_maps[0]) +def _native_local_term_mpo( + term, + support, + L, + *, + symmetry, + dtype, + max_bond=None, + cutoff=1e-12, + compress=True, + upper_ind_id="k{}", + lower_ind_id="b{}", + site_tag_id="I{}", + to_backend=None, +): + """Build an exact local-term MPO without start/done channel inflation. + + The generic native MPO assembler is designed for a collection of terms, + so it carries explicit start and done paths at every chain cut. For one + one-site or two-site term those paths are unnecessary. Factorizing the + native local array directly and propagating its operator-Schmidt bond + through identity tensors leaves only the non-zero local Schmidt sectors. + """ + support = tuple(int(site) for site in support) + if len(support) not in {1, 2} or len(set(support)) != len(support): + raise ValueError("direct local PEPO terms must act on one or two sites.") + if any(site < 0 or site >= int(L) for site in support): + raise ValueError(f"term support {support!r} is outside MPO length L={L}.") + + _require_symmray() + from symmray import utils as sr_utils # pylint: disable=import-outside-toplevel + + zero = _zero_like_charge(0 if symmetry in {"U1", "Z2"} else (0, 0)) + term_charge = _normalize_group_charge( + getattr(term, "charge", zero), symmetry + ) + indices = getattr(term, "indices", None) + if indices is None or len(indices) != 2 * len(support): + raise TypeError("direct local PEPO terms require matching native rank.") + + physical_maps = [ + _expanded_index_charges(index) for index in indices[:len(support)] + ] + input_maps = [ + _expanded_index_charges(index) for index in indices[len(support):] + ] + if physical_maps != input_maps or any( + physical_maps[site] != physical_maps[0] + for site in range(len(support)) + ): + raise ValueError( + "direct local PEPO terms require one matching physical charge map." + ) + phys_map = physical_maps[0] + phys_dim = len(phys_map) + zero_map = [zero] + + def make_array(data, index_maps, duals, *, charge=zero, label=None): + return sr_utils.from_dense( + data, + symmetry=symmetry, + index_maps=index_maps, + duals=duals, + fermionic=True, + charge=charge, + label=label, + ) + + def identity_tensor(site): + identity = np.eye(phys_dim, dtype=dtype) + if L == 1: + return make_array(identity, [phys_map, phys_map], [False, True]) + if site == 0: + return make_array( + identity.reshape(1, phys_dim, phys_dim), + [zero_map, phys_map, phys_map], + [False, False, True], + ) + if site == L - 1: + return make_array( + identity.reshape(1, phys_dim, phys_dim), + [zero_map, phys_map, phys_map], + [True, False, True], + ) + data = np.zeros((1, 1, phys_dim, phys_dim), dtype=dtype) + data[0, 0] = identity + return make_array( + data, + [zero_map, zero_map, phys_map, phys_map], + [True, False, False, True], + ) + + arrays = [identity_tensor(site) for site in range(int(L))] + local_schmidt_bond = 1 + + if len(support) == 1: + site = support[0] + dense = _dense_numpy(term, dtype=dtype) + label = site if _charged_op_needs_fermion_string(term_charge) else None + if L == 1: + arrays[site] = make_array( + dense, [phys_map, phys_map], [False, True], + charge=term_charge, label=label, + ) + elif site == 0: + arrays[site] = make_array( + dense.reshape(1, phys_dim, phys_dim), + [zero_map, phys_map, phys_map], + [False, False, True], + charge=term_charge, label=label, + ) + elif site == L - 1: + arrays[site] = make_array( + dense.reshape(1, phys_dim, phys_dim), + [zero_map, phys_map, phys_map], + [True, False, True], + charge=term_charge, label=label, + ) + else: + arrays[site] = make_array( + dense.reshape(1, 1, phys_dim, phys_dim), + [zero_map, zero_map, phys_map, phys_map], + [True, False, False, True], + charge=term_charge, label=label, + ) + else: + # Order the support by the MPO chain. A native operator's upper and + # lower legs are reordered together so its graded local signs survive. + ordered = tuple(sorted(enumerate(support), key=lambda item: item[1])) + if tuple(item[0] for item in ordered) == (0, 1): + ordered_term = term + else: + ordered_term = term.transpose((1, 0, 3, 2)) + fused = ordered_term.fuse((0, 2), (1, 3)) + # The only cutoff here removes exact numerical zero singular blocks + # left by Symmray's block SVD. It is not the user-requested PEPO + # compression cutoff and does not cap the resulting local bond. + structural_cutoff = 64.0 * np.finfo(float).eps + left, _, right = fused.svd( + absorb="right", + cutoff=structural_cutoff, + ) + left = left.unfuse(0).transpose((2, 0, 1)) + right = right.unfuse(1) + bond_map = _expanded_index_charges(left.indices[0]) + bond_dim = len(bond_map) + local_schmidt_bond = bond_dim + left_dense = _dense_numpy(left, dtype=dtype) + right_dense = _dense_numpy(right, dtype=dtype) + left_charge = _normalize_group_charge( + getattr(left, "charge", zero), symmetry + ) + right_charge = _normalize_group_charge( + getattr(right, "charge", zero), symmetry + ) + left_site, right_site = (item[1] for item in ordered) + + if left_site == 0: + arrays[left_site] = left + else: + arrays[left_site] = make_array( + left_dense.reshape(1, bond_dim, phys_dim, phys_dim), + [zero_map, bond_map, phys_map, phys_map], + [True, False, False, True], + charge=left_charge, + ) + if right_site == L - 1: + arrays[right_site] = right + else: + arrays[right_site] = make_array( + right_dense.reshape(bond_dim, 1, phys_dim, phys_dim), + [bond_map, zero_map, phys_map, phys_map], + [True, False, False, True], + charge=right_charge, + ) + identity = np.eye(phys_dim, dtype=dtype) + for site in range(left_site + 1, right_site): + data = np.zeros( + (bond_dim, bond_dim, phys_dim, phys_dim), + dtype=dtype, + ) + for bond_pos in range(bond_dim): + data[bond_pos, bond_pos] = identity + arrays[site] = make_array( + data, + [bond_map, bond_map, phys_map, phys_map], + [True, False, False, True], + ) + + mpo = qtn.MatrixProductOperator( + arrays, + shape="lrud", + upper_ind_id=upper_ind_id, + lower_ind_id=lower_ind_id, + site_tag_id=site_tag_id, + ) + if to_backend is not None: + _apply_to_tensor_network_arrays(mpo, to_backend) + raw_bond = mpo.max_bond() + raw_max_bond = 1 if raw_bond is None else int(raw_bond) + did_compress = bool(compress and L > 1) + if did_compress: + compress_opts = {"cutoff": cutoff} + if max_bond is not None: + compress_opts["max_bond"] = int(max_bond) + mpo.compress(**compress_opts) + final_bond = mpo.max_bond() + final_max_bond = 1 if final_bond is None else int(final_bond) + mpo.pepsy_compression_report = { + "direct_local": True, + "compressed": did_compress, + "cutoff": cutoff, + "requested_max_bond": None if max_bond is None else int(max_bond), + "operator_schmidt_bond": local_schmidt_bond, + "raw_max_bond": raw_max_bond, + "final_max_bond": final_max_bond, + "rank_reduced": final_max_bond < raw_max_bond, + "max_bond_exceeded": ( + did_compress + and max_bond is not None + and final_max_bond > int(max_bond) + ), + } + return mpo + + def _generic_symhamiltonian_to_mpo( hamiltonian, L, @@ -6498,6 +6724,77 @@ def to_pepo( else dtype ), ) + + # A single local term does not need the generic start/done channel + # construction used to combine a Hamiltonian. Build its native MPO + # directly from the local operator Schmidt factorization instead. In + # particular, a hopping term then has its physical rank (D=4 for the + # spinful U1 hopping operator) rather than the inflated collection + # channel count of the multi-term assembler. Charged terms keep the + # generic native route because their open boundary must carry the + # operator charge through the remaining chain. + if fermionic and len(self.terms) == 1: + raw_where, term = next(iter(self.terms.items())) + coordinate_sites = _term_mapping_uses_coordinate_sites(self.terms) + where = _as_term_where( + raw_where, + coordinate_sites=coordinate_sites, + ) + zero = _zero_like_charge( + 0 if self.symmetry in {"U1", "Z2"} else (0, 0) + ) + term_charge = _normalize_group_charge( + getattr(term, "charge", zero), + self.symmetry, + ) + if len(where) in {1, 2} and term_charge == zero: + _, coo2idx_use, mapped_L = _resolve_mpo_mapping( + mapper=builder.mapper, + ) + if mapped_L != builder.L: + raise ValueError( + f"MPO mapping length {mapped_L} does not match PEPO length " + f"{builder.L}." + ) + mapped_where = tuple( + _map_site_to_mpo_index(site, coo2idx_use) + for site in where + ) + dtype_use = ( + _dtype_from_hamiltonian_terms(self.terms) + if dtype is None + else np.dtype(dtype) + ) + mpo = _native_local_term_mpo( + term, + mapped_where, + builder.L, + symmetry=self.symmetry, + dtype=dtype_use, + max_bond=max_bond, + cutoff=cutoff, + compress=compress, + to_backend=to_backend, + ) + pepo = builder.mpo_to_pepo( + mpo, + cycle_peps=cyclic, + cycle_bond_dim=cycle_bond_dim, + inplace=True, + ) + # Keep the diagnostic on the returned PEPO after the MPO is + # relabelled and viewed as a PEPO. + pepo.pepsy_compression_report = dict( + mpo.pepsy_compression_report + ) + if charge_sectors: + charge = _normalize_group_charge( + getattr(term, "charge", 0), + self.symmetry, + ) + return {charge: pepo} + return pepo + mpo = self.to_mpo( L=builder.L, mapper=builder.mapper, @@ -10174,46 +10471,45 @@ def to_pepo( """ if Lx is None or Ly is None: raise TypeError("to_pepo requires both Lx and Ly.") + if hamiltonian is not None: + if terms_or_edges is not None: + raise TypeError( + "Pass either terms_or_edges or hamiltonian, not both." + ) + if not isinstance(hamiltonian, SymHamiltonian): + raise TypeError("hamiltonian must be a SymHamiltonian instance.") + target = hamiltonian + elif isinstance(terms_or_edges, SymHamiltonian): + target = terms_or_edges + else: + if terms_or_edges is None: + raise TypeError("to_pepo requires terms_or_edges or hamiltonian.") + target = self.hamiltonian( + terms_or_edges, + to_backend=to_backend, + **params, + ) + params = {} - from ..operators.hamiltonians import ham_tn - - builder = ham_tn( + if params: + names = ", ".join(sorted(params)) + raise TypeError( + "Model parameters cannot be supplied with an existing " + f"SymHamiltonian: {names}." + ) + return target.to_pepo( Lx=Lx, Ly=Ly, mapper=mapper, - max_bond=256 if max_bond is None else max_bond, - cutoff=cutoff, - data_type=self.dtype if dtype is None else dtype, - ) - mpo = self.to_mpo( - terms_or_edges, - hamiltonian=hamiltonian, - L=builder.L, - mapper=builder.mapper, max_bond=max_bond, cutoff=cutoff, compress=compress, + cyclic=cyclic, + cycle_bond_dim=cycle_bond_dim, dtype=dtype, fermionic=fermionic, - charge_sectors=charge_sectors, to_backend=to_backend, - **params, - ) - if charge_sectors: - return { - charge: builder.mpo_to_pepo( - sector_mpo, - cycle_peps=cyclic, - cycle_bond_dim=cycle_bond_dim, - inplace=True, - ) - for charge, sector_mpo in mpo.items() - } - return builder.mpo_to_pepo( - mpo, - cycle_peps=cyclic, - cycle_bond_dim=cycle_bond_dim, - inplace=True, + charge_sectors=charge_sectors, ) def local_terms(self, edges, *, layout="site", **params): diff --git a/tests/test_native_fermion_pepo_2x3.py b/tests/test_native_fermion_pepo_2x3.py new file mode 100644 index 0000000..cfa03cd --- /dev/null +++ b/tests/test_native_fermion_pepo_2x3.py @@ -0,0 +1,284 @@ +"""Small native-Symmray U1/U1U1 PEPO checks on a 2x3 lattice.""" + +import numpy as np +import pytest +import quimb.tensor as qtn + +import pepsy +from pepsy.tensors import OneDMap + + +def _expectation(state, operator): + acted = operator.apply(state, contract=True, compress=False) + numerator = complex( + np.asarray((state.H & acted).contract(all, optimize="auto-hq")).item() + ) + denominator = complex( + np.asarray((state.H & state).contract(all, optimize="auto-hq")).item() + ) + return numerator / denominator + + +def _state(fermion, symmetry): + lattice = (2, 3) + sites = { + (x, y): ( + 1 + if symmetry == "U1" + else ((1, 0) if (x + y) % 2 == 0 else (0, 1)) + ) + for x in range(lattice[0]) + for y in range(lattice[1]) + } + state = pepsy.ps_to_peps( + lattice, + fermion=fermion, + occupations=sites, + seed=3, + dtype="complex128", + cyclic=False, + ) + if symmetry == "U1": + # Use a nontrivial charge-one local superposition, as in the Etienne + # Neel-X state, while preserving total-U1 charge at every site. + for x, y in sites: + sign = -1.0 if (x + y) % 2 == 0 else 1.0 + tensor = state[x, y] + (sector, block), = tensor.data.blocks.items() + tensor.data.blocks[sector] = ( + np.asarray([1.0, sign], dtype=np.complex128) + .reshape(block.shape) + / np.sqrt(2.0) + ) + return state + + +@pytest.mark.parametrize("symmetry", ["U1", "U1U1"]) +def test_native_fermion_pepo_2x3_gate_simple_matches_state(symmetry): + """Native U1 and U1U1 PEPO projection agrees with direct evolution.""" + pytest.importorskip("symmray") + + fermion = pepsy.Fermion( + spinful=True, + symmetry=symmetry, + dtype="complex128", + ) + state = _state(fermion, symmetry) + mapper = OneDMap(2, 3, mode="snake") + where = ((0, 0),) + operator = fermion.to_pepo( + {where: fermion.observable("number_up")}, + Lx=2, + Ly=3, + mapper=mapper, + max_bond=64, + cutoff=0.0, + compress=False, + cyclic=False, + fermionic=True, + ) + gates = ( + ( + fermion.hopping_gate(0.11, t=0.7), + ((0, 0), (1, 0)), + ), + ( + fermion.heisenberg_gate(0.07), + ((0, 0), (0, 1)), + ), + ) + + direct = state.copy() + backward = operator.copy() + gauges = {} + backward.gauge_all_simple_(gauges=gauges, progbar=False) + for gate, gate_where in gates: + direct = pepsy.gate( + direct, + gate, + where=gate_where, + contract="split", + max_bond=64, + cutoff=0.0, + inplace=True, + ) + + # Heisenberg replay is reverse-order, while the state stream is forward. + for gate, gate_where in reversed(gates): + pepsy.gate_simple( + backward, + gate.H, + where=gate_where, + gauges=gauges, + renorm=False, + max_bond=64, + cutoff=0.0, + contract="split", + inplace=True, + ) + measured = backward.copy() + measured.gauge_simple_insert(gauges) + + direct_value = _expectation(direct, operator) + backward_value = _expectation(state, measured) + assert abs(backward_value - direct_value) < 1.0e-8 + assert all( + type(tensor.data).__name__.endswith("FermionicArray") + for tensor in operator + ) + assert all( + type(tensor.data).__name__.endswith("FermionicArray") + for tensor in measured + ) + + +@pytest.mark.parametrize("symmetry", ["U1", "U1U1"]) +def test_native_fermion_2x3_hopping_plus_u_sandwich_projection(symmetry): + """Project ``U.H @ (T + U n_up n_down) @ U`` natively on 2x3.""" + pytest.importorskip("symmray") + + fermion = pepsy.Fermion( + spinful=True, + symmetry=symmetry, + dtype="complex128", + ) + state = _state(fermion, symmetry) + mapper = OneDMap(2, 3, mode="snake") + hopping_support = ((0, 0), (1, 0)) + interaction_support = ((0, 0),) + terms = { + hopping_support: -0.7 * fermion.hopping_operator(), + interaction_support: fermion.onsite_term((0, 0), U=8.0), + } + operator = fermion.to_pepo( + terms, + Lx=2, + Ly=3, + mapper=mapper, + max_bond=64, + cutoff=0.0, + compress=False, + cyclic=False, + fermionic=True, + ) + unitary = fermion.hopping_gate(0.11, t=0.7) + + direct = pepsy.gate( + state, + unitary, + where=hopping_support, + contract="split", + max_bond=64, + cutoff=0.0, + inplace=False, + ) + backward = operator.copy() + gauges = {} + backward.gauge_all_simple_(gauges=gauges, progbar=False) + pepsy.gate_simple( + backward, + unitary.H, + where=hopping_support, + gauges=gauges, + renorm=False, + max_bond=64, + cutoff=0.0, + contract="split", + inplace=True, + ) + measured = backward.copy() + measured.gauge_simple_insert(gauges) + + direct_value = _expectation(direct, operator) + projected_value = _expectation(state, measured) + assert projected_value == pytest.approx(direct_value, abs=1.0e-8) + + # The projected sum must also agree with independently projected native + # term PEPOs, which catches a gauge/projection error hidden by cancellation. + separate_projected = 0.0j + for support, term in terms.items(): + term_operator = fermion.to_pepo( + {support: term}, + Lx=2, + Ly=3, + mapper=mapper, + max_bond=64, + cutoff=0.0, + compress=False, + cyclic=False, + fermionic=True, + ) + term_gauges = {} + term_operator.gauge_all_simple_(gauges=term_gauges, progbar=False) + pepsy.gate_simple( + term_operator, + unitary.H, + where=hopping_support, + gauges=term_gauges, + renorm=False, + max_bond=64, + cutoff=0.0, + contract="split", + inplace=True, + ) + term_operator.gauge_simple_insert(term_gauges) + separate_projected += _expectation(state, term_operator) + + assert projected_value == pytest.approx(separate_projected, abs=1.0e-8) + + +def test_native_fermion_pepo_nonchain_edge_identity_sandwich_4x2(): + """Dimension-one lattice bonds preserve a native fermionic sandwich.""" + pytest.importorskip("symmray") + + fermion = pepsy.Fermion( + spinful=True, + symmetry="U1", + dtype="complex128", + ) + state = pepsy.ps_to_peps( + (4, 2), + fermion=fermion, + occupations={(x, y): 1 for x in range(4) for y in range(2)}, + seed=17, + dtype="complex128", + cyclic=False, + ) + operator = fermion.to_pepo( + {((0, 0),): fermion.observable("identity")}, + Lx=4, + Ly=2, + mapper=OneDMap(4, 2, mode="snake"), + max_bond=64, + cutoff=0.0, + compress=False, + cyclic=True, + fermionic=True, + ) + left = operator["I2,0"] + right = operator["I3,0"] + bond = next(iter(qtn.bonds(left, right))) + left_axis = left.inds.index(bond) + right_axis = right.inds.index(bond) + assert left.data.indices[left_axis].dual != right.data.indices[right_axis].dual + + hopping = fermion.hopping_gate(0.17, t=0.73) + projected = operator.copy() + pepsy.gate( + projected, + hopping.H, + where=((3, 0), (2, 0)), + which="upper", + contract=True, + inplace=True, + ) + pepsy.gate( + projected, + hopping.T, + where=((3, 0), (2, 0)), + which="lower", + contract=True, + inplace=True, + ) + + assert _expectation(state, projected) == pytest.approx(1.0, abs=1.0e-10) diff --git a/tests/test_symmetric_tensors.py b/tests/test_symmetric_tensors.py index 4726f27..e8a2095 100644 --- a/tests/test_symmetric_tensors.py +++ b/tests/test_symmetric_tensors.py @@ -1236,6 +1236,43 @@ def test_fermion_to_pepo_builds_native_coordinate_terms(symmetry): assert all(type(tensor.data).__name__.endswith("FermionicArray") for tensor in pepo) +@pytest.mark.parametrize("symmetry", ["U1U1", "U1", "Z2"]) +def test_single_native_pepo_term_uses_local_operator_schmidt_rank(symmetry): + """A disposable one-/two-site PEPO has no generic channel inflation.""" + fermion = Fermion(spinful=True, symmetry=symmetry) + mapper = OneDMap(2, 1, mode="snake") + + hopping = fermion.to_pepo( + {((0, 0), (1, 0)): fermion.hopping_operator()}, + Lx=2, + Ly=1, + mapper=mapper, + max_bond=None, + cutoff=0.0, + compress=False, + ) + onsite = fermion.to_pepo( + {((0, 0),): fermion.onsite_term((0, 0), U=8.0)}, + Lx=2, + Ly=1, + mapper=mapper, + max_bond=None, + cutoff=0.0, + compress=False, + ) + + assert hopping.max_bond() == 4 + assert onsite.max_bond() == 1 + assert hopping.pepsy_compression_report["operator_schmidt_bond"] == 4 + assert onsite.pepsy_compression_report["operator_schmidt_bond"] == 1 + assert hopping.pepsy_compression_report["direct_local"] is True + assert onsite.pepsy_compression_report["direct_local"] is True + assert all( + type(tensor.data).__name__.endswith("FermionicArray") + for tensor in (*hopping, *onsite) + ) + + def test_fermion_to_pepo_native_result_supports_reverse_simple_update(): """Native PEPO output can take an adjoint gate through operator SU.""" fermion = Fermion(spinful=True, symmetry="U1U1") From 07fb64633e2e277b622f77487cc145dd97f7857b Mon Sep 17 00:00:00 2001 From: rezaquant Date: Fri, 31 Jul 2026 21:54:01 -0600 Subject: [PATCH 45/70] Add safe native fermionic identity PEPO --- docs/api/tensors/symmetric.md | 16 ++++ src/pepsy/tensors/constructors.py | 123 ++++++++++++++++++++++++++++- src/pepsy/tensors/symmetric.py | 40 ++++++++++ tests/test_native_identity_pepo.py | 96 ++++++++++++++++++++++ 4 files changed, 273 insertions(+), 2 deletions(-) create mode 100644 tests/test_native_identity_pepo.py diff --git a/docs/api/tensors/symmetric.md b/docs/api/tensors/symmetric.md index 9ff6233..c333b9f 100644 --- a/docs/api/tensors/symmetric.md +++ b/docs/api/tensors/symmetric.md @@ -154,6 +154,22 @@ method is not implemented yet. Both constructors return the underlying PEPS with native fermionic Symmray tensors; use ``SymPEPS`` only when wrapper methods or stored Hamiltonian metadata are needed. +For the corresponding full operator identity, use ``id_to_pepo`` with the +same model: + +```python +identity = py.id_to_pepo( + (Lx, Ly), + fermion=fh, + cyclic=True, +) +``` + +This returns a native graded PEPO containing every local charge sector and +repairs periodic bond orientations. ``occupations`` and ``site_charge`` are +intentionally rejected here: they select a product-state sector and would +make the result something other than the full identity. + ## Unified native fermion helper ``Fermion`` is the model-facing helper for both one-mode spinless fermions diff --git a/src/pepsy/tensors/constructors.py b/src/pepsy/tensors/constructors.py index e5ece60..0a6a70a 100644 --- a/src/pepsy/tensors/constructors.py +++ b/src/pepsy/tensors/constructors.py @@ -59,7 +59,57 @@ def add_cycle(peps, bond_dim, cylinder=False): return peps -def id_to_pepo(lx, ly, phys_dim=2, dtype="complex128", chi=1, rand_strength=0.0): +def _native_fermion_identity_pepo( + fermion, + lx, + ly, + *, + cyclic=False, + cycle_bond_dim=1, + mapper=None, + max_bond=None, + cutoff=1e-12, + compress=False, + dtype="complex128", + to_backend=None, +): + """Build a full native fermionic identity without state-sector slicing.""" + identity = fermion.observable("identity") + target = fermion.hamiltonian({((0, 0),): identity}, to_backend=to_backend) + return target.to_pepo( + Lx=lx, + Ly=ly, + mapper=mapper, + max_bond=max_bond, + cutoff=cutoff, + compress=compress, + cyclic=cyclic, + cycle_bond_dim=cycle_bond_dim, + dtype=dtype, + fermionic=True, + to_backend=to_backend, + ) + + +def id_to_pepo( + lx, + ly=None, + phys_dim=2, + dtype="complex128", + chi=1, + rand_strength=0.0, + *, + fermion=None, + cyclic=False, + cycle_bond_dim=1, + mapper=None, + max_bond=None, + cutoff=1e-12, + compress=False, + to_backend=None, + occupations=None, + site_charge=None, +): """Create a PEPO identity on an ``lx x ly`` lattice. Parameters @@ -67,7 +117,8 @@ def id_to_pepo(lx, ly, phys_dim=2, dtype="complex128", chi=1, rand_strength=0.0) lx : int Lattice size in x direction. ly : int - Lattice size in y direction. + Lattice size in y direction. If omitted, ``lx`` must be a two-item + ``(Lx, Ly)`` shape, matching :func:`ps_to_peps`. phys_dim : int, optional Physical dimension per site. dtype : str, optional @@ -77,12 +128,78 @@ def id_to_pepo(lx, ly, phys_dim=2, dtype="complex128", chi=1, rand_strength=0.0) expanded via ``expand_bond_dimension`` after initialization. rand_strength : float, optional Random noise strength passed to ``expand_bond_dimension``. + fermion : :class:`~pepsy.tensors.Fermion`, optional + If supplied, construct a native Symmray fermionic identity. The + physical dimension is inferred from the model when the default + ``phys_dim=2`` is left in place. + cyclic : bool, optional + If True on the native fermionic path, add repaired dimension-one + bonds around the PEPO lattice using ``cycle_bond_dim``. + cycle_bond_dim : int, optional + Periodic bond dimension for the native fermionic path. + mapper, max_bond, cutoff, compress, to_backend + Forwarded to the native Fermion PEPO construction on the native path. + occupations, site_charge : optional + Rejected on the identity path. These arguments select a state sector + for :func:`ps_to_peps`; they must not remove diagonal blocks from a + full local identity operator. Returns ------- quimb.tensor.PEPO Identity PEPO with bond dimension ``chi``. """ + if ly is None: + if not isinstance(lx, (tuple, list)) or len(lx) != 2: + raise TypeError("id_to_pepo requires Lx and Ly, or a 2-item shape.") + lx, ly = lx + lx = int(lx) + ly = int(ly) + if lx < 1 or ly < 1: + raise ValueError("PEPO dimensions must be positive integers.") + + if fermion is not None: + from .symmetric import Fermion # pylint: disable=import-outside-toplevel + + if not isinstance(fermion, Fermion): + raise TypeError("fermion must be a pepsy.tensors.Fermion instance.") + if occupations is not None or site_charge is not None: + raise ValueError( + "occupations and site_charge select product-state sectors; " + "a fermionic identity PEPO contains the full local identity." + ) + if chi != 1 or rand_strength != 0.0: + raise ValueError( + "chi and rand_strength are dense PEPO expansion controls; " + "use max_bond/cutoff/compress for a native fermionic identity." + ) + local_dim = sum(int(size) for size in fermion.physical_sectors.values()) + if phys_dim not in (None, 2, local_dim): + raise ValueError( + f"phys_dim={phys_dim!r} does not match the fermion local " + f"dimension {local_dim}." + ) + return _native_fermion_identity_pepo( + fermion, + lx, + ly, + mapper=mapper, + max_bond=max_bond, + cutoff=cutoff, + compress=compress, + cyclic=cyclic, + cycle_bond_dim=cycle_bond_dim, + dtype=dtype, + to_backend=to_backend, + ) + + if occupations is not None or site_charge is not None: + raise ValueError("occupations and site_charge require fermion=...") + if to_backend is not None: + raise ValueError("to_backend requires fermion=...") + if cyclic and not isinstance(cyclic, bool): + raise TypeError("cyclic must be a boolean for id_to_pepo.") + pepo = qtn.PEPO.rand(Lx=lx, Ly=ly, bond_dim=1, seed=666, dtype=dtype) eye = np.eye(phys_dim, dtype=dtype) @@ -94,6 +211,8 @@ def id_to_pepo(lx, ly, phys_dim=2, dtype="complex128", chi=1, rand_strength=0.0) data[tuple([0] * n_virt)] = eye tensor.modify(data=data) + if cyclic: + pepo = add_cycle(pepo, bond_dim=cycle_bond_dim) if chi > 1: pepo.expand_bond_dimension_(chi, rand_strength=rand_strength) return pepo diff --git a/src/pepsy/tensors/symmetric.py b/src/pepsy/tensors/symmetric.py index 9f8a040..e2dd14f 100644 --- a/src/pepsy/tensors/symmetric.py +++ b/src/pepsy/tensors/symmetric.py @@ -5174,6 +5174,20 @@ def _dense_numpy(value, *, dtype=None): return np.asarray(value, dtype=dtype) +def _is_single_site_identity_hamiltonian(target, local_dim, zero_charge): + """Return whether ``target`` is exactly one full local identity term.""" + if len(target.terms) != 1: + return False + term = next(iter(target.terms.values())) + if getattr(term, "charge", None) != zero_charge: + return False + dense = _dense_numpy(term) + return dense.shape == (local_dim, local_dim) and np.array_equal( + dense, + np.eye(local_dim, dtype=dense.dtype), + ) + + def _expanded_index_charges(index): chargemap = getattr(index, "chargemap", None) if chargemap is None: @@ -10497,6 +10511,32 @@ def to_pepo( "Model parameters cannot be supplied with an existing " f"SymHamiltonian: {names}." ) + if ( + fermionic + and not charge_sectors + and _is_single_site_identity_hamiltonian( + target, + sum(int(size) for size in self.physical_sectors.values()), + self.zero_charge, + ) + ): + from .constructors import ( # pylint: disable=import-outside-toplevel + _native_fermion_identity_pepo, + ) + + return _native_fermion_identity_pepo( + self, + Lx, + Ly, + cyclic=cyclic, + cycle_bond_dim=cycle_bond_dim, + mapper=mapper, + max_bond=max_bond, + cutoff=cutoff, + compress=compress, + dtype=dtype, + to_backend=to_backend, + ) return target.to_pepo( Lx=Lx, Ly=Ly, diff --git a/tests/test_native_identity_pepo.py b/tests/test_native_identity_pepo.py new file mode 100644 index 0000000..e61b8a9 --- /dev/null +++ b/tests/test_native_identity_pepo.py @@ -0,0 +1,96 @@ +"""Native fermionic identity-PEPO construction checks.""" + +import numpy as np +import pytest + +import pepsy + + +@pytest.mark.parametrize( + ("spinful", "symmetry"), + [ + (False, "U1"), + (False, "Z2"), + (True, "U1"), + (True, "U1U1"), + (True, "Z2"), + (True, "Z2Z2"), + ], +) +def test_native_id_to_pepo_is_identity_on_half_filled_state(spinful, symmetry): + """The safe native identity handles every advertised fermion space.""" + pytest.importorskip("symmray") + + fermion = pepsy.Fermion( + spinful=spinful, + symmetry=symmetry, + dtype="complex128", + ) + operator = pepsy.id_to_pepo( + (2, 3), + fermion=fermion, + cyclic=True, + ) + state = pepsy.ps_to_peps( + (2, 3), + fermion=fermion, + dtype="complex128", + seed=17, + ) + + assert all( + type(tensor.data).__name__.endswith("FermionicArray") + for tensor in operator + ) + acted = operator.apply(state, contract=True, compress=False) + numerator = complex( + np.asarray((state.H & acted).contract(all, optimize="auto-hq")).item() + ) + denominator = complex( + np.asarray((state.H & state).contract(all, optimize="auto-hq")).item() + ) + assert numerator / denominator == pytest.approx(1.0, abs=1.0e-10) + + +def test_native_id_to_pepo_rejects_state_sector_controls(): + """A full identity must not be reduced to a half-filled sector.""" + pytest.importorskip("symmray") + fermion = pepsy.Fermion(spinful=True, symmetry="U1") + + with pytest.raises(ValueError, match="full local identity"): + pepsy.id_to_pepo( + 2, + 3, + fermion=fermion, + occupations=[1] * 6, + ) + + +def test_fermion_to_pepo_identity_delegates_to_safe_constructor(monkeypatch): + """The model API shares the same native identity implementation.""" + pytest.importorskip("symmray") + from pepsy.tensors import constructors + + fermion = pepsy.Fermion(spinful=True, symmetry="U1") + called = {} + + def fake_identity(model, lx, ly, **options): + called.update(model=model, lx=lx, ly=ly, options=options) + return "native-identity" + + monkeypatch.setattr( + constructors, + "_native_fermion_identity_pepo", + fake_identity, + ) + result = fermion.to_pepo( + {((0, 0),): fermion.observable("identity")}, + Lx=2, + Ly=3, + cyclic=True, + ) + + assert result == "native-identity" + assert called["model"] is fermion + assert (called["lx"], called["ly"]) == (2, 3) + assert called["options"]["cyclic"] is True From adbc981a964b781d7d71c787b9c17bfc27514550 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Fri, 31 Jul 2026 22:24:59 -0600 Subject: [PATCH 46/70] Validate native PEPO gate compatibility --- src/pepsy/operators/gates.py | 60 +++++++++++++++++++++++++++-- tests/test_gate.py | 73 +++++++++++++++++++++++++++++++++++- 2 files changed, 128 insertions(+), 5 deletions(-) diff --git a/src/pepsy/operators/gates.py b/src/pepsy/operators/gates.py index b6b60e4..a0ba780 100644 --- a/src/pepsy/operators/gates.py +++ b/src/pepsy/operators/gates.py @@ -3160,6 +3160,61 @@ def _native_identity_mpo(length, info, *, max_bond, cutoff): ) +def _validate_native_pepo_compatibility(pepo, info): + """Validate a supplied native PEPO before applying native gates to it.""" + from ..tensors.symmetric import ( # pylint: disable=import-outside-toplevel + _expanded_index_charges, + ) + + tensors = tuple(getattr(pepo, "tensors", ())) + if not tensors: + raise TypeError( + "Native FermionicArray gates require a non-empty native Symmray PEPO." + ) + + expected_symmetry = info["symmetry"] + expected_phys_map = tuple(info["phys_map"]) + for tensor in tensors: + data = tensor.data + if not ( + _is_block_sparse_array(data) + and "FermionicArray" in type(data).__name__ + and bool(getattr(data, "fermionic", False)) + ): + raise TypeError( + "Native FermionicArray gates require a native fermionic " + "Symmray PEPO." + ) + actual_symmetry = str(getattr(data, "symmetry", "")) + if actual_symmetry != expected_symmetry: + raise ValueError( + "Native gate symmetry " + f"{expected_symmetry} does not match supplied PEPO symmetry " + f"{actual_symmetry}." + ) + + try: + outer_inds = set(pepo.outer_inds()) + except (AttributeError, TypeError) as exc: + raise TypeError( + "Native FermionicArray gates require a PEPO with physical outer " + "indices." + ) from exc + + for tensor in tensors: + for axis, ind in enumerate(tensor.inds): + if ind not in outer_inds: + continue + actual_phys_map = tuple( + _expanded_index_charges(tensor.data.indices[axis]) + ) + if actual_phys_map != expected_phys_map: + raise ValueError( + "Native gate physical charge maps do not match the " + "supplied PEPO physical charge maps." + ) + + def build_pepo_from_gates( gates, wheres=None, @@ -3272,10 +3327,7 @@ def build_pepo_from_gates( ) elif native_info is not None: pepo = pepo_.copy() - if any(not _is_block_sparse_array(tensor.data) for tensor in pepo): - raise TypeError( - "Native FermionicArray gates require a native Symmray PEPO." - ) + _validate_native_pepo_compatibility(pepo, native_info) else: pepo = pepo_.copy() if pepo_ is not None else id_to_pepo(Lx, Ly, dtype=dtype) if pepo_ is None and cyclic: diff --git a/tests/test_gate.py b/tests/test_gate.py index f655195..78f98d1 100644 --- a/tests/test_gate.py +++ b/tests/test_gate.py @@ -8,7 +8,7 @@ import pytest import quimb.tensor as qtn -from pepsy import hrs_to_peps, ps_to_3dpeps, ps_to_peps +from pepsy import hrs_to_peps, id_to_pepo, ps_to_3dpeps, ps_to_peps from pepsy.operators.gates import ( build_mpo_from_gates, build_pepo_from_gates, @@ -1116,6 +1116,77 @@ def test_native_fermion_gate_builders_preserve_symmetry(symmetry): assert pepo.Ly == 2 +def test_native_fermion_gate_builder_accepts_matching_supplied_pepo(): + """A supplied native PEPO with matching sectors is accepted.""" + pytest.importorskip("symmray") + from pepsy.tensors.symmetric import Fermion + + fermion = Fermion(spinful=True, symmetry="U1") + base = id_to_pepo((2, 2), fermion=fermion, cyclic=True) + result = build_pepo_from_gates( + fermion.hopping_gate(0.01, t=1.0), + where=((0, 0), (1, 0)), + mapper=OneDMap(2, 2, mode="snake-row-major"), + pepo_=base, + max_bond=8, + contract="split", + ) + + assert result.Lx == 2 + assert result.Ly == 2 + assert all( + type(tensor.data).__name__.endswith("FermionicArray") + for tensor in result + ) + + +@pytest.mark.parametrize( + ("gate_symmetry", "pepo_symmetry"), + [("U1U1", "U1"), ("U1", "U1U1")], +) +def test_native_fermion_gate_builder_rejects_symmetry_mismatch( + gate_symmetry, + pepo_symmetry, +): + """Mismatched native Abelian groups fail before gate contraction.""" + pytest.importorskip("symmray") + from pepsy.tensors.symmetric import Fermion + + gate_fermion = Fermion(spinful=True, symmetry=gate_symmetry) + pepo_fermion = Fermion(spinful=True, symmetry=pepo_symmetry) + base = id_to_pepo((2, 2), fermion=pepo_fermion, cyclic=True) + + with pytest.raises(ValueError, match="does not match supplied PEPO symmetry"): + build_pepo_from_gates( + gate_fermion.hopping_gate(0.01, t=1.0), + where=((0, 0), (1, 0)), + mapper=OneDMap(2, 2, mode="snake-row-major"), + pepo_=base, + max_bond=8, + contract="split", + ) + + +def test_native_fermion_gate_builder_rejects_physical_map_mismatch(): + """A same-group spinful gate cannot act on a spinless native PEPO.""" + pytest.importorskip("symmray") + from pepsy.tensors.symmetric import Fermion + + gate_fermion = Fermion(spinful=True, symmetry="U1") + pepo_fermion = Fermion(spinful=False, symmetry="U1") + base = id_to_pepo((2, 2), fermion=pepo_fermion, cyclic=True) + + with pytest.raises(ValueError, match="physical charge maps"): + build_pepo_from_gates( + gate_fermion.hopping_gate(0.01, t=1.0), + where=((0, 0), (1, 0)), + mapper=OneDMap(2, 2, mode="snake-row-major"), + pepo_=base, + max_bond=8, + contract="split", + ) + + @pytest.mark.parametrize("symmetry", ["U1U1", "U1", "Z2"]) def test_native_charged_gate_builders_require_opt_in(symmetry): """Charged native gate streams work when explicitly enabled.""" From 5296454259c6ff4369a046d1abcc3a397ce67d36 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Sat, 1 Aug 2026 08:54:27 -0600 Subject: [PATCH 47/70] Add tree energy optimization support --- docs/api/optimizers/energy.md | 27 ++- docs/api/optimizers/tree.md | 11 +- src/pepsy/optimizers/energy/tree.py | 285 ++++++++++++++++++++++++++-- tests/test_energy_tree.py | 102 ++++++++++ 4 files changed, 399 insertions(+), 26 deletions(-) create mode 100644 tests/test_energy_tree.py diff --git a/docs/api/optimizers/energy.md b/docs/api/optimizers/energy.md index d74b4f5..9ebabc7 100644 --- a/docs/api/optimizers/energy.md +++ b/docs/api/optimizers/energy.md @@ -43,10 +43,11 @@ the truncation cap. ## Tree tensor networks -``TreeEnergyOptimizer`` mirrors the ``MpsEnergyOptimizer`` measurement surface -for a :class:`~pepsy.optimizers.tree.TreeTensorNetwork`. It reports +``TreeEnergyOptimizer`` mirrors the ``MpsEnergyOptimizer`` energy surface for +a :class:`~pepsy.optimizers.tree.TreeTensorNetwork`. It reports ``sum_i / `` term by term using the tree's own exact, -fermion-safe contraction, and returns the same :class:`EnergyEstimate`: +fermion-safe contraction, returns the same :class:`EnergyEstimate`, and can +optimize the tree tensors through Quimb's autodiff ``TNOptimizer``: ```python import pepsy @@ -56,12 +57,26 @@ estimate = pepsy.TreeEnergyOptimizer( terms=hamiltonian, # {where: operator} mapping or a SymHamiltonian energy_per_site=True, ).energy() + +optimizer = pepsy.TreeEnergyOptimizer(tree_state, terms=hamiltonian) +tree_state, losses = optimizer.optimize( + n=100, + autodiff_backend="torch", + optimizer="adam", + progbar=False, + return_losses=True, +) ``` The terms are dispatched through :meth:`~pepsy.optimizers.tree.TreeTensorNetwork.local_expectations`, which shares one contraction optimiser across every term (pass a reusable ``pepsy.build_optimizer(...)`` as ``contraction_opt`` to cache paths across -same-topology contractions) and reuses the memoized graded norm, so the -result is identical to summing per-term -:meth:`~pepsy.optimizers.tree.TreeTensorNetwork.local_expectation` calls. +same-topology contractions) and reuses the memoized graded norm for ordinary +readout. During autodiff optimization, Quimb injects tensor arrays below the +TTN mutation hooks, so the loss instead sums unnormalized numerators and +divides by a freshly contracted full-tree norm on every call. Afterward the +canonical metadata is invalidated rather than rebuilt around an arbitrary +post-optimization centre; native fermionic normalized readouts therefore stay +gauge invariant. The returned state remains a ``TreeTensorNetwork`` and the +scalar history is available as ``optimizer.losses``. diff --git a/docs/api/optimizers/tree.md b/docs/api/optimizers/tree.md index d2f72db..7ecf2da 100644 --- a/docs/api/optimizers/tree.md +++ b/docs/api/optimizers/tree.md @@ -346,8 +346,15 @@ site or a tuple of sites) to its operator. It delegates each term to `pepsy.build_optimizer(...)` caches one contraction path per topology, and it reuses the memoized graded norm across the batch. Each returned value matches the corresponding single-term call exactly. For a Hamiltonian-level energy -readout, `pepsy.TreeEnergyOptimizer` wraps this batch path and returns an -`EnergyEstimate` mirroring `MpsEnergyOptimizer`. +readout or variational energy optimization, `pepsy.TreeEnergyOptimizer` wraps +this batch path, returns an `EnergyEstimate` mirroring `MpsEnergyOptimizer`, +and exposes the corresponding `make_tn_optimizer` / `optimize` methods. +Optimization updates the tree tensor parameters with Quimb's autodiff +`TNOptimizer` while retaining the exact tree local-expectation objective. For +ordinary readout the native graded norm is memoized. The optimization loss +uses a fresh full doubled-tree denominator because Quimb's direct parameter +injection cannot invalidate that cache; the post-optimization state is marked +non-canonical rather than recanonicalized around an arbitrary centre. For the package-level product-state constructor, matching `ps_to_mps`, use `pepsy.ps_to_ttn(n, theta=..., tree=...)`. It builds the requested tree, diff --git a/src/pepsy/optimizers/energy/tree.py b/src/pepsy/optimizers/energy/tree.py index d66cddd..5be0093 100644 --- a/src/pepsy/optimizers/energy/tree.py +++ b/src/pepsy/optimizers/energy/tree.py @@ -1,30 +1,35 @@ -"""Native tree-tensor-network energy measurement. +"""Native tree-tensor-network energy measurement and optimization. -This mirrors the *measurement* surface of +This mirrors the energy surface of :class:`~pepsy.optimizers.energy.MpsEnergyOptimizer` for a :class:`~pepsy.optimizers.tree.TreeTensorNetwork`. The energy is the sum of per-term local expectations ``/`` evaluated with the tree's own graded (fermion-safe) contraction, reusing a single contraction -optimiser and the cached norm across every term. It is deliberately thin: the -tree gate-stream evolution/optimization lives in -:class:`~pepsy.optimizers.tree.TreeOptimizer`; this class only reports energy. +optimiser and the cached norm across every term. Optimization uses Quimb's +``TNOptimizer`` over the tree tensors while retaining that same exact local +expectation objective. The public measurement path can cache the norm, but +the autodiff optimization path recomputes the full norm on every call: Quimb +injects tensor parameters directly and that mutation is outside the TTN's +cache-invalidation hooks. """ from __future__ import annotations from collections.abc import Mapping from typing import Any +import warnings import autoray as ar +from ...backends import infer_backend_converter_from_sample from ...tensors import build_optimizer -from .peps import EnergyEstimate +from .peps import EnergyEstimate, PepsEnergyOptimizer __all__ = ["TreeEnergyOptimizer"] -class TreeEnergyOptimizer: - """Report term-by-term local energy for a :class:`TreeTensorNetwork`. +class TreeEnergyOptimizer(PepsEnergyOptimizer): + """Measure and optimize term-by-term energy on a tree state. The objective is the normalized local expectation ``sum_i / ``. Each local term is evaluated with the @@ -62,6 +67,7 @@ def __init__( energy_per_site: bool = True, real: bool = True, contraction_opt: Any = "auto-hq", + loss_kwargs: Mapping[str, Any] | None = None, ): if hamiltonian is not None and terms is not None: raise TypeError("pass either hamiltonian or terms, not both") @@ -75,6 +81,9 @@ def __init__( "real": bool(real), "contraction_opt": self._resolve_optimize(contraction_opt), } + self.losses: list[float] = [] + if loss_kwargs is not None: + self.set_loss_kwargs(**dict(loss_kwargs)) # -- state / term / option resolution ------------------------------------- @@ -82,13 +91,19 @@ def __init__( def _as_tree_state(state): if hasattr(state, "local_expectations") and hasattr(state, "plan"): return state - inner = getattr(state, "p", None) - if inner is not None and hasattr(inner, "local_expectations"): - return inner + for name in ("tn", "p"): + inner = getattr(state, name, None) + if inner is not None and hasattr(inner, "local_expectations"): + return inner raise TypeError( "state must be a TreeTensorNetwork with local_expectations()." ) + # PepsEnergyOptimizer's shared optimization helpers use this hook for + # state validation. Keep the tree-specific validation while reusing its + # backend conversion and finite-gradient machinery. + _as_peps_state = staticmethod(_as_tree_state) + @staticmethod def _terms_from_hamiltonian(source): if source is None: @@ -154,19 +169,126 @@ def _maybe_real(value): except Exception: # pragma: no cover - defensive for scalar types return getattr(value, "real", value) - # -- energy evaluation ---------------------------------------------------- + @classmethod + def _terms_for_state_backend(cls, terms, state): + """Convert dense operator terms to the live tree backend.""" + sample = cls._sample_array_from_tn(state) + try: + converter = infer_backend_converter_from_sample(sample) + except (ImportError, TypeError, ValueError): + converter = None + if converter is None: + return terms + return { + where: cls._convert_term_array(operator, converter) + for where, operator in dict(terms).items() + } - def _total_energy(self, state, terms, opts): + @classmethod + def _loss_state( + cls, + state, + *, + terms, + normalized=True, + energy_per_site=True, + real=True, + contraction_opt="auto-hq", + ): + """Evaluate the differentiable tree energy objective.""" + state = cls._as_tree_state(state) + terms = cls._terms_from_hamiltonian(terms) + terms = cls._terms_for_state_backend(terms, state) + if contraction_opt is None: + contraction_opt = build_optimizer(progbar=False) values = state.local_expectations( terms, - optimize=opts["contraction_opt"], - normalized=bool(opts["normalized"]), + optimize=contraction_opt, + normalized=bool(normalized), ) - total = sum((complex(v) for v in values.values()), 0j) - if opts["real"]: - return self._maybe_real(total) + values = tuple(values.values()) + total = 0.0 if not values else values[0] + for value in values[1:]: + total = total + value + if real: + total = cls._maybe_real(total) + if energy_per_site: + total = total / cls._num_sites(state) return total + @classmethod + def _optimization_loss_state( + cls, + state, + *, + terms, + normalized=True, + energy_per_site=True, + real=True, + contraction_opt="auto-hq", + ): + """Evaluate the optimization objective with a fresh norm. + + ``TNOptimizer`` updates tensor arrays through ``apply_to_arrays``. + That is deliberately a low-level operation and cannot invalidate the + TTN's memoized native-fermion norm (or its canonical-region metadata). + Calling ``local_expectation(..., normalized=True)`` here would + therefore divide a newly injected state by the previous state's norm. + Compute every term unnormalized, then divide by a fresh full doubled + tree contraction. This is the gauge-invariant Rayleigh quotient and + remains differentiable through the live tensor backend. + """ + state = cls._as_tree_state(state) + terms = cls._terms_from_hamiltonian(terms) + terms = cls._terms_for_state_backend(terms, state) + if contraction_opt is None: + contraction_opt = build_optimizer(progbar=False) + + values = tuple( + state.local_expectation( + operator, + where, + optimize=contraction_opt, + normalized=False, + ) + for where, operator in terms.items() + ) + total = 0.0 if not values else values[0] + for value in values[1:]: + total = total + value + if normalized: + denominator = (state.H | state).contract( + all, + optimize=contraction_opt, + ) + total = total / denominator + if real: + total = cls._maybe_real(total) + if energy_per_site: + total = total / cls._num_sites(state) + return total + + @staticmethod + def _tnopt_loss(state, *, terms, **loss_kwargs): + """Adapter for :class:`quimb.tensor.TNOptimizer`.""" + return TreeEnergyOptimizer._optimization_loss_state( + state, + terms=terms, + **loss_kwargs, + ) + + # -- energy evaluation ---------------------------------------------------- + + def _total_energy(self, state, terms, opts): + return self._loss_state( + state, + terms=terms, + normalized=opts["normalized"], + energy_per_site=False, + real=opts["real"], + contraction_opt=opts["contraction_opt"], + ) + def _resolve(self, state, hamiltonian, terms, kwargs): state = self.state if state is None else self._as_tree_state(state) terms_use = self.terms @@ -204,3 +326,130 @@ def energy(self, state=None, *, hamiltonian=None, terms=None, **kwargs): "contraction_opt": opts["contraction_opt"], }, ) + + def set_loss_kwargs(self, **kwargs): + """Update the defaults used by :meth:`loss` and :meth:`optimize`.""" + self.loss_kwargs.update(self._pick_loss_kwargs(kwargs)) + return self + + def normalize(self, state=None, **kwargs): + """Normalize a tree state in place using its canonical norm path.""" + state = self.state if state is None else self._as_tree_state(state) + normalize = getattr(state, "normalize", None) + if not callable(normalize): + raise TypeError("state must provide normalize() for tree normalization.") + normalize(**kwargs) + return state + + @staticmethod + def _recanonicalize(state): + """Invalidate metadata after direct TN parameter updates. + + Arbitrary tensor-array updates are not guaranteed to admit an exact + canonicalization through the old centre. In particular, doing so for + native fermionic arrays can create a centre-norm shortcut inconsistent + with the updated state. Leave the post-optimization state + non-canonical and force subsequent normalized fermionic readouts down + the exact full-network norm path. + """ + plan = getattr(state, "plan", None) + node_tensor = getattr(state, "node_tensor", None) + if plan is not None and callable(node_tensor): + # TNOptimizer updates tensor data directly. Clear Quimb's local + # ``left_inds`` proofs as well as the TTN's region marker, or a + # subsequent canonicalization can incorrectly trust orientations + # inherited from the pre-optimization state. + for node in plan.nodes(): + node_tensor(node).modify(left_inds=None) + invalidate = getattr(state, "invalidate_canonical_form", None) + if callable(invalidate): + invalidate() + return state + + def make_tn_optimizer( + self, + *, + loss_kwargs: Mapping[str, Any] | None = None, + loss_constants: Mapping[str, Any] | None = None, + autodiff_backend: str = "torch", + optimizer: str = "adam", + progbar: bool = True, + jit_fn: bool = False, + device: str = "cpu", + **tnopt_kwargs, + ): + """Construct a Quimb ``TNOptimizer`` for the live tree state. + + The shared PEPS implementation supplies backend conversion and the + Quimb optimizer construction; this override exists as the explicit + tree API and documents that the objective is the tree loss adapter. + """ + return super().make_tn_optimizer( + loss_kwargs=loss_kwargs, + loss_constants=loss_constants, + autodiff_backend=autodiff_backend, + optimizer=optimizer, + progbar=progbar, + jit_fn=jit_fn, + device=device, + **tnopt_kwargs, + ) + + def optimize( + self, + *, + n=220, + loss_kwargs: Mapping[str, Any] | None = None, + loss_constants: Mapping[str, Any] | None = None, + autodiff_backend: str = "torch", + optimizer: str = "adam", + progbar: bool = True, + jit_fn: bool = False, + device: str = "cpu", + return_losses: bool = False, + normalize: bool = False, + normalize_kwargs: Mapping[str, Any] | None = None, + check_finite_gradient: bool = True, + **optimize_kwargs, + ): + """Optimize the tree tensors against the configured energy objective.""" + merged_loss_kwargs = dict(self.loss_kwargs) + merged_loss_kwargs.update(self._pick_loss_kwargs(loss_kwargs)) + tnopt = self.make_tn_optimizer( + loss_kwargs=merged_loss_kwargs, + loss_constants=loss_constants, + autodiff_backend=autodiff_backend, + optimizer=optimizer, + progbar=progbar, + jit_fn=jit_fn, + device=device, + ) + if check_finite_gradient: + finite_gradient, finite_loss = self._initial_gradient_status(tnopt) + if not finite_gradient: + self.losses = [] if finite_loss is None else [float(finite_loss)] + warnings.warn( + "Tree energy autodiff produced a non-finite initial " + "gradient; returning the unmodified state.", + RuntimeWarning, + stacklevel=2, + ) + if return_losses: + return self.state, tuple(self.losses) + return self.state + + out = tnopt.optimize(n=n, **optimize_kwargs) + self.losses = list(getattr(tnopt, "losses", ())) + out = self._state_for_autodiff_backend( + out, + self.state, + autodiff_backend, + device=device, + ) + out = self._recanonicalize(self._as_tree_state(out)) + if normalize: + out = self.normalize(out, **dict(normalize_kwargs or {})) + self.state = self._as_tree_state(out) + if return_losses: + return self.state, tuple(self.losses) + return self.state diff --git a/tests/test_energy_tree.py b/tests/test_energy_tree.py new file mode 100644 index 0000000..52ff4ad --- /dev/null +++ b/tests/test_energy_tree.py @@ -0,0 +1,102 @@ +"""Tests for tree-network energy measurement and optimization.""" + +import numpy as np +import pytest + +import pepsy + + +def _tree_state(): + plan = pepsy.TreePlan.from_order(range(2), structure="balanced") + return pepsy.TreeTensorNetwork.from_plan(plan, dtype=complex) + + +def test_tree_energy_optimizer_reports_energy(): + state = _tree_state() + z = np.diag([1.0, -1.0]).astype(complex) + optimizer = pepsy.TreeEnergyOptimizer( + state, + terms={0: -z, 1: -z}, + energy_per_site=False, + contraction_opt="auto-hq", + ) + + estimate = optimizer.energy() + + assert estimate.energy == pytest.approx(-2.0) + assert estimate.energy_per_site == pytest.approx(-1.0) + assert estimate.boundary_mode == "exact" + + +@pytest.mark.filterwarnings("ignore:The contraction tree is not a compressed one") +def test_tree_energy_optimizer_supports_tn_optimization(): + pytest.importorskip("torch") + state = _tree_state() + x = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex) + optimizer = pepsy.TreeEnergyOptimizer( + state, + terms={0: -x, 1: -x}, + energy_per_site=False, + contraction_opt="auto-hq", + ) + before = float(optimizer.energy().energy) + + out, losses = optimizer.optimize( + n=1, + autodiff_backend="torch", + progbar=False, + return_losses=True, + ) + + assert isinstance(out, pepsy.TreeTensorNetwork) + assert out.validate(check_canonical=True) is out + assert losses + assert np.isfinite(float(losses[-1])) + after = float(optimizer.energy().energy) + assert np.isfinite(after) + assert after < before + + +def test_tree_energy_optimizer_fermion_tnopt_refreshes_norm(): + """Autodiff on native U1 trees uses a fresh Rayleigh-quotient norm.""" + pytest.importorskip("symmray") + pytest.importorskip("torch") + + fermion = pepsy.Fermion( + spinful=True, + symmetry="U1", + dtype="complex128", + ) + plan = pepsy.TreePlan.from_order(range(2), structure="balanced") + state = pepsy.ps_to_ttn( + 2, + tree=plan, + fermion=fermion, + occupations=((1, 0), (0, 1)), + dtype="complex128", + ) + optimizer = pepsy.TreeEnergyOptimizer( + state, + terms={0: -fermion.observable("sx"), 1: -fermion.observable("sx")}, + normalized=True, + energy_per_site=False, + contraction_opt="auto-hq", + ) + before = float(optimizer.energy().energy) + + out, losses = optimizer.optimize( + n=8, + optimizer="l-bfgs-b", + autodiff_backend="torch", + progbar=False, + return_losses=True, + ) + + values = np.asarray(losses, dtype=float) + assert isinstance(out, pepsy.TreeTensorNetwork) + assert values.size + assert np.all(np.isfinite(values)) + assert np.min(values) >= -2.0 + assert values[-1] <= before + 1.0e-10 + assert out.orthogonality_center is None + assert np.isfinite(float(optimizer.energy().energy)) From a0c2dae58e2db101ca0c77f0236eccec31c38c89 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Sat, 1 Aug 2026 16:22:51 -0700 Subject: [PATCH 48/70] Improve tree optimizer layout and canonicalization --- .github/skills/tree-optimizer/SKILL.md | 28 +- .../references/performance-layout.md | 10 +- docs/api/optimizers/tree.md | 68 ++- src/pepsy/optimizers/tree/layout.py | 48 ++- src/pepsy/optimizers/tree/optimizer.py | 402 ++++++++++++++---- src/pepsy/optimizers/tree/ttn.py | 82 +++- src/pepsy/vmc/netket.py | 68 ++- tests/test_optimize_tree.py | 389 ++++++++++++++++- 8 files changed, 955 insertions(+), 140 deletions(-) diff --git a/.github/skills/tree-optimizer/SKILL.md b/.github/skills/tree-optimizer/SKILL.md index f5df2a4..5adaeb2 100644 --- a/.github/skills/tree-optimizer/SKILL.md +++ b/.github/skills/tree-optimizer/SKILL.md @@ -157,11 +157,13 @@ telescopes to identity between bra and ket. an unconditional O(N) recanonicalisation. - Local isometry proofs live only on each tensor's ``left_inds``. `TreeTensorNetwork.isometry_direction` / `isometry_map` derive read-only - orientations, `can_skip_canonize` recognizes an already-proven dense edge, - and `validate_isometry_metadata` checks alignment with the canonical region. + orientations, `can_skip_canonize` recognizes an already-proven dense edge or + a native Symmray edge with aligned charge maps, and + `validate_isometry_metadata` checks alignment with the canonical region. `TreeOptimizer` delegates these methods; never add a second mutable - optimizer-owned orientation map. Native fermionic edges always retain their - explicit graded QR and are never skipped through this dense metadata path. + optimizer-owned orientation map. Native fermionic edges fall back to + explicit graded QR when the proof is absent or malformed; positive-cutoff or + over-cap native compression still uses the explicit graded SVD. - `ttn.is_canonical_form(center)` verifies the invariant directly (every non-centre tensor is an isometry toward the centre) — use it in tests/diagnostics. - A freshly built product state is **already canonical at the root** (all @@ -281,15 +283,15 @@ covering range then compressed (quimb's `gate_with_submpo` is `MatrixProductStat parent together with the old state/operator bonds. No dense state tensor for the whole Steiner subtree is formed; the last node is the hub. 5. Install every routed Q factor with its ``left_inds`` isometry metadata. - Dense trees can then recover the hub centre through the normal canonical - state machine without repeating those QRs; native fermionic trees retain - their explicit graded QR recovery. Finally make one depth-first canonical - SVD sweep: every affected tree edge is truncated once, after the complete - operator has arrived. Dense path and subtree sweeps select one-sided - ``reduced="left"`` compression only when the destination tensor's live - ``left_inds`` proves the required isometry; missing proofs and native - graded tensors use the full reduction. `renormalize=True` renormalises - afterwards (for Kraus/projection). + Dense trees and charge-aligned native Symmray trees can then recover the + hub centre through the normal canonical state machine without repeating + those QRs; missing or malformed native proofs use explicit graded QR. + Finally make one depth-first canonical SVD sweep: every affected tree edge + is truncated once, after the complete operator has arrived. Dense path and + subtree sweeps select one-sided ``reduced="left"`` compression only when + the destination tensor's live ``left_inds`` proves the required isometry; + native graded compression keeps its explicit block-SVD semantics. + `renormalize=True` renormalises afterwards (for Kraus/projection). State bonds are always read from the live tensors because gate application can rename them. New state message bonds are fresh per-update names, while operator diff --git a/.github/skills/tree-optimizer/references/performance-layout.md b/.github/skills/tree-optimizer/references/performance-layout.md index 5148826..bfbbf31 100644 --- a/.github/skills/tree-optimizer/references/performance-layout.md +++ b/.github/skills/tree-optimizer/references/performance-layout.md @@ -19,10 +19,12 @@ Tree Optimizer skill so the upload-facing `SKILL.md` stays concise. canonicalization kernel. Path and subtree compression also reads that proof before selecting one-sided `reduced="left"` compression, avoiding the redundant reduction QR only when the destination tensor is proven - isometric. Missing proofs fall back to two-sided reduction. The network - derives orientation views directly from live tensors; do not cache a - duplicate map in the optimizer. Native fermionic routing deliberately - retains explicit graded QR/SVD recovery. + isometric. Missing proofs fall back to two-sided reduction. Native Symmray + routing preserves the same proof when its charge maps are aligned; native + canonical recovery skips only that proven lossless QR, while truncating + native compression remains an explicit graded SVD. The network derives + orientation views directly from live tensors; do not cache a duplicate map + in the optimizer. - `copy()` shares the immutable `TreePlan`, owns `self.tn.copy()`, resets the tid cache, and derives a deterministic child seed for an independent RNG. diff --git a/docs/api/optimizers/tree.md b/docs/api/optimizers/tree.md index 7ecf2da..cf1529a 100644 --- a/docs/api/optimizers/tree.md +++ b/docs/api/optimizers/tree.md @@ -110,12 +110,13 @@ and `is_canonical_form(center)` delegate to the state, so the optimizer and its Local isometry orientation also has one owner: each live Quimb tensor carries its proven `left_inds`, while `TreeTensorNetwork.isometry_direction(node)` and `isometry_map()` derive read-only node-to-neighbour views from those tensors. -`can_skip_canonize(a, b)` exposes the exact dense-edge condition used to avoid -an already-proven QR, and `validate_isometry_metadata()` checks the local +`can_skip_canonize(a, b)` exposes the exact local condition used to avoid an +already-proven QR, and `validate_isometry_metadata()` checks the local orientations against the tracked canonical region. `TreeOptimizer` delegates the same four methods without maintaining another mutable map. Native -fermionic trees retain explicit graded QR and therefore never report a -skippable edge through this API. +fermionic edges use this shortcut only when Symmray reports a fermionic array +with aligned charge maps and a complete `left_inds` proof; otherwise they +retain the explicit graded QR path. `TreeTensorNetwork.validate()` checks the live tensor set, physical legs, tree edges, and bond ownership against the `TreePlan`; pass @@ -174,10 +175,12 @@ nodes. Application then proceeds recursively from subtree leaves to a hub: each local state/operator message is losslessly QR-split on one edge and absorbed by its parent, carrying every still-open operator virtual leg. No dense state tensor -for the whole Steiner subtree is formed. Each dense routed Q tensor retains its -`left_inds` isometry metadata, so canonical recovery recognizes that it already -points toward the hub instead of repeating the same QR; native fermionic trees -retain explicit graded QR recovery. Once all MPO factors have arrived, every +for the whole Steiner subtree is formed. Each routed Q tensor, including native +Symmray graded Q factors, retains its `left_inds` isometry metadata when +available, so canonical recovery recognizes that it already points toward the +hub instead of repeating the same QR. The native predicate additionally +validates charge-map alignment before skipping. Once all MPO factors have +arrived, every touched edge is compressed once. A bond that remains within its configured `max_bond` uses a lossless QR, avoiding repeated cutoff loss of tiny state components across successive sub-MPO events; when the MPO expands an edge @@ -410,8 +413,11 @@ For circuits with gates of different operator-Schmidt ranks, use congestion-aware, and balanced candidates using the predicted log bond growth on every edge. A gate crossing an edge contributes `log2(k)`, where `k` is its operator-Schmidt rank across that edge; the maximum edge load therefore -predicts the worst-case multiplicative bond growth. The default -`objective="path"` remains the co-occurrence/path-length heuristic. +predicts the worst-case multiplicative bond growth. `TreeOptimizer` uses +`layout_objective="congestion"` by default because it is a better +execution-oriented choice at finite `chi`; a bare `TreeLayoutFinder` retains +`objective="path"` as its fast, backward-compatible default. The path +objective remains the co-occurrence/path-length heuristic. `objective="hybrid"` is useful when both replay cost and bond pressure matter: it combines normalized path score, maximum edge load, and total edge load with `hybrid_weights=(path, max_edge_load, total_edge_load)`. The @@ -431,16 +437,25 @@ set explicit budgets for a larger search. Dense operators wider than For a whole-tree optimization, use `objective="full_tree"` (also accepted as `"tree"` or `"cotengra"`). This evaluates dynamic operator-Schmidt demand, -working tensor width, estimated work/write volume, and route length across -every hierarchical scale, not only the root cut. It enables bounded subtree -reconfiguration and simulated annealing by default; override these with -`topology_refine="subtree"`, `topology_budget=`, `search="anneal"`, and +cap overflow, working tensor width, estimated work/write volume, and route +length across every hierarchical scale, not only the root cut. Finite-`chi` +overflow and edge demand are ranked before tensor-work proxies, since avoiding +unnecessary truncation is the primary execution concern. It enables bounded +subtree reconfiguration and simulated annealing by default; override these +with `topology_refine="subtree"`, `topology_budget=`, `search="anneal"`, and `search_budget=`. The result is still a cheap layout proxy rather than a real TTN replay, so the state-aware pilot remains the final accuracy check. The default `chi=None` leaves this as a static, chi-blind objective; supplying `chi` only adds cap-aware ranking and does not change the no-tensor nature of layout discovery. +For the 6×6 periodic square-lattice calibration stream (Hadamards followed by +periodic controlled-phase gates), predicted total overflow ranked binary, +ternary, and four-way candidates in the same order as actual capped replay +pressure and truncation counts. This validates the profile as a layout-ranking +proxy; use the state-aware pilot when the circuit has strong cancellations or +state-dependent rank loss. + Use `order="quality"` with `finder.run()` (or set it on the finder) for the MPS-style high-quality offline search. Quality mode now means `objective="full_tree"`: it evaluates every hierarchy scale, enables bounded @@ -693,6 +708,7 @@ choice = opt.optimize_layout( rounds=2, pilot_candidates=4, pilot_steps=64, + pilot_workers=2, topology_budget=32, search_budget=64, ) @@ -707,6 +723,11 @@ edges. `objective="full_tree"` combines all-scale static work/bond estimates with this short state-aware replay. `install=True` remounts the product state on the final plan; it remains rejected for an entangled state. +Independent product-state pilots can be evaluated concurrently with +pilot_workers greater than one; the default is one for minimal overhead and +deterministic resource use. Candidate order and tie-breaking remain +deterministic. + Both helpers are also available from the package-level API: ```python @@ -924,7 +945,11 @@ scale. the global spectrum; native Symmray states compare the full and actually retained charge-block spectra using the same sector-aware truncation rule as the live update. Spectrum probes are opt-in because they add local SVD work - per truncation edge. The report also contains gate-level `updates`, grouping + per truncation edge. Enabling it emits a one-time warning because the + diagnostic spectrum probes can add substantial SVD work. It remains + disabled by default. Lossless zero-cutoff edges that are already within + their bond cap use QR and do not probe a spectrum even when tracking is on. + The report also contains gate-level `updates`, grouping edge events by support and reporting the cumulative relative loss. - `TreeOptimizer.convergence_sweep(gates, n, chi_values, ops=...)` replays the stream at several `chi` on one fixed tree and returns per-`chi` `max_bond`, @@ -967,6 +992,19 @@ available through `truncation_report()`, `get_infidelities()`, and ## Performance and stability +- **Lossless QR fast paths.** Zero-cutoff splits and edge updates whose + rank is already within the active bond cap use QR rather than SVD. This + includes the sibling-leaf split and remains valid for native graded QR. + A positive cutoff retains the existing rank-revealing compression semantics. +- **Repeated-gate cache.** Direct gate SVDs and MPO factorizations are cached + by immutable payload identity, backend signature, support, and local + dimensions. The bounded cache returns fresh-index tensor copies, so it does + not share mutable network indices with the live state. +- **Subtree and pilot parallelism.** Set `subtree_workers>1` to evaluate + independent dense leaf-to-hub QR messages in a wave, with deterministic + merging. Native fermionic routing stays serial because graded Symmray phase + bookkeeping has not been established as thread-safe. Set `pilot_workers>1` + for independent layout pilot replays; both options default to one. - **Sibling fast path.** A two-qubit gate on two leaves that share a parent is applied as a single two-site update: the two leaves and their parent are contracted into one blob, the gate is applied, and the blob is re-split by diff --git a/src/pepsy/optimizers/tree/layout.py b/src/pepsy/optimizers/tree/layout.py index a6b178f..508ce20 100644 --- a/src/pepsy/optimizers/tree/layout.py +++ b/src/pepsy/optimizers/tree/layout.py @@ -1332,8 +1332,9 @@ class TreeLayoutFinder: objective. `"hypergraph"` is the direct multi-site mode: it ranks plans from the full support hyperedges and per-edge Schmidt loads, then applies bounded leaf and binary-topology refinement by default. - `"full_tree"` evaluates dynamic bond pressure, tensor width, estimated - work, write volume, and route length across every tree scale. It is + `"full_tree"` evaluates dynamic bond pressure, predicted ``chi`` + overflow, tensor width, estimated work, write volume, and route length + across every tree scale. It is the high-quality, Cotengra-inspired mode; ``order="quality"`` selects it automatically and enables its bounded search stages. order : {None, "quality"}, optional @@ -2023,12 +2024,20 @@ def _objective_key(self, plan): if self.objective == "full_tree": profile = self.full_tree_profile(plan) return ( + # For a finite-chi optimizer, predicted cut overflow is the + # first-order performance and accuracy failure. Prefer a + # layout that stays within the cap, then distinguish the + # remaining candidates by uncapped edge demand before + # considering tensor work. With chi=None, overflow is zero + # and this naturally reduces to uncapped demand ordering. + profile["peak_overflow_log2"], + profile["total_overflow_log2"], + profile["peak_edge_demand_log2"], + profile["total_edge_demand_log2"], profile["peak_tensor_log2"], profile["peak_work_log2"], profile["log_total_write"], profile["log_total_work"], - profile["peak_edge_demand_log2"], - profile["total_edge_demand_log2"], profile["total_route_length"], self.score(plan), ) @@ -3749,9 +3758,11 @@ def full_tree_profile(self, plan=None): The profile is a cheap layout proxy, not a replacement for replaying the circuit. It accumulates uncapped operator-Schmidt demand on every - tree edge, tracks the capped working bond pressure at the configured - ``chi``, estimates tensor widths and write/work volume for every - touched node, and groups those quantities by hierarchical tree scale. + tree edge, tracks capped working-bond pressure and predicted ``chi`` + overflow at the configured ``chi``, estimates tensor widths and + write/work volume for every touched node, and groups those quantities + by hierarchical tree scale. It never allocates a TTN or performs a + tensor truncation. """ if plan is None: plan = self.run() @@ -3802,6 +3813,8 @@ def full_tree_profile(self, plan=None): "log_total_tensor_size": -np.inf, "peak_edge_demand_log2": 0.0, "total_edge_demand_log2": 0.0, + "peak_bond_log2": 0.0, + "peak_overflow_log2": 0.0, } for node, scale in node_scales.items(): scales[scale]["node_count"] += 1 @@ -3903,6 +3916,14 @@ def node_log_size(node): scale["peak_edge_demand_log2"] = max( scale["peak_edge_demand_log2"], demand_log[edge] ) + scale["peak_bond_log2"] = max( + scale["peak_bond_log2"], bond_log[edge] + ) + if np.isfinite(log_chi): + scale["peak_overflow_log2"] = max( + scale["peak_overflow_log2"], + max(0.0, demand_log[edge] - log_chi), + ) for node in plan.children: log_size = node_log_size(node) @@ -3918,8 +3939,19 @@ def node_log_size(node): scale["peak_edge_demand_log2"] = max( scale["peak_edge_demand_log2"], demand ) + scale["peak_bond_log2"] = max( + scale["peak_bond_log2"], bond_log[edge] + ) scale["total_edge_demand_log2"] += demand + overflow_log = { + edge: ( + max(0.0, demand - log_chi) + if np.isfinite(log_chi) else 0.0 + ) + for edge, demand in demand_log.items() + } + profile = { "event_count": event_count, "peak_tensor_log2": float(peak_tensor_log2), @@ -3933,6 +3965,8 @@ def node_log_size(node): "peak_edge_demand_log2": float(max(demand_log.values(), default=0.0)), "total_edge_demand_log2": float(sum(demand_log.values())), "peak_bond_log2": float(max(bond_log.values(), default=0.0)), + "peak_overflow_log2": float(max(overflow_log.values(), default=0.0)), + "total_overflow_log2": float(sum(overflow_log.values())), "total_route_length": int(total_route_length), "exact_events": int(exact_events), "bounded_events": int(bounded_events), diff --git a/src/pepsy/optimizers/tree/optimizer.py b/src/pepsy/optimizers/tree/optimizer.py index 52858e5..deba716 100644 --- a/src/pepsy/optimizers/tree/optimizer.py +++ b/src/pepsy/optimizers/tree/optimizer.py @@ -339,8 +339,8 @@ class TreeOptimizer: estimated local tensor cost at ``chi``; ``"hypergraph"`` directly scores every original multi-qubit support across every crossed tree edge and enables bounded leaf/NNI refinement by default; - ``"full_tree"`` evaluates dynamic all-scale tensor width, work, write, - bond pressure, and route costs; + ``"full_tree"`` evaluates dynamic all-scale bond pressure, overflow, + tensor width, work, write, and route costs; ``"hybrid"`` combines normalized path, peak-load, and total-load costs. Pass a configured :class:`TreeLayoutFinder` through ``layout=`` to customize its hybrid weights or enable pre-simulation refinement. @@ -452,11 +452,12 @@ def __init__(self, gates=None, n=None, *, chi=64, structure="quality", max_arity=2, top_arity=_DEFAULT_TOP_ARITY, community_frac=0.35, - star_frac=0.75, layout_objective="path", + star_frac=0.75, layout_objective="congestion", layout_weight_mode="count", layout_time_decay=None, layout_time_window=None, layout=None, tree=None, root_qubit=None, - dtype=complex, threads=1, seed=None, run=True, tn=None, + dtype=complex, threads=1, subtree_workers=1, seed=None, + run=True, tn=None, state=None, track_truncation=False, track_infidelity=True, max_intermediate_bond=None, max_operator_qubits=_DEFAULT_MAX_OPERATOR_QUBITS, @@ -644,8 +645,12 @@ def __init__(self, gates=None, n=None, *, chi=64, self.threads = None if threads is None else int(threads) if self.threads is not None and self.threads < 1: raise ValueError("threads must be positive or None.") + self.subtree_workers = self._positive_limit( + subtree_workers, "subtree_workers" + ) self.rng = np.random.default_rng(seed) self.track_truncation = bool(track_truncation) + self._track_warning_emitted = False self.track_infidelity = bool(track_infidelity) self.max_intermediate_bond = self._positive_limit( max_intermediate_bond, "max_intermediate_bond" @@ -668,6 +673,13 @@ def __init__(self, gates=None, n=None, *, chi=64, self.normalizations = [] self.projection_diagnostics = [] self._backend_conversion_warnings = set() + # Gate payloads are treated as immutable during replay. Keeping a + # small identity-keyed cache avoids repeating the same operator SVD or + # MPO factorization for repeated circuit gates without hashing/copying + # large backend arrays. Entries retain the payload object so Python id + # reuse cannot return a stale factorization. + self._gate_factor_cache = {} + self._gate_factor_cache_limit = 64 self._active_update = None self._truncation_survival = 1.0 @@ -1001,6 +1013,20 @@ def _warn_backend_conversion(self, source_signature, target_signature): stacklevel=3, ) + def _warn_track_truncation_slow(self): + """Warn once that complete-spectrum diagnostics add SVD work.""" + if self.track_truncation and not self._track_warning_emitted: + warnings.warn( + "TreeOptimizer track_truncation=True enables complete " + "singular-spectrum probes and can add extra SVDs for each " + "compressed edge; this diagnostic mode can substantially " + "slow replay. Use track_truncation=False for performance " + "runs.", + UserWarning, + stacklevel=3, + ) + self._track_warning_emitted = True + @staticmethod def _backend_converter(like): """Build one converter for a stream targeting ``like``.""" @@ -1459,7 +1485,7 @@ def isometry_map(self): return self.tn.isometry_map() def can_skip_canonize(self, a, b, *, absorb="right"): - """Whether local metadata proves this dense edge QR is redundant.""" + """Whether local metadata proves this edge QR is redundant.""" return self.tn.can_skip_canonize(a, b, absorb=absorb) def validate_isometry_metadata(self, region=None): @@ -1563,6 +1589,7 @@ def _pilot_layout_candidate( tree=plan, dtype=self.dtype, threads=self.threads, + subtree_workers=self.subtree_workers, track_truncation=True, track_infidelity=True, max_intermediate_bond=self.max_intermediate_bond, @@ -1579,6 +1606,10 @@ def _pilot_layout_candidate( trial.G = trial.G[:pilot_steps] trial.where = trial.where[:pilot_steps] trial.event_types = trial.event_types[:pilot_steps] + # Pilots intentionally enable full diagnostics, but their parent + # optimizer already owns the decision to pay that cost. Suppress the + # user-facing warning for these internal comparison replays. + trial._track_warning_emitted = True try: trial.run(progbar=progbar) except Exception as exc: # pragma: no cover - backend-specific @@ -1630,6 +1661,7 @@ def optimize_layout( pilot_candidates=4, candidate_budget=None, pilot_steps=None, + pilot_workers=1, include_quality=True, rounds=2, topology_budget=None, @@ -1678,6 +1710,12 @@ def optimize_layout( ) from exc if pilot_steps < 1: raise ValueError("pilot_steps must be a positive integer or None.") + try: + pilot_workers = int(pilot_workers) + except (TypeError, ValueError) as exc: + raise ValueError("pilot_workers must be a positive integer.") from exc + if pilot_workers < 1: + raise ValueError("pilot_workers must be a positive integer.") if not _is_product_tensor_network(self.tn): raise ValueError( "Tree layout pilots require a product initial state when " @@ -1780,15 +1818,34 @@ def optimize_layout( else: ranked = ranked_static[:pilot_candidates] - reports = {} - successful = [] - for name in ranked: - report = self._pilot_layout_candidate( - candidates[name]["plan"], + pilot_jobs = [ + (name, candidates[name]["plan"]) + for name in ranked + ] + + def run_pilot(job): + name, plan = job + return name, self._pilot_layout_candidate( + plan, objective=finder.objective, pilot_steps=pilot_steps, progbar=progbar, ) + + if pilot_workers > 1 and len(pilot_jobs) > 1: + from concurrent.futures import ThreadPoolExecutor + + with ThreadPoolExecutor( + max_workers=min(pilot_workers, len(pilot_jobs)), + thread_name_prefix="pepsy-tree-pilot", + ) as pool: + pilot_results = list(pool.map(run_pilot, pilot_jobs)) + else: + pilot_results = [run_pilot(job) for job in pilot_jobs] + + reports = {} + successful = [] + for name, report in pilot_results: reports[name] = report if report["status"] != "ok": continue @@ -1855,6 +1912,7 @@ def select_layout_for_compression( pilot_candidates=4, candidate_budget=None, pilot_steps=None, + pilot_workers=1, include_quality=True, rounds=1, topology_budget=None, @@ -1876,6 +1934,7 @@ def select_layout_for_compression( pilot_candidates=pilot_candidates, candidate_budget=candidate_budget, pilot_steps=pilot_steps, + pilot_workers=pilot_workers, include_quality=include_quality, rounds=rounds, topology_budget=topology_budget, @@ -2142,6 +2201,7 @@ def _finish_update(self): def apply_gate(self, gate, where, *, renormalize=False): """Apply a gate and aggregate its edge truncation diagnostics.""" + self._warn_track_truncation_slow() started = self._begin_update("gate", _normalize_where(where)) try: result = self._apply_gate_impl(gate, where, renormalize=renormalize) @@ -2223,7 +2283,14 @@ def run(self, gates=None, *, progbar=False, mode=None, non_unitary=False, the progress bar omits the norm-based infidelity field and avoids the per-event norm readout. Truncation-spectrum diagnostics remain controlled independently by :attr:`track_truncation`. + + Notes + ----- + track_truncation remains False by default. Enabling it is a + diagnostic mode: complete singular spectra require additional + factorization work and can substantially slow replay. """ + self._warn_track_truncation_slow() if mode is not None: requested_mode = str(mode).strip().lower() if requested_mode in {"tree", "ttn", "tree_tensor_network"}: @@ -2251,6 +2318,7 @@ def run(self, gates=None, *, progbar=False, mode=None, non_unitary=False, self.rng = np.random.default_rng(seed) if gates is not None: self.G, self.where, self.event_types = self._normalize_gate_queue(gates) + self._gate_factor_cache.clear() self._validate_event_stream_for_run() self._validate_mode_for_stream() # Prepare the executable payloads without mutating the public queue. @@ -2379,6 +2447,7 @@ def run(self, gates=None, *, progbar=False, mode=None, non_unitary=False, def set_gates(self, gates): """Replace the queued gate stream and return ``self``.""" self.G, self.where, self.event_types = self._normalize_gate_queue(gates) + self._gate_factor_cache.clear() return self def add_gates(self, gates): @@ -2387,6 +2456,7 @@ def add_gates(self, gates): self.G.extend(G_new) self.where.extend(where_new) self.event_types.extend(event_types_new) + self._gate_factor_cache.clear() return self @staticmethod @@ -2611,6 +2681,7 @@ def _apply_2q_mpo_impl( else: da = int(self.tn.ind_size(self._phys(qa))) db = int(self.tn.ind_size(self._phys(qb))) + cache_source = gate gate = self._as_gate_tensor4(gate, da, db) # ``MatrixProductOperator.from_dense`` is backend-generic: for a @@ -2619,10 +2690,26 @@ def _apply_2q_mpo_impl( # orientations. Its Tensor.split default has a nonzero cutoff, so make # the gate factorisation explicitly lossless; ``chi``/``self.cutoff`` # are applied only after the complete gate reaches the TTN path. - submpo = qtn.MatrixProductOperator.from_dense( - gate, dims=(da, db), sites=(qa, qb), L=self.n, - max_bond=None, cutoff=0.0, + cache_key = ( + "mpo", + id(cache_source), + _array_backend_signature(cache_source), + _array_backend_signature(self._state_like()), + qa, + qb, + da, + db, + self.n, ) + cached = self._gate_factor_cache.get(cache_key) + if cached is not None and cached[0] is cache_source: + submpo = cached[1] + else: + submpo = qtn.MatrixProductOperator.from_dense( + gate, dims=(da, db), sites=(qa, qb), L=self.n, + max_bond=None, cutoff=0.0, + ) + self._cache_gate_factorization(cache_key, cache_source, submpo) factors = self._two_site_mpo_factors(submpo, qa, qb) if factors is None: raise TypeError( @@ -2638,6 +2725,15 @@ def _apply_2q_mpo_impl( preserve_subcap=True, ) + def _cache_gate_factorization(self, key, source, value): + """Store one immutable gate factorization in the bounded cache.""" + if key in self._gate_factor_cache: + self._gate_factor_cache[key] = (source, value) + return + if len(self._gate_factor_cache) >= self._gate_factor_cache_limit: + self._gate_factor_cache.pop(next(iter(self._gate_factor_cache))) + self._gate_factor_cache[key] = (source, value) + def _two_site_mpo_factors(self, submpo, qa, qb, *, site_where=None): """Extract and normalize a two-site Quimb MPO into local factors. @@ -2685,15 +2781,38 @@ def _two_site_mpo_factors(self, submpo, qa, qb, *, site_where=None): if len(bonds) != 1: return None + cache_key = ( + "mpo_factors", + id(submpo), + _array_backend_signature(self._state_like()), + qa, + qb, + site_where, + self.n, + ) + cached = self._gate_factor_cache.get(cache_key) + if cached is not None and cached[0] is submpo: + raw_factors, shared_bond = cached[1] + else: + raw_factors = { + qubit: (factor.copy(), upper, lower) + for qubit, (factor, upper, lower) in raw_factors.items() + } + shared_bond = bonds[0] + self._cache_gate_factorization( + cache_key, submpo, (raw_factors, shared_bond) + ) + thread_ind = qtn.rand_uuid() factors = {} outputs = {} - for qubit, (factor, upper, lower) in raw_factors.items(): + for qubit, (factor_template, upper, lower) in raw_factors.items(): + factor = factor_template.copy() output = qtn.rand_uuid() factor.reindex_({ upper: output, lower: self._phys(qubit), - bonds[0]: thread_ind, + shared_bond: thread_ind, }) factors[qubit] = factor outputs[qubit] = output @@ -2704,6 +2823,7 @@ def _apply_2q_path_thread_impl( ): """Apply a two-qubit gate without opening another thread context.""" gate = self._as_state_backend(gate) + cache_source = gate qa = self._validate_qubit(qa) qb = self._validate_qubit(qb) if qa == qb: @@ -2712,22 +2832,12 @@ def _apply_2q_path_thread_impl( pa, pb = self._phys(qa), self._phys(qb) da, db = int(self.tn.ind_size(pa)), int(self.tn.ind_size(pb)) gate = self._as_gate_tensor4(gate, da, db) - thread_ind = qtn.rand_uuid() - output_a, output_b = qtn.rand_uuid(), qtn.rand_uuid() - gate_tensor = qtn.Tensor( - gate, inds=(output_a, output_b, pa, pb), - ) - left, right = gate_tensor.split( - left_inds=(output_a, pa), - method="svd", - cutoff=0.0, - absorb="both", - get="tensors", - bond_ind=thread_ind, + factors, outputs, thread_ind = self._cached_direct_gate_factors( + gate, cache_source, qa, qb, pa, pb, da, db, ) return self._apply_2q_factors_impl( - {qa: left, qb: right}, - {qa: output_a, qb: output_b}, + factors, + outputs, thread_ind, qa, qb, @@ -2735,6 +2845,70 @@ def _apply_2q_path_thread_impl( cutoff=cutoff, ) + def _cached_direct_gate_factors( + self, gate, source, qa, qb, pa, pb, da, db, + ): + """Return fresh-index copies of a cached direct gate factorization.""" + key = ( + "direct", + id(source), + _array_backend_signature(source), + _array_backend_signature(self._state_like()), + qa, + qb, + da, + db, + ) + cached = self._gate_factor_cache.get(key) + if cached is not None and cached[0] is source: + left_template, right_template = cached[1] + else: + template_out_a = "_pepsy_gate_out_a" + template_out_b = "_pepsy_gate_out_b" + template_in_a = "_pepsy_gate_in_a" + template_in_b = "_pepsy_gate_in_b" + template_thread = "_pepsy_gate_thread" + gate_tensor = qtn.Tensor( + gate, + inds=( + template_out_a, + template_out_b, + template_in_a, + template_in_b, + ), + ) + left_template, right_template = gate_tensor.split( + left_inds=(template_out_a, template_in_a), + method="svd", + cutoff=0.0, + absorb="both", + get="tensors", + bond_ind=template_thread, + ) + self._cache_gate_factorization( + key, source, (left_template, right_template) + ) + + thread_ind = qtn.rand_uuid() + output_a, output_b = qtn.rand_uuid(), qtn.rand_uuid() + left = left_template.copy() + right = right_template.copy() + left.reindex_({ + "_pepsy_gate_out_a": output_a, + "_pepsy_gate_in_a": pa, + "_pepsy_gate_thread": thread_ind, + }) + right.reindex_({ + "_pepsy_gate_out_b": output_b, + "_pepsy_gate_in_b": pb, + "_pepsy_gate_thread": thread_ind, + }) + return ( + {qa: left, qb: right}, + {qa: output_a, qb: output_b}, + thread_ind, + ) + def _apply_2q_factors_impl( self, factors, outputs, thread_ind, qa, qb, *, max_bond=None, cutoff=None, preserve_subcap=False, @@ -3100,10 +3274,14 @@ def _split_with_diagnostics( cutoff = self._subtree_cutoff_for_size( before_bond, max_bond=max_bond, cutoff=cutoff, ) - # An uncapped zero-cutoff routing split is provably lossless. Avoid an - # otherwise redundant full SVD spectrum probe when diagnostics are on; - # the subsequent final compression remains fully tracked. - lossless = max_bond is None and float(cutoff) == 0.0 + # A zero-cutoff split whose rank bound is already within the requested + # cap cannot truncate. Use QR in that case, including capped splits: + # the previous uncapped-only condition needlessly sent sibling leaf + # updates through a full SVD even when the cap could not bind. + lossless = ( + float(cutoff) == 0.0 + and (max_bond is None or before_bond <= max_bond) + ) full_spectrum = ( self._probe_split_spectrum( tensor, @@ -3114,11 +3292,21 @@ def _split_with_diagnostics( ) if self.track_truncation and not lossless else None ) - left, right = tensor.split( - left_inds=left_inds, method="svd", max_bond=max_bond, - cutoff=cutoff, cutoff_mode=self.cutoff_mode, - absorb="right", get="tensors", bond_ind=bond_ind, - ) + if lossless: + left, right = tensor.split( + left_inds=left_inds, + method="qr", + cutoff=0.0, + absorb="right", + get="tensors", + bond_ind=bond_ind, + ) + else: + left, right = tensor.split( + left_inds=left_inds, method="svd", max_bond=max_bond, + cutoff=cutoff, cutoff_mode=self.cutoff_mode, + absorb="right", get="tensors", bond_ind=bond_ind, + ) after_bond = self._tensor_ind_size(left, bond_ind) self._record_truncation( kind="split", edge=edge, before_bond=before_bond, @@ -3154,13 +3342,14 @@ def _compress_edge_with_diagnostics( cutoff = self.cutoff if cutoff is None else float(cutoff) bond_before = self.tn.bond(u, v) before_bond = int(self.tn.ind_size(bond_before)) - if ( - not self.track_truncation - and cutoff == 0.0 + lossless = ( + cutoff == 0.0 and (max_bond is None or before_bond <= max_bond) - ): + ) + if lossless: # No singular value can be removed in this case. A lossless QR - # still moves the centre across the edge, but avoids an SVD. + # still moves the centre across the edge, but avoids both the + # diagnostic spectrum probe and the compression SVD. self.tn.canonize_edge_(u, v, absorb="right") bond_after = self.tn.bond(u, v) self._record_truncation( @@ -3171,6 +3360,7 @@ def _compress_edge_with_diagnostics( return full_spectrum = None if self.track_truncation: + self._warn_track_truncation_slow() ta = self.tn.tensor_map[self._tid(u)] tb = self.tn.tensor_map[self._tid(v)] full_spectrum = self._probe_bond_spectrum( @@ -3200,9 +3390,11 @@ def _metadata_aware_reduction(self, u, v): Every caller compresses ``u -> v`` with ``absorb="right"``. If the live ``left_inds`` on ``v`` prove that it is already isometric toward - ``u``, Quimb can SVD only ``u`` and reuse ``v`` directly. Missing or - native graded metadata conservatively falls back to the usual - two-sided QR reduction. + ``u``, Quimb can SVD only ``u`` and reuse ``v`` directly. Native + graded metadata is accepted only after the TreeTensorNetwork charge + alignment guard; missing or malformed metadata falls back to the + usual two-sided reduction. The native compression kernel still uses + its explicit graded SVD regardless of this advisory ``reduced`` flag. """ if self.tn.can_skip_canonize(u, v, absorb="left"): return "left" @@ -3295,23 +3487,73 @@ def _qr_route_message(tensor, left_inds, *, bond_ind): def _route_subtree_messages( self, local, state_inds, operator_inds, order, *, token, + workers=None, ): """QR-route open MPO bonds from subtree leaves to their common hub.""" - for u, v in order: - state_bond = self.tn.bond(u, v) - left_inds = [ - index for index in local[u].inds - if index != state_bond and index not in operator_inds[u] + if workers is None: + workers = self.subtree_workers + workers = self._positive_limit(workers, "subtree_workers") + # Native graded contractions remain deliberately serial. The QR + # factorization is algebraically independent on a peel wave, but + # Symmray's global index/phase bookkeeping has not been proven + # thread-safe and correctness takes precedence over throughput there. + if self.tn.fermionic: + workers = 1 + + pending = list(order) + while pending: + ready = [ + (index, u, v) + for index, (u, v) in enumerate(pending) + if u not in {dst for _, dst in pending} + ] + if not ready: + raise RuntimeError( + "subtree peel order contains a cyclic message dependency." + ) + if workers == 1: + ready = ready[:1] + + def split_message(item): + index, u, v = item + state_bond = self.tn.bond(u, v) + left_inds = [ + ix for ix in local[u].inds + if ix != state_bond and ix not in operator_inds[u] + ] + new_bond = f"_ttn_mpo_route_{token}_{u}_{v}" + kept, message = self._qr_route_message( + local[u], left_inds, bond_ind=new_bond, + ) + return index, u, v, state_bond, new_bond, kept, message + + if workers > 1 and len(ready) > 1: + from concurrent.futures import ThreadPoolExecutor + + with ThreadPoolExecutor( + max_workers=min(workers, len(ready)), + thread_name_prefix="pepsy-ttn-qr", + ) as pool: + results = list(pool.map(split_message, ready)) + else: + results = [split_message(item) for item in ready] + + # Merge in peel-order order. At this point worker tensors are + # private and no destination tensor has been modified yet, so + # equal-destination messages remain deterministic. + for index, u, v, state_bond, new_bond, kept, message in sorted( + results + ): + local[u] = kept + local[v] = qtn.tensor_contract(local[v], message) + state_inds[v].discard(state_bond) + state_inds[v].add(new_bond) + operator_inds[v] = set(local[v].inds) - state_inds[v] + removed = {index for index, *_ in results} + pending = [ + edge for index, edge in enumerate(pending) + if index not in removed ] - new_bond = f"_ttn_mpo_route_{token}_{u}_{v}" - kept, message = self._qr_route_message( - local[u], left_inds, bond_ind=new_bond, - ) - local[u] = kept - local[v] = qtn.tensor_contract(local[v], message) - state_inds[v].discard(state_bond) - state_inds[v].add(new_bond) - operator_inds[v] = set(local[v].inds) - state_inds[v] def _install_routed_subtree(self, local, snodes, hub): """Install routed tensors and recover their proven hub centre. @@ -3319,28 +3561,30 @@ def _install_routed_subtree(self, local, snodes, hub): Dense routing already QR-isometrizes every peeled non-hub tensor toward ``hub``. Retaining each Q factor's ``left_inds`` lets Quimb's canonical recovery walk short-circuit those decompositions while still advancing - the canonical-region state machine honestly. Native fermionic tensors - deliberately keep the prior behavior: their graded QR recovery remains - explicit inside :class:`TreeTensorNetwork`. + the canonical-region state machine honestly. Native graded routing + retains the same metadata when Symmray supplied it; the + :class:`TreeTensorNetwork` predicate validates the charge maps and + falls back to explicit graded QR otherwise. """ - dense = not self.tn.fermionic for nid in snodes: routed = local[nid] modify_opts = { "data": routed.data, "inds": routed.inds, } - if dense: - if nid == hub: - # The accumulated operator and state norm live here. - modify_opts["left_inds"] = None - else: - if routed.left_inds is None: - raise RuntimeError( - "dense subtree routing lost QR isometry metadata " - f"for non-hub node {nid}." - ) - modify_opts["left_inds"] = routed.left_inds + if nid == hub: + # The accumulated operator and state norm live here. + modify_opts["left_inds"] = None + elif routed.left_inds is not None: + # Both dense and native QR return a Q factor with a live + # isometry proof. The network-level predicate performs the + # stricter native charge-map check when this is installed. + modify_opts["left_inds"] = routed.left_inds + elif not self.tn.fermionic: + raise RuntimeError( + "dense subtree routing lost QR isometry metadata " + f"for non-hub node {nid}." + ) self.tn.tensor_map[self._tid(nid)].modify(**modify_opts) self.tn.canonical_region = frozenset(snodes) @@ -3351,6 +3595,7 @@ def _install_routed_subtree(self, local, snodes, hub): def apply_subtree_operator(self, op, where, *, max_bond=None, cutoff=None, renormalize=False): """Apply a subtree operator and aggregate its edge truncations.""" + self._warn_track_truncation_slow() self._invalidate_state_norm_cache() started = self._begin_update("subtree", _normalize_where(where)) try: @@ -3376,6 +3621,7 @@ def apply_submpo(self, submpo, where, *, max_bond=None, cutoff=None): one compression sweep. Opaque MPO-like payloads fall back to dense :meth:`apply_subtree_operator` lowering. """ + self._warn_track_truncation_slow() self._invalidate_state_norm_cache() logical_where = _normalize_where(where) where = self._validate_support(logical_where) @@ -3778,6 +4024,7 @@ def _try_apply_native_submpo( self._route_subtree_messages( local, state_inds, operator_inds, order, token=qtn.rand_uuid(), + workers=self.subtree_workers, ) if operator_inds[hub]: @@ -3824,6 +4071,7 @@ def _apply_factorized_subtree_operator_impl( self._route_subtree_messages( local, state_inds, operator_inds, order, token=qtn.rand_uuid(), + workers=self.subtree_workers, ) if operator_inds[hub]: raise ValueError( @@ -4867,6 +5115,7 @@ def copy(self): tree=self.plan, dtype=self.dtype, threads=self.threads, + subtree_workers=self.subtree_workers, layout_objective=self.layout_objective, layout_weight_mode=self.layout_weight_mode, layout_time_decay=self.layout_time_decay, @@ -4895,6 +5144,7 @@ def copy(self): other._backend_conversion_warnings = set( self._backend_conversion_warnings ) + other._track_warning_emitted = self._track_warning_emitted other._logical_qubits = list(self._logical_qubits) other._logical_positions = dict(self._logical_positions) other._truncation_survival = self._truncation_survival @@ -4931,7 +5181,7 @@ def get_projection_diagnostics(self): @classmethod def find_tree_layout(cls, gates, n=None, *, structure="quality", max_arity=2, community_frac=0.35, - star_frac=0.75, layout_objective="path", + star_frac=0.75, layout_objective="congestion", layout_weight_mode="count", layout_time_decay=None, layout_time_window=None, root_qubit=None, top_arity=_DEFAULT_TOP_ARITY, diff --git a/src/pepsy/optimizers/tree/ttn.py b/src/pepsy/optimizers/tree/ttn.py index 3e796d4..a8c1fd6 100644 --- a/src/pepsy/optimizers/tree/ttn.py +++ b/src/pepsy/optimizers/tree/ttn.py @@ -709,10 +709,10 @@ def bond(self, a, b): def isometry_direction(self, nid): """Return the neighbour proven by ``left_inds`` to receive node ``nid``. - A dense tree tensor is an isometry toward exactly one adjacent node - when its ``left_inds`` contain every leg except that shared tree bond. - ``None`` means no usable local proof is currently recorded. This is a - derived view of the live tensor metadata, not separately tracked state. + A tree tensor is an isometry toward exactly one adjacent node when its + ``left_inds`` contain every leg except that shared tree bond. ``None`` + means no usable local proof is currently recorded. This is a derived + view of the live tensor metadata, not separately tracked state. """ if nid not in self._plan.children: raise ValueError(f"{nid!r} is not a node of the tree.") @@ -764,14 +764,16 @@ def can_skip_canonize(self, a, b, *, absorb="right"): """Whether local metadata proves edge canonicalisation is redundant. With ``absorb="right"`` node ``a`` must already be isometric toward - ``b``; the ``"left"`` orientation is symmetric. Native fermionic trees - always return ``False`` because their graded QR path remains explicit. + ``b``; the ``"left"`` orientation is symmetric. For native fermionic + tensors the same local proof is accepted only when the live data is a + Symmray fermionic array with aligned charge maps. This deliberately + checks structure, not numerical isometry: the metadata is written by + the native graded QR path, while malformed or unknown metadata falls + back to explicit graded QR. """ if absorb not in {"right", "left"}: raise ValueError("absorb must be 'right' or 'left'.") bond = self.bond(a, b) # validate the requested tree edge - if self.fermionic: - return False if absorb == "right": node = a else: @@ -779,7 +781,38 @@ def can_skip_canonize(self, a, b, *, absorb="right"): tensor = self.node_tensor(node) if tensor.left_inds is None: return False - return set(tensor.left_inds) == set(tensor.inds) - {bond} + if set(tensor.left_inds) != set(tensor.inds) - {bond}: + return False + if not self.fermionic: + return True + + # Native graded arrays carry the phase/charge convention in their + # index metadata. Do not infer it from dense data or trust arbitrary + # ``left_inds`` supplied by a caller: only the native array contract + # is eligible for this QR-free move. + data = tensor.data + if not getattr(data, "fermionic", False): + return False + indices = getattr(data, "indices", None) + duals = getattr(data, "duals", None) + charges = getattr(data, "charges", None) + if ( + indices is None + or duals is None + or charges is None + or len(indices) != len(tensor.inds) + or len(duals) != len(tensor.inds) + or len(charges) != len(tensor.inds) + ): + return False + check_aligned = getattr(data, "check_chargemaps_aligned", None) + if check_aligned is None: + return False + try: + check_aligned() + except (AttributeError, TypeError, ValueError, KeyError): + return False + return True def validate_isometry_metadata(self, region=None): """Validate local ``left_inds`` against a canonical region. @@ -1173,11 +1206,26 @@ def _track_edge_center(self, a, b, absorb, *, previous=None): def canonize_edge_(self, a, b, absorb="right"): """Canonicalise across the tree edge ``a -> b`` in place. - Dense/nonfermionic trees delegate to Quimb's ``canonize_between``. - Native fermionic trees use the explicit graded QR helper above. + Dense/nonfermionic trees delegate to Quimb's ``canonize_between``; + native fermionic trees use the explicit graded QR helper above. ``absorb="right"`` leaves node ``a`` isometric and pushes the tracked - orthogonality centre onto node ``b``. + orthogonality centre onto node ``b``. Edges whose live metadata proves + the required isometry are metadata-only moves for both dense and + native graded arrays. """ + if self.can_skip_canonize(a, b, absorb=absorb): + # The local QR is already represented by the tensor's proven + # ``left_inds``. Keep centre bookkeeping honest, but do not touch + # tensor data or invalidate the norm cache. + previous = self.orthogonality_center + source, target = ( + (a, b) if absorb == "right" else (b, a) + ) + if previous == source: + self._canonical_region = frozenset({target}) + elif previous not in (None, target): + self._canonical_region = None + return self previous = self.orthogonality_center if self.fermionic: self._fermionic_canonize_edge_(a, b, absorb) @@ -1217,6 +1265,16 @@ def compress_edge_( non-shared legs. Native fermionic compression ignores this option and retains its explicit graded split. """ + bond = self.bond(a, b) + before_bond = int(self.ind_size(bond)) + if cutoff == 0.0 and ( + max_bond is None or before_bond <= int(max_bond) + ): + # No singular value can be removed, so a QR gauge move is exact + # for both dense and native graded trees. This also protects direct + # TreeTensorNetwork callers that do not go through the optimizer's + # diagnostic-aware wrapper. + return self.canonize_edge_(a, b, absorb=absorb) previous = self.orthogonality_center if self.fermionic: self._fermionic_compress_edge_( diff --git a/src/pepsy/vmc/netket.py b/src/pepsy/vmc/netket.py index c3efe5a..00bfb24 100644 --- a/src/pepsy/vmc/netket.py +++ b/src/pepsy/vmc/netket.py @@ -4359,6 +4359,54 @@ def netket_fermion_operator(hilbert, terms, *, constant=0.0, conserving=False): return total +def _fermi_hubbard_terms(edges, n_sites, *, t, U): + """Return second-quantized terms for the spinful Hubbard Hamiltonian.""" + terms = [] + for site in range(int(n_sites)): + terms.append( + ( + U, + ( + (site, +1, True), + (site, +1, False), + (site, -1, True), + (site, -1, False), + ), + ) + ) + for left, right in edges: + left, right = int(left), int(right) + for spin in (+1, -1): + terms.extend( + ( + (-t, ((left, spin, True), (right, spin, False))), + (-t, ((right, spin, True), (left, spin, False))), + ) + ) + return terms + + +def _build_netket_fermi_hubbard_operator(hilbert, graph, *, n_sites, t, U): + """Build Hubbard metadata across NetKet's native API generations. + + NetKet releases before the removal of ``FermiHubbardJax`` expose the + equivalent second-quantized operator instead. The fallback preserves the + fixed spin sector and uses NetKet's conserving implementation when + available; it is not a dense or Jordan--Wigner conversion. + """ + nk = _require_netket() + native = getattr(nk.operator, "FermiHubbardJax", None) + if native is not None: + return native(hilbert, graph=graph, t=t, U=U, dtype=float) + + edges = tuple(graph.edges()) + return netket_fermion_operator( + hilbert, + _fermi_hubbard_terms(edges, n_sites, t=t, U=U), + conserving="auto", + ) + + def compile_operator_sum_netket(hilbert, terms, *, site_order=None, conserving=False): """Compile a backend-neutral :class:`OperatorSum` for NetKet. @@ -5044,9 +5092,11 @@ def build_fermi_hubbard_vmc( internally, and its variational parameters use quimb's native ``qtn.pack``/``qtn.unpack`` representation. - NetKet's ``FermiHubbardJax`` Hamiltonian is the canonical fixed-sector - Hamiltonian: no chemical-potential term is added, which is equivalent to - setting ``MU=0`` once the spin sector is fixed. + When available, NetKet's ``FermiHubbardJax`` Hamiltonian is used. Newer + NetKet releases that removed that convenience constructor use the + equivalent conserving second-quantized operator instead. No + chemical-potential term is added, which is equivalent to setting ``MU=0`` + once the spin sector is fixed. When ``register_stable_svd`` is True (default) and ``contraction`` is an SVD-based approximation (``hotrg``/``ctmrg``/``boundary``), Pepsy installs @@ -5109,12 +5159,12 @@ def build_fermi_hubbard_vmc( edges = square_lattice_edges(Lx, Ly, pbc=pbc) graph = nk.graph.Graph(edges=tuple(edges), n_nodes=n_sites) - hamiltonian = nk.operator.FermiHubbardJax( + hamiltonian = _build_netket_fermi_hubbard_operator( hilbert, - graph=graph, + graph, + n_sites=n_sites, t=t, U=U, - dtype=float, ) columns = netket_spin_orbital_columns(hilbert) if verify_columns: @@ -5681,12 +5731,12 @@ def build_sparse_fermi_hubbard_vmc( edges = square_lattice_edges(Lx, Ly, pbc=pbc) graph = nk.graph.Graph(edges=tuple(edges), n_nodes=n_sites) - hamiltonian = nk.operator.FermiHubbardJax( + hamiltonian = _build_netket_fermi_hubbard_operator( hilbert, - graph=graph, + graph, + n_sites=n_sites, t=t, U=U, - dtype=float, ) columns = netket_spin_orbital_columns(hilbert) if verify_columns: diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index 7bdc146..ea335b1 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -1086,6 +1086,30 @@ def test_tree_layout_pilot_feedback_is_iterative_and_edge_aware(): assert opt.max_bond() == 1 +def test_tree_layout_pilots_can_run_in_parallel(): + """Independent layout pilots preserve the normal report contract.""" + gates = [(pepsy.cnot(), (0, 3)), (pepsy.cnot(), (3, 1))] + opt = TreeOptimizer(gates, n=4, chi=1, run=False) + + selected = opt.optimize_layout( + objective="full_tree", + pilot_candidates=2, + pilot_workers=2, + pilot_steps=2, + rounds=1, + topology_budget=1, + refine_budget=1, + search_budget=2, + ) + + assert selected["pilot"]["selected_candidate"] + assert len(selected["pilot"]["reports"]) == 2 + assert all( + report["status"] == "ok" + for report in selected["pilot"]["reports"].values() + ) + + def test_tree_layout_targeted_candidates_are_static_and_bounded(): """Hot-edge proposal generation never allocates or mutates a TTN.""" finder = TreeLayoutFinder( @@ -1173,7 +1197,97 @@ def test_full_tree_profile_reports_dynamic_cost_at_all_scales(): report = finder.report(plan) assert report["objective"] == "full_tree" assert report["full_tree"] == profile - assert len(report["objective_key"]) == 8 + assert len(report["objective_key"]) == 10 + assert report["objective_key"][:4] == ( + profile["peak_overflow_log2"], + profile["total_overflow_log2"], + profile["peak_edge_demand_log2"], + profile["total_edge_demand_log2"], + ) + + +def test_full_tree_6x6_pbc_calibrates_against_actual_replay(): + """All-scale overflow ranking tracks real capped-tree replay pressure.""" + Lx = Ly = 6 + n = Lx * Ly + + def site(x, y): + return x * Ly + y + + edges = [] + for x in range(Lx): + for y in range(Ly): + for dx, dy in ((1, 0), (0, 1)): + edge = tuple(sorted(( + site(x, y), + site((x + dx) % Lx, (y + dy) % Ly), + ))) + if edge[0] != edge[1] and edge not in edges: + edges.append(edge) + + gates = ( + [(pepsy.h(), q) for q in range(n)] + + [(pepsy.cphase(np.pi / 4), edge) for edge in edges] + ) + finder = TreeLayoutFinder( + gates, + n=n, + max_arity=(2, 3, 4), + top_arity=3, + objective="full_tree", + chi=4, + seed=11, + ) + recommendation = finder.recommend_arities( + (2, 3, 4), + refine=None, + topology_refine=None, + search=None, + ) + + calibrated = [] + for candidate in recommendation["candidates"]: + optimizer = TreeOptimizer( + gates, + n=n, + tree=candidate["plan"], + chi=4, + cutoff=0.0, + track_truncation=False, + run=False, + ) + optimizer.run() + history = optimizer.truncation_history + calibrated.append({ + "arity": candidate["max_arity"], + "predicted_total_overflow": candidate["full_tree_profile"][ + "total_overflow_log2" + ], + "actual_total_excess": sum( + max(0, event["before_bond"] - 4) for event in history + ), + "actual_truncations": sum( + bool(event["truncated"]) for event in history + ), + }) + + by_arity = {item["arity"]: item for item in calibrated} + assert recommendation["recommended_max_arity"] == 4 + assert ( + by_arity[4]["predicted_total_overflow"] + < by_arity[3]["predicted_total_overflow"] + < by_arity[2]["predicted_total_overflow"] + ) + assert ( + by_arity[4]["actual_total_excess"] + < by_arity[3]["actual_total_excess"] + < by_arity[2]["actual_total_excess"] + ) + assert ( + by_arity[4]["actual_truncations"] + < by_arity[3]["actual_truncations"] + < by_arity[2]["actual_truncations"] + ) def test_full_tree_anneals_subtrees_without_changing_binary_contract(): @@ -1301,6 +1415,15 @@ def test_optimizer_exposes_congestion_layout_objective(): assert opt.plan.is_binary() +def test_optimizer_defaults_to_congestion_layout_for_replay_performance(): + """Automatic optimizer layouts default to finite-chi edge pressure.""" + opt = TreeOptimizer(None, n=6, run=False) + + assert opt.layout_objective == "congestion" + assert opt.layout_finder.objective == "congestion" + assert opt.track_truncation is False + + def test_layout_recommends_arity_and_reports_tree_shape(): """The finder compares binary/wider candidates and exposes their costs.""" rng = np.random.default_rng(103) @@ -2317,6 +2440,62 @@ def test_truncation_history_keeps_fast_untracked_path_cheap(): assert all(event["discarded_weight"] is None for event in report["events"]) +def test_track_truncation_warns_and_skips_impossible_spectra(): + """Tracking warns, while lossless within-cap edges still use QR only.""" + with pytest.warns(UserWarning, match="track_truncation=True"): + opt = TreeOptimizer( + [(pepsy.cnot(), (0, 3))], + n=4, + chi=16, + cutoff=0.0, + track_truncation=True, + ) + + assert opt.track_truncation is True + assert opt.truncation_history + assert all(event["spectrum_rank"] is None for event in opt.truncation_history) + + +def test_repeated_direct_gate_reuses_factorization(monkeypatch): + """Repeated immutable gate objects do not repeat their operator SVD.""" + original_split = qtn.Tensor.split + svd_calls = [] + + def traced_split(tensor, *args, **kwargs): + if kwargs.get("method") == "svd": + svd_calls.append(tensor) + return original_split(tensor, *args, **kwargs) + + monkeypatch.setattr(qtn.Tensor, "split", traced_split) + gate = pepsy.cnot() + opt = TreeOptimizer(None, n=4, chi=16, cutoff=0.0, run=False) + opt.apply_2q(gate, 0, 3) + first_count = len(svd_calls) + opt.apply_2q(gate, 0, 3) + + assert len(svd_calls) == first_count + assert sum(key[0] == "direct" for key in opt._gate_factor_cache) == 1 + + +def test_parallel_subtree_messages_match_serial(): + """Independent dense QR message waves preserve the serial result.""" + rng = np.random.default_rng(818) + operator, _ = np.linalg.qr( + rng.standard_normal((8, 8)) + 1j * rng.standard_normal((8, 8)) + ) + serial = TreeOptimizer( + None, n=8, chi=16, cutoff=0.0, subtree_workers=1, run=False, + ) + parallel = TreeOptimizer( + None, n=8, chi=16, cutoff=0.0, subtree_workers=3, run=False, + ) + + serial.apply_subtree_operator(operator, (0, 2, 5)) + parallel.apply_subtree_operator(operator, (0, 2, 5)) + + assert _fidelity(serial.to_dense(), parallel.to_dense()) > 1 - 1e-12 + + def test_convergence_sweep_reports_rising_fidelity(): """convergence_sweep reuses one tree and reports monotone fidelity.""" rng = np.random.default_rng(13) @@ -3873,6 +4052,36 @@ def test_shift_center_idempotent_touches_nothing(): assert np.array_equal(ttn.node_tensor(nid).data, snap[nid]) +def test_dense_edge_canonization_skips_proven_isometry(monkeypatch): + """A direct lossless edge move reuses a live dense ``left_inds`` proof.""" + ttn = _entangled_ttn(seed=40) + leaf = ttn.leaf_of_qubit(0) + parent = ttn.parent(leaf) + ttn.shift_orthogonality_center(parent) + assert ttn.orthogonality_center == parent + calls = [] + canonize_between = ttn.canonize_between + + def traced_canonize_between(*args, **kwargs): + calls.append((args, kwargs)) + return canonize_between(*args, **kwargs) + + monkeypatch.setattr(ttn, "canonize_between", traced_canonize_between) + before = { + nid: np.array(ttn.node_tensor(nid).data) + for nid in ttn.plan.nodes() + } + + assert ttn.can_skip_canonize(leaf, parent) + ttn.canonize_edge_(leaf, parent) + + assert not calls + assert all( + np.array_equal(ttn.node_tensor(nid).data, data) + for nid, data in before.items() + ) + + def test_shift_center_from_unknown_canonicalises_once(): """An unknown centre falls back to a full canonicalisation about the target.""" plan = TreePlan.from_order(range(6), structure="balanced") @@ -4346,8 +4555,84 @@ def build_stream(): assert float(tensors.tn_fidelity(engine.p, mps_exact.p)) > 1 - 1e-8 +@pytest.mark.parametrize( + ("symmetry", "occupations"), + [ + ("U1", (1, 1, 1, 1)), + ("U1U1", ((1, 0), (0, 1), (1, 0), (0, 1))), + ], +) +def test_native_tree_direct_mpo_and_submpo_match_without_global_fidelity( + symmetry, occupations, +): + """Native gate kernels agree without Cotengra's process-based fidelity.""" + pytest.importorskip("symmray") + L = 4 + fermion = pepsy.Fermion( + spinful=True, + symmetry=symmetry, + dtype="complex128", + ) + plan = TreePlan.from_order(range(L), structure="balanced") + seed = pepsy.ps_to_ttn( + L, + tree=plan, + fermion=fermion, + occupations=occupations, + dtype="complex128", + ) + hopping = fermion.hopping_gate(0.05, t=1.0, imaginary=False) + onsite = fermion.onsite_gate( + 0.03, site=1, U=8.0, mu=0.0, imaginary=False, + ) + submpo = qtn.MatrixProductOperator.from_dense( + hopping, dims=(4, 4), sites=(0, 2), L=L, + ) + + def dense_vector(opt): + tensor = opt.tn.contract(all, optimize="greedy").transpose( + *(opt.tn.site_ind(q) for q in range(L)) + ) + return np.asarray(tensor.data.to_dense()).reshape(-1) + + outputs = [] + for mode in ("direct", "mpo"): + opt = TreeOptimizer( + None, + n=L, + tree=plan, + state=seed.copy(), + chi=64, + cutoff=0.0, + mode=mode, + run=False, + ) + opt.apply_1q(onsite, 1) + opt.apply_2q(hopping, 0, 2) + outputs.append(dense_vector(opt)) + assert opt.tn.validate(check_canonical=True) is opt.tn + + submpo_opt = TreeOptimizer( + None, + n=L, + tree=plan, + state=seed.copy(), + chi=64, + cutoff=0.0, + mode="submpo", + run=False, + ) + submpo_opt.apply_1q(onsite, 1) + submpo_opt.apply_submpo(submpo, (0, 2)) + outputs.append(dense_vector(submpo_opt)) + assert submpo_opt.tn.validate(check_canonical=True) is submpo_opt.tn + + for output in outputs[1:]: + assert _fidelity(outputs[0], output) > 1 - 1e-10 + + def test_native_fermionic_submpo_keeps_graded_hub_recovery(monkeypatch): - """Dense isometry metadata must not replace native graded subtree QR.""" + """Native routed Q metadata skips only already-proven graded QR.""" pytest.importorskip("symmray") L = 4 fermion = pepsy.Fermion( @@ -4434,7 +4719,8 @@ def dense_vector(opt): assert installs assert compressions - assert all(reduced is True for reduced in compressions) + assert all(reduced in {True, "left"} for reduced in compressions) + assert any(reduced == "left" for reduced in compressions) assert ( _fidelity(dense_vector(candidate), dense_vector(reference)) > 1 - 1e-10 @@ -4442,10 +4728,105 @@ def dense_vector(opt): assert candidate.validate_isometry_metadata() is candidate for nid, toward in candidate.isometry_map().items(): if toward is not None: - assert not candidate.can_skip_canonize(nid, toward) + assert candidate.can_skip_canonize(nid, toward) assert candidate.tn.validate(check_canonical=True) is candidate.tn +@pytest.mark.parametrize( + ("symmetry", "spinful", "occupations"), + [ + ("U1", False, (1, 0, 1, 0)), + ("U1U1", True, ((1, 0), (0, 1), (1, 0), (0, 1))), + ], +) +def test_native_fermionic_left_inds_skips_lossless_qr( + symmetry, spinful, occupations, monkeypatch, +): + """Symmray U1 variants reuse native QR isometry metadata safely.""" + pytest.importorskip("symmray") + L = 4 + fermion = pepsy.Fermion( + spinful=spinful, + symmetry=symmetry, + dtype="complex128", + ) + plan = TreePlan.from_order(range(L), structure="balanced") + ttn = pepsy.ps_to_ttn( + L, + tree=plan, + fermion=fermion, + occupations=occupations, + dtype="complex128", + ) + target = plan.leaf_of_qubit[0] + ttn.shift_orthogonality_center(target) + + source = next( + nid for nid, toward in ttn.isometry_map().items() + if toward == target and ttn.can_skip_canonize(nid, toward) + ) + before = { + nid: np.asarray(ttn.node_tensor(nid).data.to_dense()).copy() + for nid in plan.nodes() + } + calls = [] + graded_qr = ttn._fermionic_canonize_edge_ + + def traced_qr(*args, **kwargs): + calls.append((args, kwargs)) + return graded_qr(*args, **kwargs) + + monkeypatch.setattr(ttn, "_fermionic_canonize_edge_", traced_qr) + ttn.canonize_edge_(source, target) + + assert not calls + assert ttn.orthogonality_center == target + assert ttn.is_canonical_form(target) + assert ttn.validate(check_canonical=True) is ttn + for nid in plan.nodes(): + np.testing.assert_array_equal( + np.asarray(ttn.node_tensor(nid).data.to_dense()), before[nid] + ) + + +def test_native_truncating_compression_keeps_explicit_svd(monkeypatch): + """A positive cutoff never turns a native truncation into metadata-only.""" + pytest.importorskip("symmray") + fermion = pepsy.Fermion( + spinful=True, + symmetry="U1U1", + dtype="complex128", + ) + plan = TreePlan.from_order(range(4), structure="balanced") + ttn = pepsy.ps_to_ttn( + 4, + tree=plan, + fermion=fermion, + occupations=((1, 0), (0, 1), (1, 0), (0, 1)), + dtype="complex128", + ) + target = plan.leaf_of_qubit[0] + ttn.shift_orthogonality_center(target) + source = next( + nid for nid, toward in ttn.isometry_map().items() + if toward == target and ttn.can_skip_canonize(nid, toward) + ) + + calls = [] + compress = ttn._fermionic_compress_edge_ + + def traced_compress(*args, **kwargs): + calls.append((args, kwargs)) + return compress(*args, **kwargs) + + monkeypatch.setattr(ttn, "_fermionic_compress_edge_", traced_compress) + ttn.compress_edge_( + source, target, max_bond=64, cutoff=1e-10, cutoff_mode="rel", + ) + + assert calls + + def test_tree_stable_labels_route_submpo_by_payload_sites(monkeypatch): """Stable logical labels do not disable native structured MPO routing.""" x = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex) From dbb6dda479e795873d32bf3369ebd414b55309c6 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Sat, 1 Aug 2026 18:02:10 -0700 Subject: [PATCH 49/70] Fix backend-safe canonical checks and fidelity accumulation --- docs/api/optimizers/mps.md | 6 ++- docs/api/optimizers/tree.md | 3 ++ src/pepsy/optimizers/_fidelity.py | 53 +++++++++++++++++++++++++ src/pepsy/optimizers/mps/optimizer.py | 42 ++++++++++++-------- src/pepsy/optimizers/tree/optimizer.py | 54 ++++++++++++++++++-------- src/pepsy/optimizers/tree/ttn.py | 19 ++++++--- tests/test_optimize_mps.py | 10 +++++ tests/test_optimize_tree.py | 46 ++++++++++++++++++++++ 8 files changed, 192 insertions(+), 41 deletions(-) create mode 100644 src/pepsy/optimizers/_fidelity.py diff --git a/docs/api/optimizers/mps.md b/docs/api/optimizers/mps.md index df15e63..b8a7b63 100644 --- a/docs/api/optimizers/mps.md +++ b/docs/api/optimizers/mps.md @@ -100,8 +100,10 @@ For unitary streams, `get_infidelities()` uses the current retained norm divided by the initial run norm, squared; this is the global retained-fidelity estimate and is equivalent to multiplying the per-gate fidelities without requiring a -pre-gate target measurement. Local ratios remain available in detailed -samples for diagnostics. The trace is populated by default. Set +pre-gate target measurement. Local products and norm-ratio evaluations are +accumulated in the log domain and exponentiated only for readout, so long +streams and very small retained norms do not lose fidelity to underflow. Local +ratios remain available in detailed samples for diagnostics. The trace is populated by default. Set ``track_infidelity=False`` in the constructor, or pass ``track_infidelity=False`` to ``run()``, to skip target-norm construction, retained-norm calculations, samples, and progress-bar infidelity fields. For diff --git a/docs/api/optimizers/tree.md b/docs/api/optimizers/tree.md index cf1529a..94b64f2 100644 --- a/docs/api/optimizers/tree.md +++ b/docs/api/optimizers/tree.md @@ -949,6 +949,9 @@ scale. diagnostic spectrum probes can add substantial SVD work. It remains disabled by default. Lossless zero-cutoff edges that are already within their bond cap use QR and do not probe a spectrum even when tracking is on. + Per-edge retained survival and cumulative infidelity are accumulated in log + space and exponentiated only for readout, avoiding product underflow on long + streams. The report also contains gate-level `updates`, grouping edge events by support and reporting the cumulative relative loss. - `TreeOptimizer.convergence_sweep(gates, n, chi_values, ops=...)` replays the diff --git a/src/pepsy/optimizers/_fidelity.py b/src/pepsy/optimizers/_fidelity.py new file mode 100644 index 0000000..c32dc51 --- /dev/null +++ b/src/pepsy/optimizers/_fidelity.py @@ -0,0 +1,53 @@ +"""Numerically stable fidelity and infidelity helpers for optimizers.""" + +from __future__ import annotations + +import math + +import numpy as np + + +def log_fidelity_from_norms(approx_norm, target_norm): + """Return ``log((approx_norm / target_norm) ** 2)`` clipped at zero. + + Norm ratios are used as retained-fidelity proxies by both the MPS and tree + optimizers. Computing the ratio first can overflow or underflow even when + the logarithm of the ratio is representable, so the subtraction is done in + log space. Invalid or zero targets follow the historical clipped-fidelity + convention: two zero norms have fidelity one, otherwise the fidelity is + zero. + """ + approx = float(np.real(approx_norm)) + target = float(np.real(target_norm)) + + if target <= 0.0: + return 0.0 if approx <= 0.0 else -math.inf + if approx <= 0.0 or math.isnan(approx) or math.isnan(target): + return -math.inf + if approx == target: + return 0.0 + + # The normal path is finite. Handle infinities conservatively without + # creating ``inf - inf == nan``. + if not math.isfinite(approx) or not math.isfinite(target): + return 0.0 if approx > target else -math.inf + + return min(0.0, 2.0 * (math.log(approx) - math.log(target))) + + +def fidelity_from_log(log_fidelity): + """Convert a clipped log-fidelity to a finite value in ``[0, 1]``.""" + log_fidelity = float(log_fidelity) + if math.isnan(log_fidelity) or log_fidelity == -math.inf: + return 0.0 + return min(1.0, max(0.0, math.exp(min(0.0, log_fidelity)))) + + +def infidelity_from_log(log_fidelity): + """Return ``1 - exp(log_fidelity)`` stably for small losses.""" + log_fidelity = float(log_fidelity) + if math.isnan(log_fidelity) or log_fidelity == -math.inf: + return 1.0 + # ``-expm1`` preserves small positive infidelities that ``1 - exp`` + # would round to zero. Fidelity is clipped at one by construction. + return min(1.0, max(0.0, -math.expm1(min(0.0, log_fidelity)))) diff --git a/src/pepsy/optimizers/mps/optimizer.py b/src/pepsy/optimizers/mps/optimizer.py index ecea199..20c4525 100644 --- a/src/pepsy/optimizers/mps/optimizer.py +++ b/src/pepsy/optimizers/mps/optimizer.py @@ -74,6 +74,11 @@ gate as apply_gate, gate_simple as apply_gate_simple, ) +from .._fidelity import ( + fidelity_from_log, + infidelity_from_log, + log_fidelity_from_norms, +) from .layout import ( MpsGateStreamLayoutFinder, _normalize_layout_support, @@ -3429,15 +3434,19 @@ def _append_compression_infidelity_sample( ratios. Accumulating ``log(F)`` avoids losing information to floating-point underflow on long non-unitary streams. """ - local_fidelity = self._norm_ratio_fidelity(approx_norm, target_norm) - if local_fidelity <= 0.0 or np.isneginf(self._infidelity_log_fidelity): + local_log_fidelity = log_fidelity_from_norms( + approx_norm, target_norm, + ) + if ( + np.isneginf(local_log_fidelity) + or np.isneginf(self._infidelity_log_fidelity) + ): self._infidelity_log_fidelity = -np.inf else: - self._infidelity_log_fidelity += float(np.log(local_fidelity)) - cumulative_infidelity = ( - 1.0 - if np.isneginf(self._infidelity_log_fidelity) - else float(-np.expm1(self._infidelity_log_fidelity)) + self._infidelity_log_fidelity += local_log_fidelity + local_fidelity = fidelity_from_log(local_log_fidelity) + cumulative_infidelity = infidelity_from_log( + self._infidelity_log_fidelity, ) sample = { "step": int(step), @@ -3498,11 +3507,12 @@ def _append_unitary_compression_infidelity_sample( if self._unitary_initial_norm is None: self._unitary_initial_norm = previous_norm if self._unitary_global_norm_tracking: - global_fidelity = self._norm_ratio_fidelity( + global_log_fidelity = log_fidelity_from_norms( approx_norm, self._unitary_initial_norm, ) - global_infidelity = 1.0 - global_fidelity + global_fidelity = fidelity_from_log(global_log_fidelity) + global_infidelity = infidelity_from_log(global_log_fidelity) sample["fidelity"] = global_fidelity sample["global_fidelity"] = global_fidelity sample["global_infidelity"] = global_infidelity @@ -3714,14 +3724,12 @@ def _gate_target_norm_from_expectation(self, p, gate, where): @staticmethod def _norm_ratio_fidelity(approx_norm, target_norm): """Return clipped ``(||approx|| / ||target||)**2``.""" - approx = MpsOptimizer._real_float(approx_norm) - target = MpsOptimizer._real_float(target_norm) - - if target <= 0.0: - return 1.0 if approx <= 0.0 else 0.0 - - fidelity = (approx / target) ** 2 - return min(1.0, max(0.0, float(fidelity))) + return fidelity_from_log( + log_fidelity_from_norms( + MpsOptimizer._real_float(approx_norm), + MpsOptimizer._real_float(target_norm), + ) + ) def _build_norm_target(self, p, gate, where, cutoff, cutoff_mode="rsum2"): """Build the pre-chi-compression target used for norm diagnostics.""" diff --git a/src/pepsy/optimizers/tree/optimizer.py b/src/pepsy/optimizers/tree/optimizer.py index deba716..6b531ad 100644 --- a/src/pepsy/optimizers/tree/optimizer.py +++ b/src/pepsy/optimizers/tree/optimizer.py @@ -55,6 +55,11 @@ to_float, ) from ...operators.gates import _normalize_gate_entries +from .._fidelity import ( + fidelity_from_log, + infidelity_from_log, + log_fidelity_from_norms, +) from ..mps.optimizer import ( _control_event_parts as _mps_control_event_parts, normalize_submpo_where, @@ -681,7 +686,7 @@ def __init__(self, gates=None, n=None, *, chi=64, self._gate_factor_cache = {} self._gate_factor_cache_limit = 64 self._active_update = None - self._truncation_survival = 1.0 + self._truncation_log_survival = 0.0 if tree is None: self.layout_finder = TreeLayoutFinder( @@ -1281,7 +1286,7 @@ def set_tn(self, tn): self.infidelity_samples.clear() self.normalizations.clear() self.projection_diagnostics.clear() - self._truncation_survival = 1.0 + self._truncation_log_survival = 0.0 return self def set_p(self, tn): @@ -2138,11 +2143,17 @@ def _finish_update(self): if event["discarded_fraction"] is not None ] if tracked: - edge_survival = float(np.prod([ - max(0.0, 1.0 - event["discarded_fraction"]) - for event in tracked - ])) - relative_loss = float(1.0 - edge_survival) + edge_log_survival = 0.0 + for event in tracked: + edge_survival = min( + 1.0, + max(0.0, 1.0 - float(event["discarded_fraction"])), + ) + if edge_survival <= 0.0: + edge_log_survival = -np.inf + break + edge_log_survival += float(np.log(edge_survival)) + relative_loss = infidelity_from_log(edge_log_survival) absolute_loss = float( sum(event["discarded_weight"] for event in tracked) ) @@ -2152,15 +2163,25 @@ def _finish_update(self): max_edge_fraction = float( max(event["discarded_fraction"] for event in tracked) ) - self._truncation_survival *= edge_survival - cumulative_loss = float(1.0 - self._truncation_survival) + if ( + np.isneginf(self._truncation_log_survival) + or np.isneginf(edge_log_survival) + ): + self._truncation_log_survival = -np.inf + else: + self._truncation_log_survival += edge_log_survival + cumulative_loss = infidelity_from_log( + self._truncation_log_survival, + ) else: if self.track_truncation: relative_loss = 0.0 absolute_loss = 0.0 max_edge_loss = 0.0 max_edge_fraction = 0.0 - cumulative_loss = float(1.0 - self._truncation_survival) + cumulative_loss = infidelity_from_log( + self._truncation_log_survival, + ) else: relative_loss = None absolute_loss = None @@ -2191,7 +2212,7 @@ def _finish_update(self): "step": len(self.infidelity_samples) + 1, "where": active["support"], "edge_count": len(edge_events), - "local_fidelity": float(1.0 - local_infidelity), + "local_fidelity": fidelity_from_log(edge_log_survival), "local_infidelity": local_infidelity, "infidelity": float(cumulative_loss), "cumulative_infidelity": float(cumulative_loss), @@ -2418,10 +2439,11 @@ def run(self, gates=None, *, progbar=False, mode=None, non_unitary=False, elif progress_reference_norm in (None, 0.0): truncation_infidelity = 0.0 else: - survival = state_norm / progress_reference_norm - truncation_infidelity = max( - 0.0, - 1.0 - survival * survival, + truncation_infidelity = infidelity_from_log( + log_fidelity_from_norms( + state_norm, + progress_reference_norm, + ) ) postfix["infidelity"] = self._format_progress_infidelity( truncation_infidelity @@ -5147,7 +5169,7 @@ def copy(self): other._track_warning_emitted = self._track_warning_emitted other._logical_qubits = list(self._logical_qubits) other._logical_positions = dict(self._logical_positions) - other._truncation_survival = self._truncation_survival + other._truncation_log_survival = self._truncation_log_survival return other def get_infidelities(self): diff --git a/src/pepsy/optimizers/tree/ttn.py b/src/pepsy/optimizers/tree/ttn.py index a8c1fd6..6309f98 100644 --- a/src/pepsy/optimizers/tree/ttn.py +++ b/src/pepsy/optimizers/tree/ttn.py @@ -46,6 +46,7 @@ from quimb.tensor.tensor_core import TensorNetwork from numbers import Integral +from ...backends import to_float from .layout import TreePlan, _DEFAULT_TOP_ARITY __all__ = ["TreeTensorNetwork"] @@ -1515,14 +1516,20 @@ def is_subtree_canonical_form(self, nodes=None, *, span=False, tol=1e-9): prod = qtn.tensor_contract(t, tc, output_inds=output_inds) d = int(prod.shape[0]) data = prod.data - # Autoray's generic NumPy conversion deliberately leaves a - # Symmray block array intact. Its explicit dense readout is only - # used here for this diagnostic, never in a simulation hot path. + # Keep the diagnostic on the live backend. In particular, + # ``ar.to_numpy`` cannot move a CUDA tensor to the host, while a + # scalar reduction can be transferred safely and cheaply. if hasattr(data, "to_dense"): data = data.to_dense() - else: - data = ar.to_numpy(data) - if not np.allclose(data, np.eye(d), atol=tol): + identity = ar.do("eye", d, like=data) + close = ar.do( + "allclose", + data, + identity, + atol=float(tol), + rtol=1.0e-5, + ) + if not bool(to_float(close, real=True)): return False return True diff --git a/tests/test_optimize_mps.py b/tests/test_optimize_mps.py index 8fc91f4..9253c6a 100644 --- a/tests/test_optimize_mps.py +++ b/tests/test_optimize_mps.py @@ -145,6 +145,16 @@ def test_mps_optimizer_run_can_override_infidelity_tracking(): assert len(opt.get_infidelity_samples()) == 1 +def test_mps_norm_fidelity_uses_log_ratio_for_extreme_scales(): + """Norm-ratio fidelity remains finite without forming a huge quotient.""" + p0 = qtn.MPS_computational_state("0", dtype="complex128") + opt = py.MpsOptimizer(p0, gates=[], chi=2, mode="svd") + + fidelity = opt._norm_ratio_fidelity(1.0e300, 1.0e300 * (1.0 + 1.0e-8)) + + assert fidelity == pytest.approx((1.0 / (1.0 + 1.0e-8)) ** 2) + + def test_mps_optimizer_simple_update_routes_torch_u1u1_long_range_gate(): """SU routed SWAPs should stay on the live Torch Symmray backend.""" torch = pytest.importorskip("torch") diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index ea335b1..5eba443 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -651,6 +651,34 @@ def test_tree_truncation_infidelity_compatibility_trace(): assert opt.get_normalizations() == [] +def test_tree_truncation_survival_accumulates_in_log_space(): + """Many local survival factors remain stable in the cumulative trace.""" + opt = TreeOptimizer(None, n=2, chi=2, track_truncation=True, run=False) + opt._active_update = { + "kind": "gate", + "support": (0, 1), + "edge_start": 0, + "started_at": 0.0, + } + local_survival = 0.999999999999 + count = 1000 + opt.truncation_history = [ + { + "discarded_fraction": 1.0 - local_survival, + "discarded_weight": 1.0 - local_survival, + "truncated": True, + } + for _ in range(count) + ] + + opt._finish_update() + + expected_log = count * np.log(local_survival) + expected_infidelity = -np.expm1(expected_log) + assert opt._truncation_log_survival == pytest.approx(expected_log) + assert opt.get_infidelities()[-1] == pytest.approx(expected_infidelity) + + def test_tree_run_supports_shared_non_unitary_normalization_controls(): """Tree replay accepts the shared non-unitary normalization contract.""" half = 0.5 * np.eye(2, dtype=complex) @@ -3579,6 +3607,24 @@ def test_tree_torch_state_stays_native_across_public_operations(): assert all(torch.is_tensor(tensor.data) for tensor in opt.tn.tensor_map.values()) +def test_tree_canonical_check_uses_backend_scalar_reduction(monkeypatch): + """Canonical diagnostics must not convert device tensors with NumPy.""" + torch = pytest.importorskip("torch") + import importlib + + ttn_module = importlib.import_module("pepsy.optimizers.tree.ttn") + state = TreeTensorNetwork.from_plan(TreePlan.from_order(range(3))) + state.apply_to_arrays( + pepsy.backend_torch(device="cpu", dtype=torch.complex128) + ) + + def fail_to_numpy(_value): + raise AssertionError("canonical checks must stay on the live backend") + + monkeypatch.setattr(ttn_module.ar, "to_numpy", fail_to_numpy) + assert state.is_canonical_form() + + def test_tree_warns_once_when_a_gate_does_not_match_the_state_backend(): """User payload mismatches are explicit while compatibility is preserved.""" torch = pytest.importorskip("torch") From 8d9403729d18c0319a74e3a4568fafc7035ad075 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Sat, 1 Aug 2026 20:46:50 -0600 Subject: [PATCH 50/70] Harden native fermionic tree QR stabilization --- src/pepsy/optimizers/tree/optimizer.py | 12 +++ src/pepsy/optimizers/tree/ttn.py | 14 ++- tests/test_optimize_tree.py | 129 +++++++++++++++++++++++++ 3 files changed, 153 insertions(+), 2 deletions(-) diff --git a/src/pepsy/optimizers/tree/optimizer.py b/src/pepsy/optimizers/tree/optimizer.py index 6b531ad..036394b 100644 --- a/src/pepsy/optimizers/tree/optimizer.py +++ b/src/pepsy/optimizers/tree/optimizer.py @@ -3118,6 +3118,12 @@ def _fermionic_thread_hop(self, u, v): method="qr", absorb="right", cutoff=0.0, + # Threaded native blocks can be rank deficient by symmetry. The + # stabilized Symmray QR normalizes every R diagonal entry to a + # phase, so a structural zero becomes 0 / |0| -> NaN in + # complex64. Plain QR is still lossless here and leaves those + # zero sectors finite. + stabilized=False, get="tensors", ) merged_v = qtn.tensor_contract(carry, tv) @@ -3320,6 +3326,9 @@ def _split_with_diagnostics( method="qr", cutoff=0.0, absorb="right", + # Native Symmray QR must not phase-normalize structural-zero + # diagonal entries; retain Quimb's dense default otherwise. + stabilized=not self.tn.fermionic, get="tensors", bond_ind=bond_ind, ) @@ -3503,6 +3512,9 @@ def _qr_route_message(tensor, left_inds, *, bond_ind): left_inds=left_inds, method="qr", absorb="right", + # Keep the dense Quimb phase convention, but avoid 0 / |0| in + # native Symmray structural-zero sectors. + stabilized=not _is_symmray_array(tensor.data), get="tensors", bond_ind=bond_ind, ) diff --git a/src/pepsy/optimizers/tree/ttn.py b/src/pepsy/optimizers/tree/ttn.py index 6309f98..278691c 100644 --- a/src/pepsy/optimizers/tree/ttn.py +++ b/src/pepsy/optimizers/tree/ttn.py @@ -1134,6 +1134,9 @@ def _fermionic_canonize_edge_(self, a, b, absorb): method="qr", absorb="right", cutoff=0.0, + # Native graded blocks can have exact structural-zero R diagonals; + # skip Quimb's phase normalization so those sectors stay finite. + stabilized=False, get="tensors", ) merged = qtn.tensor_contract(carry, reduced) @@ -1332,12 +1335,19 @@ def canonize_subtree_(self, nodes, *, span=False, absorb="right"): """ region = self._validated_region(nodes, span=span) tags = [self.node_tag(n) for n in region] + canonize_opts = { + "method": "qr", + "cutoff": 0.0, + } + if self.fermionic: + # The Symmray QR backend can encounter exact structural-zero + # diagonals while gauging a native fermionic region. + canonize_opts["stabilized"] = False self.canonize_around_( tags, which="any", absorb=absorb, - method="qr", - cutoff=0.0, + **canonize_opts, ) self._canonical_region = region return self diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index 5eba443..cf47fbb 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -4601,6 +4601,135 @@ def build_stream(): assert float(tensors.tn_fidelity(engine.p, mps_exact.p)) > 1 - 1e-8 +@pytest.mark.filterwarnings( + "ignore:TreeOptimizer is converting a gate/operator payload" +) +def test_native_complex64_threading_disables_zero_phase_stabilization(monkeypatch): + """Native path threading keeps structural-zero QR sectors finite.""" + pytest.importorskip("symmray") + L = 4 + fermion = pepsy.Fermion( + spinful=True, + symmetry="U1U1", + dtype="complex64", + ) + plan = TreePlan.from_order(range(L), structure="balanced") + state = pepsy.ps_to_ttn( + L, + tree=plan, + fermion=fermion, + occupations=((1, 0), (0, 1), (1, 0), (0, 1)), + dtype="complex64", + ) + gate = fermion.hopping_gate(0.05, t=1.0, imaginary=False) + + threaded = False + qr_stabilized = [] + original_hop = TreeOptimizer._fermionic_thread_hop + original_split = qtn.Tensor.split + + def traced_hop(self, u, v): + nonlocal threaded + threaded = True + try: + return original_hop(self, u, v) + finally: + threaded = False + + def traced_split(self, *args, **kwargs): + if threaded and kwargs.get("method") == "qr": + qr_stabilized.append(kwargs.get("stabilized")) + return original_split(self, *args, **kwargs) + + monkeypatch.setattr(TreeOptimizer, "_fermionic_thread_hop", traced_hop) + monkeypatch.setattr(qtn.Tensor, "split", traced_split) + + optimizer = TreeOptimizer( + None, + n=L, + tree=plan, + state=state, + chi=16, + cutoff=0.0, + mode="direct", + run=False, + ) + optimizer.apply_2q(gate, 0, 2) + + assert qr_stabilized + assert qr_stabilized == [False] * len(qr_stabilized) + assert all( + np.isfinite(np.asarray(block)).all() + for tensor in optimizer.tn.tensors + for block in tensor.data.blocks.values() + ) + + +@pytest.mark.filterwarnings( + "ignore:TreeOptimizer is converting a gate/operator payload" +) +def test_native_complex64_all_lossless_qr_routes_skip_zero_phase(monkeypatch): + """All native tree QR routes avoid phase division on zero sectors.""" + pytest.importorskip("symmray") + L = 4 + fermion = pepsy.Fermion( + spinful=True, + symmetry="U1U1", + dtype="complex64", + ) + plan = TreePlan.from_order(range(L), structure="balanced") + state = pepsy.ps_to_ttn( + L, + tree=plan, + fermion=fermion, + occupations=((1, 0), (0, 1), (1, 0), (0, 1)), + dtype="complex64", + ) + hopping = fermion.hopping_gate(0.05, t=1.0, imaginary=False) + routed_ops = [ + fermion.onsite_gate( + 0.01, site=site, U=8.0, mu=0.0, imaginary=False, + ) + for site in (0, 1, 2) + ] + submpo = qtn.MPO_product_operator( + routed_ops, + sites=(0, 1, 2), + L=L, + upper_ind_id="k{}", + lower_ind_id="b{}", + ) + + observed = [] + original_split = qtn.Tensor.split + + def traced_split(self, *args, **kwargs): + if kwargs.get("method") == "qr" and hasattr(self.data, "blocks"): + observed.append(kwargs.get("stabilized")) + return original_split(self, *args, **kwargs) + + monkeypatch.setattr(qtn.Tensor, "split", traced_split) + + for operation in ( + lambda: TreeOptimizer( + None, n=L, tree=plan, state=state.copy(), chi=16, + cutoff=0.0, mode="direct", run=False, + ).apply_2q(hopping, 0, 1), + lambda: TreeOptimizer( + None, n=L, tree=plan, state=state.copy(), chi=16, + cutoff=0.0, mode="direct", run=False, + ).apply_2q(hopping, 0, 2), + lambda: TreeOptimizer( + None, n=L, tree=plan, state=state.copy(), chi=16, + cutoff=0.0, mode="submpo", run=False, + ).apply_submpo(submpo, (0, 1, 2)), + ): + operation() + + assert observed + assert observed == [False] * len(observed) + + @pytest.mark.parametrize( ("symmetry", "occupations"), [ From 98d73ffd85c644cf179772a565c89eb323d712a7 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Sat, 1 Aug 2026 20:53:54 -0600 Subject: [PATCH 51/70] Centralize native tree QR stabilization policy --- src/pepsy/optimizers/tree/optimizer.py | 27 ++++++---------------- src/pepsy/optimizers/tree/ttn.py | 32 ++++++++++++++++++-------- 2 files changed, 30 insertions(+), 29 deletions(-) diff --git a/src/pepsy/optimizers/tree/optimizer.py b/src/pepsy/optimizers/tree/optimizer.py index 036394b..11b0c38 100644 --- a/src/pepsy/optimizers/tree/optimizer.py +++ b/src/pepsy/optimizers/tree/optimizer.py @@ -3112,18 +3112,12 @@ def _fermionic_thread_hop(self, u, v): for index in tu.inds if index not in left_inds ] - keep, carry = tu.split( + keep, carry = self.tn._native_qr_split( + tu, left_inds=left_inds, right_inds=right_inds, - method="qr", absorb="right", cutoff=0.0, - # Threaded native blocks can be rank deficient by symmetry. The - # stabilized Symmray QR normalizes every R diagonal entry to a - # phase, so a structural zero becomes 0 / |0| -> NaN in - # complex64. Plain QR is still lossless here and leaves those - # zero sectors finite. - stabilized=False, get="tensors", ) merged_v = qtn.tensor_contract(carry, tv) @@ -3321,14 +3315,11 @@ def _split_with_diagnostics( if self.track_truncation and not lossless else None ) if lossless: - left, right = tensor.split( + left, right = self.tn._native_qr_split( + tensor, left_inds=left_inds, - method="qr", cutoff=0.0, absorb="right", - # Native Symmray QR must not phase-normalize structural-zero - # diagonal entries; retain Quimb's dense default otherwise. - stabilized=not self.tn.fermionic, get="tensors", bond_ind=bond_ind, ) @@ -3505,16 +3496,12 @@ def descend(node, parent): descend(hub, None) self.center = hub - @staticmethod - def _qr_route_message(tensor, left_inds, *, bond_ind): + def _qr_route_message(self, tensor, left_inds, *, bond_ind): """Split one subtree message losslessly while carrying operator legs.""" - return tensor.split( + return self.tn._native_qr_split( + tensor, left_inds=left_inds, - method="qr", absorb="right", - # Keep the dense Quimb phase convention, but avoid 0 / |0| in - # native Symmray structural-zero sectors. - stabilized=not _is_symmray_array(tensor.data), get="tensors", bond_ind=bond_ind, ) diff --git a/src/pepsy/optimizers/tree/ttn.py b/src/pepsy/optimizers/tree/ttn.py index 278691c..2aec3e5 100644 --- a/src/pepsy/optimizers/tree/ttn.py +++ b/src/pepsy/optimizers/tree/ttn.py @@ -1113,6 +1113,26 @@ def _toward_region(self, nid, region): best = p return best[1] + def _native_qr_options(self, tensor=None): + """Return the QR options needed by native graded tree tensors. + + Symmray's stabilized QR phase-normalizes every diagonal of ``R``. + A structural-zero diagonal therefore creates ``0 / |0|`` and can + produce a NaN in complex64. Dense tensors retain Quimb's default + phase convention; network-level canonicalisation uses ``fermionic`` + because it does not expose the individual tensors here. + """ + native = self.fermionic if tensor is None else _is_symmray_array( + tensor.data + ) + return {"stabilized": False} if native else {} + + def _native_qr_split(self, tensor, **kwargs): + """Perform a QR split with the native graded zero-sector safeguard.""" + kwargs.update(self._native_qr_options(tensor)) + kwargs.setdefault("method", "qr") + return tensor.split(**kwargs) + # -- edge-level canonical / compression helpers --------------------------- def _fermionic_canonize_edge_(self, a, b, absorb): @@ -1128,15 +1148,12 @@ def _fermionic_canonize_edge_(self, a, b, absorb): reduced = self.node_tensor(reduced_node) bond = self.bond(isometric_node, reduced_node) left_inds = [index for index in isometric.inds if index != bond] - kept, carry = isometric.split( + kept, carry = self._native_qr_split( + isometric, left_inds=left_inds, right_inds=(bond,), - method="qr", absorb="right", cutoff=0.0, - # Native graded blocks can have exact structural-zero R diagonals; - # skip Quimb's phase normalization so those sectors stay finite. - stabilized=False, get="tensors", ) merged = qtn.tensor_contract(carry, reduced) @@ -1339,10 +1356,7 @@ def canonize_subtree_(self, nodes, *, span=False, absorb="right"): "method": "qr", "cutoff": 0.0, } - if self.fermionic: - # The Symmray QR backend can encounter exact structural-zero - # diagonals while gauging a native fermionic region. - canonize_opts["stabilized"] = False + canonize_opts.update(self._native_qr_options()) self.canonize_around_( tags, which="any", From 213f9fa6f98fa1155551b7f58565c931a18d9941 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Sat, 1 Aug 2026 20:56:19 -0600 Subject: [PATCH 52/70] Document native tree QR stability policy --- AGENTS.md | 23 +++++++++++++++++++++++ docs/api/optimizers/tree.md | 24 ++++++++++++++++++++++++ 2 files changed, 47 insertions(+) diff --git a/AGENTS.md b/AGENTS.md index 56b4382..d4d88f9 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -65,6 +65,29 @@ Before changing a specialized subsystem, read its skill: Keep domain-specific invariants in those skills or their direct references; do not duplicate them here. +## Native fermionic tree QR policy + +The native `TreeTensorNetwork` QR policy is centralized in +`TreeTensorNetwork._native_qr_split` and `_native_qr_options`: + +- Every lossless QR split on native Symmray tree tensors must use + `_native_qr_split`; do not add a direct `tensor.split(method="qr")` call to a + native tree route. +- The helper sets `stabilized=False` only for Symmray block-sparse tensors. + Symmray's stabilized QR phase-normalizes each diagonal of `R`; an exact + structural-zero diagonal makes that phase `0 / |0|`, which can become NaN in + `complex64`. Dense tensors retain Quimb's normal stabilized-QR default. +- Network-level canonicalization, which does not expose one tensor at a time, + obtains the same option from `_native_qr_options()` when the tree is + fermionic. Keep this policy aligned if another native canonicalization route + is added. +- Skipping the phase convention is lossless: `Q @ R` is unchanged, and the + resulting `left_inds` isometry metadata remains valid. Native truncating + compression still uses the explicit graded SVD and its configured cutoff. +- This safeguard is scoped to `TreeTensorNetwork` / `TreeOptimizer`. It does + not change the separate `MpsOptimizer` QR implementation or globally patch + Quimb/Symmray. + ## Dependency and backend rules - Prefer public `quimb`, `cotengra`, `cotengrust`, and `autoray` APIs over diff --git a/docs/api/optimizers/tree.md b/docs/api/optimizers/tree.md index 94b64f2..ca4289c 100644 --- a/docs/api/optimizers/tree.md +++ b/docs/api/optimizers/tree.md @@ -139,6 +139,30 @@ fermionic centre uses a one-tensor `TensorNetwork.H` contraction (which applies the required outer-leg phase flips), while an unknown centre falls back to an exact complete doubled-network contraction. +### Native fermionic QR stability + +Native Symmray tree routes use Pepsy's internal +`TreeTensorNetwork._native_qr_split` policy for every lossless QR gauge move, +including two-qubit path threading, edge canonicalization, lossless path +splits, and sub-MPO message routing. The corresponding network-level subtree +canonicalization uses the same policy through `_native_qr_options()`. + +For native block-sparse tensors, the policy passes `stabilized=False` to +Quimb's QR split. Symmray's stabilized QR phase-normalizes each diagonal of +`R`; symmetry can make a diagonal an exact structural zero, so the phase +`0 / |0|` can produce a NaN in `complex64`. Plain QR avoids that undefined +phase while preserving the exact factorization (`Q @ R`) and the tensor's +`left_inds` isometry metadata. This is a gauge choice, not a truncation or a +change to the represented state, and native `complex64` trees therefore do not +need to be promoted to `complex128` as a workaround for this issue. + +The safeguard is tensor-aware: dense TTNs retain Quimb's ordinary stabilized +QR convention. It is internal to the tree implementation, so callers do not +need to pass a QR flag. Native truncating compression continues to use the +graded block SVD and the configured `chi`, `cutoff`, and `cutoff_mode`. This +policy is specific to `TreeTensorNetwork` / `TreeOptimizer`; the separate MPS +optimizer implementation is unchanged. + ## Range / subtree canonicalisation The single orthogonality centre generalises to a connected **canonical region** From 18ef51aadc67afef480f9552ec5dfb87a06d8068 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Sun, 2 Aug 2026 13:51:16 -0600 Subject: [PATCH 53/70] Optimize native tree compression and MPO expectations --- src/pepsy/optimizers/tree/optimizer.py | 139 ++++++++++++++- src/pepsy/optimizers/tree/ttn.py | 225 +++++++++++++++++++++---- tests/test_optimize_tree.py | 54 ++++++ 3 files changed, 376 insertions(+), 42 deletions(-) diff --git a/src/pepsy/optimizers/tree/optimizer.py b/src/pepsy/optimizers/tree/optimizer.py index 11b0c38..63c3aa5 100644 --- a/src/pepsy/optimizers/tree/optimizer.py +++ b/src/pepsy/optimizers/tree/optimizer.py @@ -404,6 +404,10 @@ class TreeOptimizer: Maximum Steiner-subtree size allowed for multi-qubit application and preflight. The default limits the number of recursive local messages; pass ``None`` to disable it. + profile : bool + Whether to collect opt-in kernel timing records in + :meth:`profile_report`. Profiling is disabled by default and adds no + synchronization or timing calls to the normal replay path. tn : TreeTensorNetwork or product MatrixProductState, optional Initial coefficient state. A tree state is copied and canonicalised if needed. Its plan must match ``tree``/``layout`` when it is entangled; @@ -467,7 +471,7 @@ def __init__(self, gates=None, n=None, *, chi=64, max_intermediate_bond=None, max_operator_qubits=_DEFAULT_MAX_OPERATOR_QUBITS, max_subtree_nodes=_DEFAULT_MAX_SUBTREE_NODES, - record_history=True): + record_history=True, profile=False): # Preserve one-shot streams for both queue normalization and automatic # layout discovery. Materializing only inside # ``_normalize_gate_queue`` would leave the finder with an exhausted @@ -667,6 +671,8 @@ def __init__(self, gates=None, n=None, *, chi=64, max_subtree_nodes, "max_subtree_nodes" ) self.record_history = bool(record_history) + self.profile = bool(profile) + self.profile_events = [] self.measurements = [] self.truncation_history = [] self.update_history = [] @@ -733,6 +739,7 @@ def __init__(self, gates=None, n=None, *, chi=64, self.center = self.plan.root else: self._install_tn(tn) + self._attach_profile_sink() self._thread_ind = None self.backend_info() @@ -1275,6 +1282,10 @@ def _install_tn(self, tn): ): self.tn.canonize_around_node_(self.plan.root) + def _attach_profile_sink(self): + """Attach the optimizer's optional timing sink to the live TTN.""" + self.tn._profile_sink = self.profile_events if self.profile else None + def set_tn(self, tn): """Replace the live tree state with a canonical independent copy.""" if not isinstance(tn, TreeTensorNetwork): @@ -1287,6 +1298,7 @@ def set_tn(self, tn): self.normalizations.clear() self.projection_diagnostics.clear() self._truncation_log_survival = 0.0 + self._attach_profile_sink() return self def set_p(self, tn): @@ -2133,7 +2145,14 @@ def _finish_update(self): active = self._active_update if active is None: return + elapsed = time.perf_counter() - active["started_at"] if not self.record_history: + if self.profile: + self.profile_events.append({ + "kind": "update", + "support": active["support"], + "seconds": elapsed, + }) self._active_update = None return start = active["edge_start"] @@ -2193,9 +2212,7 @@ def _finish_update(self): "update": len(self.update_history), "kind": active["kind"], "support": active["support"], - "elapsed_seconds": float( - time.perf_counter() - active["started_at"] - ), + "elapsed_seconds": float(elapsed), "edge_event_indices": list(range(start, len(self.truncation_history))), "edge_count": len(edge_events), "truncated_edges": sum(event["truncated"] for event in edge_events), @@ -2205,6 +2222,12 @@ def _finish_update(self): "max_edge_discarded_weight": max_edge_loss, "max_edge_discarded_fraction": max_edge_fraction, }) + if self.profile: + self.profile_events.append({ + "kind": "update", + "support": active["support"], + "seconds": elapsed, + }) if tracked: local_infidelity = float(relative_loss) self.infidelities.append(float(cumulative_loss)) @@ -3358,6 +3381,7 @@ def _compress_edge_with_diagnostics( self, u, v, *, max_bond=None, cutoff=None, reduced=True, ): """Compress one live tree edge and record its truncation diagnostics.""" + profile_started = time.perf_counter() if self.profile else None max_bond = ( self.chi if max_bond is None else self._normalize_max_bond(max_bond) ) @@ -3379,6 +3403,14 @@ def _compress_edge_with_diagnostics( after_bond=int(self.tn.ind_size(bond_after)), bond_ind=bond_after, max_bond=max_bond, cutoff=cutoff, ) + if profile_started is not None: + self.profile_events.append({ + "kind": "edge_canonize", + "edge": (u, v), + "before_bond": before_bond, + "after_bond": int(self.tn.ind_size(bond_after)), + "seconds": time.perf_counter() - profile_started, + }) return full_spectrum = None if self.track_truncation: @@ -3406,17 +3438,26 @@ def _compress_edge_with_diagnostics( after_bond=after_bond, bond_ind=bond_after, full_spectrum=full_spectrum, max_bond=max_bond, cutoff=cutoff, ) + if profile_started is not None: + self.profile_events.append({ + "kind": "edge_compress", + "edge": (u, v), + "before_bond": before_bond, + "after_bond": after_bond, + "reduced": reduced, + "native": bool(self.tn.fermionic), + "seconds": time.perf_counter() - profile_started, + }) def _metadata_aware_reduction(self, u, v): """Choose one-sided compression when ``v`` is proven isometric. Every caller compresses ``u -> v`` with ``absorb="right"``. If the live ``left_inds`` on ``v`` prove that it is already isometric toward - ``u``, Quimb can SVD only ``u`` and reuse ``v`` directly. Native - graded metadata is accepted only after the TreeTensorNetwork charge - alignment guard; missing or malformed metadata falls back to the - usual two-sided reduction. The native compression kernel still uses - its explicit graded SVD regardless of this advisory ``reduced`` flag. + ``u``, the compression kernel can SVD only ``u`` and reuse ``v`` + directly. Native graded metadata is accepted only after the + TreeTensorNetwork charge alignment guard; missing or malformed + metadata falls back to the usual two-sided reduced compression. """ if self.tn.can_skip_canonize(u, v, absorb="left"): return "left" @@ -3651,6 +3692,56 @@ def apply_submpo(self, submpo, where, *, max_bond=None, cutoff=None): max_bond=max_bond, cutoff=cutoff, ) + def expectation_mpo( + self, submpo, where, *, max_bond=None, cutoff=0.0, + normalized=True, optimize="auto", + ): + """Evaluate ```` through one structured tree-MPO pass. + + The live state is not modified. A private branch routes the MPO with + :meth:`apply_submpo`, so MPO site tensors remain blockwise native on a + Symmray tree and no ``to_dense`` conversion is needed. ``max_bond`` + defaults to this optimizer's ``chi``; pass a larger cap when the + operator application must retain more of the exact MPO-transformed + state. ``cutoff=0.0`` is the default because this is a measurement, + not a variational update. + """ + event_start = len(self.profile_events) + work = self.copy() + work.apply_submpo( + submpo, + where, + max_bond=max_bond, + cutoff=cutoff, + ) + + # The bra and ket are separate TTNs. Keep their physical indices shared + # for the inner product, but rename the ket's virtual bonds so each + # layer remains an ordinary tree and no bond is accidentally merged + # across the bra/ket boundary. + ket = work.tn.copy() + outer = set(ket.outer_inds()) + virtual = { + index + for tensor in ket.tensors + for index in tensor.inds + if index not in outer + } + ket.reindex_({index: qtn.rand_uuid() for index in virtual}) + numerator = (self.tn.H | ket).contract(all, optimize=optimize) + if not normalized: + result = numerator + elif self.tn.fermionic: + result = numerator / self.tn._fermionic_norm_squared() + else: + result = numerator / (self.tn.norm() ** 2) + + if self.profile and len(work.profile_events) > event_start: + self.profile_events.extend( + deepcopy(work.profile_events[event_start:]) + ) + return result + def _apply_submpo_resolved(self, submpo, where, *, max_bond=None, cutoff=None, logical_where=None): """Apply a sub-MPO whose support is already in compact TTN positions.""" @@ -5147,6 +5238,7 @@ def copy(self): max_operator_qubits=self.max_operator_qubits, max_subtree_nodes=self.max_subtree_nodes, record_history=self.record_history, + profile=self.profile, seed=child_seed, run=False, tn=self.tn, @@ -5169,6 +5261,8 @@ def copy(self): other._logical_qubits = list(self._logical_qubits) other._logical_positions = dict(self._logical_positions) other._truncation_log_survival = self._truncation_log_survival + other.profile_events = deepcopy(self.profile_events) + other._attach_profile_sink() return other def get_infidelities(self): @@ -5185,6 +5279,33 @@ def get_infidelity_samples(self): """Return detailed cumulative tree-truncation sample records.""" return self.infidelity_samples + def profile_report(self): + """Return opt-in tree kernel timings grouped by operation kind. + + Construct the optimizer with ``profile=True`` to collect records. + Timing is deliberately kept separate from truncation history so the + normal replay and diagnostic APIs remain unchanged. The returned + ``events`` list is a deep copy and can safely be serialized alongside + a benchmark result. + """ + events = deepcopy(self.profile_events) + grouped = {} + for event in events: + kind = str(event.get("kind", "unknown")) + summary = grouped.setdefault( + kind, {"count": 0, "seconds": 0.0} + ) + summary["count"] += 1 + summary["seconds"] += float(event.get("seconds", 0.0)) + return { + "enabled": self.profile, + "events": events, + "by_kind": grouped, + "total_seconds": float( + sum(float(event.get("seconds", 0.0)) for event in events) + ), + } + def get_normalizations(self): """Return automatic normalization records. diff --git a/src/pepsy/optimizers/tree/ttn.py b/src/pepsy/optimizers/tree/ttn.py index 2aec3e5..ec1239d 100644 --- a/src/pepsy/optimizers/tree/ttn.py +++ b/src/pepsy/optimizers/tree/ttn.py @@ -39,6 +39,7 @@ from __future__ import annotations import re +import time import autoray as ar import numpy as np @@ -405,24 +406,34 @@ def _fermionic_local_expectation( self, operator, where, *, optimize, normalized, ): """Evaluate a native observable with the complete graded exterior.""" - inds = [self.site_ind(site) for site in where] - operated = qtn.tensor_network_gate_inds( - self, - operator, - inds, - contract=False, - tags=[], - info=None, - inplace=False, - ) - numerator = (self.H | operated).contract( - all, - optimize=optimize, - ) - if not normalized: - return numerator - denominator = self._fermionic_norm_squared() - return numerator / denominator + profile_sink = getattr(self, "_profile_sink", None) + profile_started = time.perf_counter() if profile_sink is not None else None + try: + inds = [self.site_ind(site) for site in where] + operated = qtn.tensor_network_gate_inds( + self, + operator, + inds, + contract=False, + tags=[], + info=None, + inplace=False, + ) + numerator = (self.H | operated).contract( + all, + optimize=optimize, + ) + if not normalized: + return numerator + denominator = self._fermionic_norm_squared() + return numerator / denominator + finally: + if profile_started is not None: + profile_sink.append({ + "kind": "native_observable", + "support": tuple(where), + "seconds": time.perf_counter() - profile_started, + }) def _restore_readout_region(self, region): """Restore a dense readout's tracked canonical region.""" @@ -591,16 +602,60 @@ def local_expectations(self, terms, *, optimize="auto", normalized=True): Returns a ``{where: value}`` dict following the iteration order of ``terms``. """ + profile_sink = getattr(self, "_profile_sink", None) + profile_started = time.perf_counter() if profile_sink is not None else None results = {} - for where, operator in terms.items(): - if isinstance(where, Integral): - support = (int(where),) - else: - support = tuple(int(site) for site in where) - results[where] = self.local_expectation( - operator, support, optimize=optimize, normalized=normalized, - ) - return results + try: + for where, operator in terms.items(): + if isinstance(where, Integral): + support = (int(where),) + else: + support = tuple(int(site) for site in where) + results[where] = self.local_expectation( + operator, support, optimize=optimize, normalized=normalized, + ) + return results + finally: + if profile_started is not None: + profile_sink.append({ + "kind": "observable_batch", + "count": len(terms), + "seconds": time.perf_counter() - profile_started, + }) + + def expectation_mpo( + self, mpo, where, *, max_bond=None, cutoff=0.0, + normalized=True, optimize="auto", + ): + """Evaluate a structured MPO expectation without changing this TTN. + + This is a convenience wrapper around + :meth:`TreeOptimizer.expectation_mpo`. The MPO is routed once over + the tree, preserving native Symmray blocks and avoiding a dense + operator on the full support. The transformed-state bond cap defaults + to this TTN's current maximum bond; pass ``max_bond`` explicitly when + a larger measurement workspace is acceptable. + """ + from .optimizer import TreeOptimizer + + current_bond = int(self.max_bond()) + engine = TreeOptimizer( + None, + n=self.nqubits, + tree=self.plan, + state=self, + chi=current_bond if max_bond is None else max_bond, + cutoff=0.0, + run=False, + ) + return engine.expectation_mpo( + mpo, + where, + max_bond=max_bond, + cutoff=cutoff, + normalized=normalized, + optimize=optimize, + ) @property def orthogonality_center(self): @@ -1170,9 +1225,18 @@ def _fermionic_canonize_edge_(self, a, b, absorb): return self def _fermionic_compress_edge_( - self, a, b, *, max_bond, cutoff, cutoff_mode, absorb, + self, a, b, *, max_bond, cutoff, cutoff_mode, absorb, reduced=True, ): - """Compress one native graded tree cut by an explicit two-node SVD.""" + """Compress one native graded tree cut with a reduced graded SVD. + + Native Symmray tensors cannot use Quimb's generic compression helper: + its QR phase convention is not safe for structural zero sectors in + complex64. We nevertheless retain the same reduction that makes the + dense path fast. A proven one-sided isometry lets us SVD only the + non-isometric endpoint; otherwise both endpoints are QR-reduced with + the native zero-sector-safe QR before the small graded core is SVD'd. + The complete two-node graded SVD remains the conservative fallback. + """ if absorb == "right": isometric_node, reduced_node = a, b elif absorb == "left": @@ -1184,6 +1248,99 @@ def _fermionic_compress_edge_( reduced = self.node_tensor(reduced_node) bond = self.bond(isometric_node, reduced_node) left_inds = [index for index in isometric.inds if index != bond] + + # ``reduced="left"`` is the metadata value emitted by the optimizer + # when ``reduced_node`` is already isometric towards ``isometric``. + # In that case the environment is an exact identity and decomposing + # the complete two-node tensor is unnecessary. Validate the proof at + # this low-level boundary as well, since direct TTN callers can pass + # arbitrary reduction hints. + if ( + reduced == "left" + and self.can_skip_canonize( + isometric_node, reduced_node, absorb="left" + ) + ): + kept, remainder = isometric.split( + left_inds=left_inds, + method="svd", + max_bond=max_bond, + cutoff=cutoff, + cutoff_mode=cutoff_mode, + absorb="right", + get="tensors", + bond_ind=bond, + ) + merged = qtn.tensor_contract(remainder, reduced) + isometric.modify( + data=kept.data, + inds=kept.inds, + left_inds=kept.left_inds, + ) + reduced.modify( + data=merged.data, + inds=merged.inds, + left_inds=None, + ) + return self + + if reduced is True: + # Mirror ``qtn.tensor_compress_bond(reduced=True)`` while routing + # both QR decompositions through the native policy above. This + # keeps the expensive SVD on the reduced core and avoids the + # O((Dl * d) x (Dr * d)) full two-node matrix in the common case. + right_inds = [index for index in reduced.inds if index != bond] + left_bond = qtn.rand_uuid() + right_bond = qtn.rand_uuid() + isometric_q, isometric_r = self._native_qr_split( + isometric, + left_inds=left_inds, + right_inds=(bond,), + absorb="right", + cutoff=0.0, + get="tensors", + bond_ind=left_bond, + ) + reduced_l, reduced_q = self._native_qr_split( + reduced, + left_inds=(bond,), + right_inds=right_inds, + absorb="left", + cutoff=0.0, + get="tensors", + bond_ind=right_bond, + ) + core = qtn.tensor_contract(isometric_r, reduced_l) + core_left, core_right = core.split( + left_inds=(left_bond,), + method="svd", + max_bond=max_bond, + cutoff=cutoff, + cutoff_mode=cutoff_mode, + absorb="right", + get="tensors", + bond_ind=bond, + ) + isometric_compressed = qtn.tensor_contract( + isometric_q, core_left, output_inds=isometric.inds, + ) + reduced_compressed = qtn.tensor_contract( + core_right, reduced_q, output_inds=reduced.inds, + ) + isometric.modify( + data=isometric_compressed.data, + inds=isometric_compressed.inds, + left_inds=left_inds, + ) + reduced.modify( + data=reduced_compressed.data, + inds=reduced_compressed.inds, + left_inds=None, + ) + return self + + # Keep the old complete graded split as a compatibility fallback for + # direct callers that provide an unrecognised reduction hint. theta = qtn.tensor_contract(isometric, reduced) kept, remainder = theta.split( left_inds=left_inds, @@ -1281,10 +1438,11 @@ def compress_edge_( :meth:`canonize_edge_`. ``cutoff_mode`` selects Quimb's singular-value cutoff convention. - ``reduced`` is forwarded only on the dense path. Quimb's one-sided - ``"left"`` mode is exact when node ``b`` is already isometric on its - non-shared legs. Native fermionic compression ignores this option and - retains its explicit graded split. + ``reduced`` selects the dense Quimb reduction and the corresponding + native graded reduction. Quimb's one-sided ``"left"`` mode is exact + when node ``b`` is already isometric on its non-shared legs; native + trees use the same proof, with the zero-sector-safe QR policy retained + for the two-sided reduced path. """ bond = self.bond(a, b) before_bond = int(self.ind_size(bond)) @@ -1305,6 +1463,7 @@ def compress_edge_( cutoff=cutoff, cutoff_mode=cutoff_mode, absorb=absorb, + reduced=reduced, ) else: self.compress_between( diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index cf47fbb..2b10a50 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -3109,6 +3109,52 @@ def test_tree_public_submpo_and_pauli_backend_operations(): ) +def test_tree_expectation_mpo_is_batched_and_non_mutating(): + """A structured MPO expectation uses one tree pass and preserves state.""" + h = np.array([[1.0, 1.0], [1.0, -1.0]], dtype=complex) / np.sqrt(2.0) + cnot = np.array( + [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]], + dtype=complex, + ) + zz = np.diag([1.0, -1.0, -1.0, 1.0]) + mpo = qtn.MatrixProductOperator.from_dense( + zz, dims=(2, 2), sites=(0, 1), L=4, + ) + opt = TreeOptimizer([(h, 0), (cnot, (0, 1))], n=4, chi=16) + before = opt.to_dense().copy() + + value = opt.expectation_mpo(mpo, (0, 1), max_bond=16) + + assert value == pytest.approx(1.0) + assert np.allclose(opt.to_dense(), before) + assert opt.tn.validate(check_canonical=True) is opt.tn + + +def test_tree_profile_report_is_opt_in(): + """Kernel timings are empty by default and available when requested.""" + x = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex) + cnot = np.array( + [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]], + dtype=complex, + ) + quiet = TreeOptimizer([(x, 0)], n=4, chi=4) + assert quiet.profile_report() == { + "enabled": False, + "events": [], + "by_kind": {}, + "total_seconds": 0.0, + } + + profiled = TreeOptimizer( + [(x, 0), (cnot, (0, 3))], n=4, chi=4, profile=True, + ) + report = profiled.profile_report() + assert report["enabled"] is True + assert report["events"] + assert report["by_kind"]["update"]["count"] == 2 + assert report["total_seconds"] > 0.0 + + def test_tree_pauli_expectation_and_projection_are_public(): """Pauli expectation/projection share the measurement backend semantics.""" h = np.array([[1.0, 1.0], [1.0, -1.0]], dtype=complex) / np.sqrt(2.0) @@ -4805,6 +4851,14 @@ def dense_vector(opt): for output in outputs[1:]: assert _fidelity(outputs[0], output) > 1 - 1e-10 + before = dense_vector(submpo_opt) + mpo_value = submpo_opt.expectation_mpo( + submpo, (0, 2), max_bond=64, + ) + direct_value = submpo_opt.tn.local_expectation(hopping, (0, 2)) + assert complex(mpo_value) == pytest.approx(complex(direct_value), abs=1e-5) + assert np.allclose(dense_vector(submpo_opt), before) + def test_native_fermionic_submpo_keeps_graded_hub_recovery(monkeypatch): """Native routed Q metadata skips only already-proven graded QR.""" From a91568331632c39d802cfc7a68dac28fb12554c2 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Sun, 2 Aug 2026 13:09:11 -0700 Subject: [PATCH 54/70] Add fixed lattice orders for layout finders --- docs/api/optimizers/mps.md | 13 +++++ docs/api/optimizers/tree.md | 14 ++++++ src/pepsy/__init__.py | 1 + src/pepsy/optimizers/__init__.py | 1 + src/pepsy/optimizers/_layout_orders.py | 69 ++++++++++++++++++++++++++ src/pepsy/optimizers/mps/layout.py | 65 +++++++++++++++--------- src/pepsy/optimizers/tree/layout.py | 28 +++++++++-- tests/test_optimize_mps.py | 11 ++++ tests/test_optimize_tree.py | 16 ++++++ 9 files changed, 192 insertions(+), 26 deletions(-) create mode 100644 src/pepsy/optimizers/_layout_orders.py diff --git a/docs/api/optimizers/mps.md b/docs/api/optimizers/mps.md index b8a7b63..d105af8 100644 --- a/docs/api/optimizers/mps.md +++ b/docs/api/optimizers/mps.md @@ -154,6 +154,19 @@ default to `weight_mode="auto"`: angle metadata when present, otherwise a cheap operator-Schmidt proxy for small dense two-site gates, falling back to count weights. Pass `weight_fn(payload, support, event_type)` for explicit weights. +For a prescribed baseline rather than a searched order, pass an explicit site +permutation as `order`. The returned plan is marked `selected_order="fixed"` +and keeps the original gate stream unchanged: + +```python +zigzag = py.square_lattice_zigzag(6, 6) +fixed_plan = finder.run(order=zigzag) +``` + +`square_lattice_zigzag` scans x across each row and reverses direction on +successive rows. It is a deterministic comparison layout; it performs no +refinement or tensor work. + For compression-oriented selection, pass `objective="compression"`. This uses operator-Schmidt load over every MPS cut crossed by each support, with support span retained as a replay-cost tie-breaker. Exact small dense ranks diff --git a/docs/api/optimizers/tree.md b/docs/api/optimizers/tree.md index 94b64f2..73955c9 100644 --- a/docs/api/optimizers/tree.md +++ b/docs/api/optimizers/tree.md @@ -526,6 +526,20 @@ three-virtual-bond root convention described above. `TreePlan.max_arity()` and `TreePlan.is_binary()` report the shape; `TreePlan.is_strictly_binary()` is the strict two-child-at-every-internal-node predicate. +Both layout finders also accept an explicit site permutation as `order` for a +fixed baseline. For example, this builds the same binary/ternary-root geometry +using a square-lattice snake order without refinement: + +```python +zigzag = py.square_lattice_zigzag(6, 6) +tree_plan = TreeLayoutFinder( + gates, n=36, max_arity=2, top_arity=3, +).run(order=zigzag) +``` + +The explicit order must cover every site exactly once and cannot be combined +with iterable `max_arity` candidate search. + For an automatic arity search, call `finder.recommend_arities((2, 3, 4))` explicitly. The default `TreeOptimizer(gate_stream, n=n, chi=chi)` uses the fixed binary/ternary-root geometry; it does not allocate tensors or perform diff --git a/src/pepsy/__init__.py b/src/pepsy/__init__.py index 4241f91..e69b031 100644 --- a/src/pepsy/__init__.py +++ b/src/pepsy/__init__.py @@ -155,6 +155,7 @@ "TreeLayoutFinder": ".optimizers", "TreeOptimizer": ".optimizers", "TreePlan": ".optimizers", + "square_lattice_zigzag": ".optimizers", "TreeStabOptimizer": ".optimizers", "TreeTensorNetwork": ".optimizers", "compile_stim_circuit": ".optimizers", diff --git a/src/pepsy/optimizers/__init__.py b/src/pepsy/optimizers/__init__.py index 5f5cced..528baa4 100644 --- a/src/pepsy/optimizers/__init__.py +++ b/src/pepsy/optimizers/__init__.py @@ -41,6 +41,7 @@ "TreeLayoutFinder": ".tree", "TreeOptimizer": ".tree", "TreePlan": ".tree", + "square_lattice_zigzag": "._layout_orders", "TreeStabOptimizer": ".tree_stabilizer", "TreeTensorNetwork": ".tree", "CoalescedMeasurementRecord": ".noise", diff --git a/src/pepsy/optimizers/_layout_orders.py b/src/pepsy/optimizers/_layout_orders.py new file mode 100644 index 0000000..1bd95f7 --- /dev/null +++ b/src/pepsy/optimizers/_layout_orders.py @@ -0,0 +1,69 @@ +"""Reusable deterministic site orders for layout diagnostics.""" + +from __future__ import annotations + +from collections.abc import Callable, Iterable + + +def normalize_fixed_order(order: Iterable, sites, *, name="order"): + """Validate an explicit permutation of layout sites.""" + if isinstance(order, (str, bytes)): + raise TypeError(f"{name} must be a site sequence, not a string.") + try: + order = tuple(order) + except TypeError as exc: + raise TypeError(f"{name} must be an iterable site sequence.") from exc + sites = tuple(sites) + if len(order) != len(sites) or set(order) != set(sites): + raise ValueError( + f"{name} must be a permutation of the layout sites; expected " + f"{len(sites)} unique sites." + ) + if len(set(order)) != len(order): + raise ValueError(f"{name} must not contain duplicate sites.") + return order + + +def square_lattice_zigzag( + Lx: int, + Ly: int, + *, + site: Callable[[int, int], int] | None = None, + x_first: bool = True, +): + """Return a deterministic serpentine order for a rectangular lattice. + + The default scans across ``x`` and alternates direction on each successive + ``y`` row. ``x_first=False`` scans across ``y`` and alternates direction on + each successive ``x`` column. ``site`` maps Cartesian coordinates to the + caller's logical site labels and defaults to row-major integer labels. + + This helper only creates a fixed leaf order. It does not infer PBC edges, + and it does not allocate tensors or optimize the order. + """ + if isinstance(Lx, bool) or isinstance(Ly, bool): + raise ValueError("Lx and Ly must be positive integers.") + try: + Lx, Ly = int(Lx), int(Ly) + except (TypeError, ValueError) as exc: + raise ValueError("Lx and Ly must be positive integers.") from exc + if Lx < 1 or Ly < 1: + raise ValueError("Lx and Ly must be positive integers.") + if site is None: + site = lambda x, y: y * Lx + x + if not callable(site): + raise TypeError("site must be callable or None.") + + order = [] + if x_first: + for y in range(Ly): + xs = range(Lx) if y % 2 == 0 else range(Lx - 1, -1, -1) + order.extend(site(x, y) for x in xs) + else: + for x in range(Lx): + ys = range(Ly) if x % 2 == 0 else range(Ly - 1, -1, -1) + order.extend(site(x, y) for y in ys) + return tuple(order) + + +__all__ = ["normalize_fixed_order", "square_lattice_zigzag"] diff --git a/src/pepsy/optimizers/mps/layout.py b/src/pepsy/optimizers/mps/layout.py index be1e788..6f3a796 100644 --- a/src/pepsy/optimizers/mps/layout.py +++ b/src/pepsy/optimizers/mps/layout.py @@ -10,6 +10,7 @@ import numpy as np from ...operators.gates import _normalize_gate_entries +from .._layout_orders import normalize_fixed_order from .._layout_visualization import ( coordinate_lattice_edge_keys, coordinate_lattice_edges, @@ -1580,8 +1581,18 @@ def run( schmidt_max_dim=4, max_operator_qubits=8, ): - """Return a layout plan for the stored gate stream.""" - order_name = _normalize_gate_stream_layout_order(order) + """Return a layout plan for the stored gate stream. + + ``order`` can also be an explicit permutation of the layout sites. + In that case the permutation is returned as a fixed comparison plan + and no layout search or refinement is performed. + """ + fixed_order = None + if isinstance(order, (str, type(None))): + order_name = _normalize_gate_stream_layout_order(order) + else: + fixed_order = normalize_fixed_order(order, self.sites) + order_name = "fixed" objective = _gate_stream_layout_objective(objective) if max_operator_qubits is not None: try: @@ -1622,27 +1633,35 @@ def run( event_weights, ) score_event_weights = event_weights - include_nevergrad = ( - order_name == "auto" or order_name.startswith("nevergrad") - ) - include_kahypar = ( - order_name == "auto" or order_name.startswith("kahypar") - ) - candidates = _gate_stream_layout_candidates( - self.sites, - pair_weights, - refine_passes=refine_passes, - refine_numba=refine_numba, - spectral_dense_max=spectral_dense_max, - recursive_dense_max=recursive_dense_max, - include_nevergrad=include_nevergrad, - nevergrad_budget=nevergrad_budget, - nevergrad_seed=nevergrad_seed, - nevergrad_optimizer=nevergrad_optimizer, - include_kahypar=include_kahypar, - kahypar_config_path=kahypar_config_path, - kahypar_seed=kahypar_seed, - ) + if fixed_order is None: + include_nevergrad = ( + order_name == "auto" or order_name.startswith("nevergrad") + ) + include_kahypar = ( + order_name == "auto" or order_name.startswith("kahypar") + ) + candidates = _gate_stream_layout_candidates( + self.sites, + pair_weights, + refine_passes=refine_passes, + refine_numba=refine_numba, + spectral_dense_max=spectral_dense_max, + recursive_dense_max=recursive_dense_max, + include_nevergrad=include_nevergrad, + nevergrad_budget=nevergrad_budget, + nevergrad_seed=nevergrad_seed, + nevergrad_optimizer=nevergrad_optimizer, + include_kahypar=include_kahypar, + kahypar_config_path=kahypar_config_path, + kahypar_seed=kahypar_seed, + ) + else: + # Keep the input baseline in diagnostics so a fixed order can be + # compared directly with the original logical site order. + candidates = { + "input": list(self.sites), + "fixed": list(fixed_order), + } candidate_stats = {} for name, candidate in candidates.items(): diff --git a/src/pepsy/optimizers/tree/layout.py b/src/pepsy/optimizers/tree/layout.py index 508ce20..6cdb162 100644 --- a/src/pepsy/optimizers/tree/layout.py +++ b/src/pepsy/optimizers/tree/layout.py @@ -41,6 +41,7 @@ _operator_schmidt_rank_info as _mps_operator_schmidt_rank_info, _normalize_weight_mode, ) +from .._layout_orders import normalize_fixed_order from ..mps.optimizer import _control_event_parts as _mps_control_event_parts from .._layout_visualization import ( add_order_colorbar, @@ -181,9 +182,11 @@ def _normalize_layout_search(search): def _normalize_layout_order(order): - """Normalize the optional high-quality layout mode.""" + """Normalize the high-quality mode or preserve an explicit site order.""" if order is None: return None + if not isinstance(order, (str, bytes)): + return tuple(order) name = str(order).replace("-", "_").strip().lower() aliases = { "auto": "quality", @@ -1337,12 +1340,13 @@ class TreeLayoutFinder: across every tree scale. It is the high-quality, Cotengra-inspired mode; ``order="quality"`` selects it automatically and enables its bounded search stages. - order : {None, "quality"}, optional + order : {None, "quality"} or sequence, optional Optional high-quality offline mode. `"quality"` means `objective="full_tree"` and enables bounded greedy leaf refinement, all-scale subtree topology refinement, and hybrid Nevergrad/annealing search. Omitted keeps the fast deterministic - objective selected by `objective`. + objective selected by `objective`. An explicit site permutation builds + a fixed tree without refinement. hybrid_weights : mapping or sequence of three floats, optional Weights for the hybrid path, maximum edge load, and total edge load. The default is ``(1.0, 1.0, 0.25)``. @@ -3082,11 +3086,29 @@ def run( hybrid topology annealing plus Nevergrad leaf search. When Nevergrad is unavailable, it selects dependency-free simulated annealing. Pass ``search=None`` or ``refine=None`` explicitly to disable either stage. + + An explicit site permutation can also be passed as ``order``. This + returns the corresponding fixed tree immediately, without layout + refinement or offline search. """ if order is _DEFAULT_ORDER: order = self.order else: order = _normalize_layout_order(order) + if not isinstance(order, str) and order is not None: + if self.arity_candidates is not None: + raise ValueError( + "an explicit site order requires scalar max_arity; " + "pass one fixed arity instead of arity candidates." + ) + fixed_order = normalize_fixed_order(order, self.leaf_qubits) + return TreePlan.from_order( + fixed_order, + structure=self.structure, + max_arity=self.max_arity, + root_qubit=self.root_qubit, + top_arity=self.top_arity, + ) if order == "quality": if self.objective != "full_tree": self.objective = "full_tree" diff --git a/tests/test_optimize_mps.py b/tests/test_optimize_mps.py index 9253c6a..772510e 100644 --- a/tests/test_optimize_mps.py +++ b/tests/test_optimize_mps.py @@ -692,6 +692,17 @@ def test_mps_optimizer_gate_stream_layout_remaps_long_range_path(): ) +def test_mps_layout_accepts_explicit_fixed_site_order(): + """An explicit site permutation bypasses search and is preserved.""" + gates = [(qu.CNOT(), (0, 3)), (qu.CNOT(), (1, 2))] + order = (2, 0, 3, 1) + plan = py.MpsOptimizer.LayoutFinder(gates, L=4).run(order=order) + + assert plan["selected_order"] == "fixed" + assert plan["site_order"] == order + assert plan["mapped_where"] == ((1, 2), (3, 0)) + + def test_mps_compression_layout_reports_operator_cut_load(): """Compression objective exposes cut-load diagnostics and rank bounds.""" gate = np.eye(8, dtype=complex) diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index cf47fbb..c4e4ad4 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -955,6 +955,22 @@ def test_layout_finder_builds_strict_binary_tree_when_requested(): assert plan.tree_distance(a, b) >= 1 +def test_tree_layout_accepts_explicit_fixed_site_order(): + """An explicit order builds the requested binary/ternary-root tree.""" + order = pepsy.square_lattice_zigzag(2, 2) + finder = TreeLayoutFinder( + [(pepsy.cnot(), (0, 1)), (pepsy.cnot(), (2, 3))], + n=4, + max_arity=2, + top_arity=3, + ) + plan = finder.run(order=order) + + assert tuple(plan.qubit_of_leaf.values()) == order + assert plan.top_arity == 3 + assert plan.is_binary() + + def test_quality_layout_not_worse_than_balanced(): """Entanglement-adapted structure scores no worse than balanced order.""" rng = np.random.default_rng(9) From bc211919e82ec686b78960521b1b18fd62adb778 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Sun, 2 Aug 2026 15:30:03 -0600 Subject: [PATCH 55/70] Fix native Tree central-edge compression --- src/pepsy/optimizers/tree/ttn.py | 70 ++++++++++++++++++++++---------- tests/test_optimize_tree.py | 47 +++++++++++++++++++++ 2 files changed, 96 insertions(+), 21 deletions(-) diff --git a/src/pepsy/optimizers/tree/ttn.py b/src/pepsy/optimizers/tree/ttn.py index ec1239d..fad8e02 100644 --- a/src/pepsy/optimizers/tree/ttn.py +++ b/src/pepsy/optimizers/tree/ttn.py @@ -1232,9 +1232,9 @@ def _fermionic_compress_edge_( Native Symmray tensors cannot use Quimb's generic compression helper: its QR phase convention is not safe for structural zero sectors in complex64. We nevertheless retain the same reduction that makes the - dense path fast. A proven one-sided isometry lets us SVD only the - non-isometric endpoint; otherwise both endpoints are QR-reduced with - the native zero-sector-safe QR before the small graded core is SVD'd. + dense path fast. A proven one-sided isometry first gets a native QR + reduction of the active endpoint, then only the small graded core is + SVD'd; otherwise both endpoints are QR-reduced before the core SVD. The complete two-node graded SVD remains the conservative fallback. """ if absorb == "right": @@ -1244,8 +1244,9 @@ def _fermionic_compress_edge_( else: raise ValueError("absorb must be 'right' or 'left'.") + reduction_hint = reduced isometric = self.node_tensor(isometric_node) - reduced = self.node_tensor(reduced_node) + reduced_tensor = self.node_tensor(reduced_node) bond = self.bond(isometric_node, reduced_node) left_inds = [index for index in isometric.inds if index != bond] @@ -1256,40 +1257,66 @@ def _fermionic_compress_edge_( # this low-level boundary as well, since direct TTN callers can pass # arbitrary reduction hints. if ( - reduced == "left" + reduction_hint == "left" and self.can_skip_canonize( isometric_node, reduced_node, absorb="left" ) ): - kept, remainder = isometric.split( + # The destination endpoint is already an isometry toward the + # active endpoint, so only the active endpoint needs reducing. + # QR it first: after fusing its external legs, a central Tree + # tensor can be thousands by thousands even though its shared + # bond is only O(chi). The QR leaves a core whose right dimension + # is the live bond, avoiding a full SVD of that large matrix. + reduced_bond = qtn.rand_uuid() + isometric_q, isometric_r = self._native_qr_split( + isometric, left_inds=left_inds, + right_inds=(bond,), + absorb="right", + cutoff=0.0, + get="tensors", + bond_ind=reduced_bond, + ) + compressed_bond = qtn.rand_uuid() + core_left, core_right = isometric_r.split( + left_inds=(reduced_bond,), method="svd", max_bond=max_bond, cutoff=cutoff, cutoff_mode=cutoff_mode, absorb="right", get="tensors", - bond_ind=bond, + bond_ind=compressed_bond, + ) + kept = qtn.tensor_contract( + isometric_q, core_left, + ) + merged = qtn.tensor_contract( + core_right, reduced_tensor, ) - merged = qtn.tensor_contract(remainder, reduced) + kept.reindex_({compressed_bond: bond}) + merged.reindex_({compressed_bond: bond}) isometric.modify( data=kept.data, inds=kept.inds, - left_inds=kept.left_inds, + left_inds=left_inds, ) - reduced.modify( + reduced_tensor.modify( data=merged.data, inds=merged.inds, left_inds=None, ) return self - if reduced is True: + if reduction_hint is True: # Mirror ``qtn.tensor_compress_bond(reduced=True)`` while routing # both QR decompositions through the native policy above. This # keeps the expensive SVD on the reduced core and avoids the # O((Dl * d) x (Dr * d)) full two-node matrix in the common case. - right_inds = [index for index in reduced.inds if index != bond] + right_inds = [ + index for index in reduced_tensor.inds if index != bond + ] left_bond = qtn.rand_uuid() right_bond = qtn.rand_uuid() isometric_q, isometric_r = self._native_qr_split( @@ -1302,7 +1329,7 @@ def _fermionic_compress_edge_( bond_ind=left_bond, ) reduced_l, reduced_q = self._native_qr_split( - reduced, + reduced_tensor, left_inds=(bond,), right_inds=right_inds, absorb="left", @@ -1322,17 +1349,17 @@ def _fermionic_compress_edge_( bond_ind=bond, ) isometric_compressed = qtn.tensor_contract( - isometric_q, core_left, output_inds=isometric.inds, + isometric_q, core_left, ) reduced_compressed = qtn.tensor_contract( - core_right, reduced_q, output_inds=reduced.inds, + core_right, reduced_q, ) isometric.modify( data=isometric_compressed.data, inds=isometric_compressed.inds, left_inds=left_inds, ) - reduced.modify( + reduced_tensor.modify( data=reduced_compressed.data, inds=reduced_compressed.inds, left_inds=None, @@ -1341,7 +1368,7 @@ def _fermionic_compress_edge_( # Keep the old complete graded split as a compatibility fallback for # direct callers that provide an unrecognised reduction hint. - theta = qtn.tensor_contract(isometric, reduced) + theta = qtn.tensor_contract(isometric, reduced_tensor) kept, remainder = theta.split( left_inds=left_inds, method="svd", @@ -1357,7 +1384,7 @@ def _fermionic_compress_edge_( inds=kept.inds, left_inds=kept.left_inds, ) - reduced.modify( + reduced_tensor.modify( data=remainder.data, inds=remainder.inds, left_inds=None, @@ -1433,9 +1460,10 @@ def compress_edge_( """Compress the tree edge ``a -> b`` in place. Dense/nonfermionic trees delegate to Quimb's ``compress_between``. - Native fermionic trees explicitly SVD the complete two-node tensor. - The tracked :attr:`orthogonality_center` advances as for - :meth:`canonize_edge_`. + Native fermionic trees use the same reduced-core decomposition while + routing every lossless factorization through the zero-sector-safe + native QR helper. The tracked :attr:`orthogonality_center` advances as + for :meth:`canonize_edge_`. ``cutoff_mode`` selects Quimb's singular-value cutoff convention. ``reduced`` selects the dense Quimb reduction and the corresponding diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index 2b10a50..21ed9c2 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -5056,6 +5056,53 @@ def traced_compress(*args, **kwargs): assert calls +def test_native_one_sided_compression_qr_reduces_before_svd(monkeypatch): + """A proven native isometry sends only its reduced core to the SVD.""" + pytest.importorskip("symmray") + fermion = pepsy.Fermion( + spinful=True, + symmetry="U1U1", + dtype="complex128", + ) + plan = TreePlan.from_order(range(4), structure="balanced") + ttn = pepsy.ps_to_ttn( + 4, + tree=plan, + fermion=fermion, + occupations=((1, 0), (0, 1), (1, 0), (0, 1)), + dtype="complex128", + ) + target = plan.leaf_of_qubit[0] + ttn.shift_orthogonality_center(target) + source = next( + nid for nid, toward in ttn.isometry_map().items() + if toward == target and ttn.can_skip_canonize(nid, toward) + ) + + split_methods = [] + original_split = qtn.Tensor.split + + def traced_split(self, *args, **kwargs): + method = kwargs.get("method") + if method in {"qr", "svd"} and hasattr(self.data, "blocks"): + split_methods.append((method, tuple(self.shape))) + return original_split(self, *args, **kwargs) + + monkeypatch.setattr(qtn.Tensor, "split", traced_split) + ttn._fermionic_compress_edge_( + target, + source, + max_bond=64, + cutoff=1e-10, + cutoff_mode="rel", + absorb="right", + reduced="left", + ) + + assert [method for method, _shape in split_methods] == ["qr", "svd"] + assert ttn.validate(check_canonical=True) is ttn + + def test_tree_stable_labels_route_submpo_by_payload_sites(monkeypatch): """Stable logical labels do not disable native structured MPO routing.""" x = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex) From 2b496d6cdf064bd1b13e2a1a22b85a2f381a2593 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Sun, 2 Aug 2026 15:40:39 -0600 Subject: [PATCH 56/70] Document native Tree compression performance --- .github/skills/tree-optimizer/SKILL.md | 12 ++++-- .../references/performance-layout.md | 32 ++++++++++++++++ docs/api/optimizers/tree.md | 38 ++++++++++++++++++- 3 files changed, 77 insertions(+), 5 deletions(-) diff --git a/.github/skills/tree-optimizer/SKILL.md b/.github/skills/tree-optimizer/SKILL.md index 5adaeb2..d2966a8 100644 --- a/.github/skills/tree-optimizer/SKILL.md +++ b/.github/skills/tree-optimizer/SKILL.md @@ -243,9 +243,15 @@ forwarded to the SVD and raises `TypeError`). Unique `rand_uuid()` bonds avoid Native fermionic trees take an isolated version of this kernel: `_fermionic_thread_hop` explicitly calls the native Symmray QR and carries its -graded factor, while `TreeTensorNetwork._fermionic_compress_edge_` forms the -two-node tensor and performs the native block SVD. Dense/nonfermionic trees -retain the generic Quimb edge wrappers. +graded factor. `TreeTensorNetwork._fermionic_compress_edge_` uses a reduced +graded core: when the destination endpoint is proven isometric, it QR-splits +the active endpoint and SVDs only its `R` factor; otherwise it QR-reduces both +endpoints and SVDs their contracted core. Only an unrecognised reduction hint +falls back to the complete two-node SVD. The reduction hint must remain +separate from the destination tensor object, and the one-sided split must use +fresh intermediate bond names before restoring the live edge label. See the +performance reference for the algebra and profiling evidence. Dense/nonfermionic +trees retain the generic Quimb edge wrappers. ### Sibling-leaf fast path (`_apply_2q_sibling_factors`) diff --git a/.github/skills/tree-optimizer/references/performance-layout.md b/.github/skills/tree-optimizer/references/performance-layout.md index bfbbf31..0540c82 100644 --- a/.github/skills/tree-optimizer/references/performance-layout.md +++ b/.github/skills/tree-optimizer/references/performance-layout.md @@ -13,6 +13,38 @@ Tree Optimizer skill so the upload-facing `SKILL.md` stays concise. `threadpoolctl` when available. Only raise `threads` in a large-`chi` regime. - The self-healing tid cache (`_nid_to_tid`, `_tid`) validates cached tensor ids against `self.tn.tensor_map`; a stale entry is recomputed safely. +- **Native central-edge compression.** A native compression call receives a + reduction hint separately from its destination tensor. For a proven + one-sided reduction (`reduced="left"`), let `A` be the active endpoint and + `B` the destination endpoint, with `B†B = I` on the non-shared legs. The + implementation computes the lossless graded factorization + `A = Q_A R_A`, then SVDs only `R_A = U S V†`. It installs + `Q_A U` on `A` and absorbs `S V†` into `B`. Thus the expensive SVD scales + with the active endpoint's QR carry and the live bond, rather than the + full fused two-node tensor. The proof is structural: + `can_skip_canonize(A, B, absorb="left")` must accept the destination's + `left_inds` and aligned Symmray charge maps. +- If that one-sided proof is absent but `reduced=True`, both endpoints are + QR-reduced and only the contracted `R_A L_B` core is SVD'd. Unknown hints + retain the complete two-node graded SVD as a compatibility fallback. The + truncating step always remains the explicit native block SVD with the + configured `max_bond`, `cutoff`, and `cutoff_mode`; the native + `stabilized=False` policy applies only to lossless QR. +- The one-sided path uses fresh intermediate QR/SVD bond names because the + original live edge label is still present in `R_A` while that factor is + decomposed. After both contractions, the new compressed bond is reindexed + to the original live edge. Reusing the old label during the SVD creates a + repeated index and can route the contraction into an unsupported Symmray + hyperedge. +- A 6x6, 48-gate, chi=64, complex64 Torch-CPU calibration with 12 Torch/tree + threads reduced Tree evolution from 127.77 s to 6.21 s. The same run's MPS + evolution was 2.85 s, so the remaining Tree/MPS ratio was 2.18x. The + pre-fix central SVDs were 4096x4096; after the QR reduction they were + 384x384. Layout planning (~26.6 s in that run) is setup cost and is not + included in the evolution comparison. The remaining replay gap is mainly + gate-update/threading/contraction work; use `profile=True` to inspect + `update` envelopes and nested `edge_compress` events before changing path + geometry or observable contractions. - Dense path and subtree routing preserve each QR-produced Q tensor's `left_inds`. Canonical recovery therefore recognizes an already-isometric routed branch without repeating its decomposition or entering Quimb's dense diff --git a/docs/api/optimizers/tree.md b/docs/api/optimizers/tree.md index ca4289c..4da4417 100644 --- a/docs/api/optimizers/tree.md +++ b/docs/api/optimizers/tree.md @@ -131,8 +131,12 @@ operations that prove canonicality is preserved. Native fermionic trees use a separate graded edge path. Centre moves explicitly QR-split the Symmray tensor and absorb the native carry into the next node; -edge compression explicitly forms the two-node tensor and performs its native -block SVD. Dense and nonfermionic trees continue to use Quimb's generic +edge compression uses a reduced graded core whenever the destination endpoint +is already proven isometric: the active endpoint is QR-split first, and only +its `R` factor is sent to the truncating native block SVD. If that proof is +absent, both endpoints are QR-reduced and their contracted core is SVD'd; +only an unrecognised reduction hint forms the complete two-node tensor. +Dense and nonfermionic trees continue to use Quimb's generic `canonize_between` / `compress_between` wrappers. A graded exterior is not assumed to be an ordinary Frobenius identity for readout: a known native fermionic centre uses a one-tensor `TensorNetwork.H` contraction (which applies @@ -163,6 +167,36 @@ graded block SVD and the configured `chi`, `cutoff`, and `cutoff_mode`. This policy is specific to `TreeTensorNetwork` / `TreeOptimizer`; the separate MPS optimizer implementation is unchanged. +### Native central-edge compression and profiling + +For a compression from active endpoint `A` to destination endpoint `B`, the +one-sided native path is valid only when `B` is structurally proven isometric +toward `A` (`can_skip_canonize(A, B, absorb="left")`). It uses + +```text +A = Q_A R_A +R_A = U S V† +new_A = Q_A U +new_B = (S V†) B +``` + +The first QR is lossless and uses `_native_qr_split`; the second factorization +is the only truncating SVD. This avoids SVD'ing the fully fused `A B` tensor, +which can be thousands by thousands at moderate `chi` even when the live +edge is small. The implementation keeps the reduction hint separate from the +destination tensor, uses fresh intermediate bond labels while the old edge is +still present in `R_A`, and restores the original live edge label after the +factors are contracted. `reduced=True` uses the analogous two-sided QR/core +reduction. A positive cutoff never turns this into metadata-only compression. + +On the calibrated 6x6 χ=64 complex64 Torch-CPU run (12 threads, 48 gates), +Tree evolution improved from 127.77 s to 6.21 s; MPS took 2.85 s in the same +post-fix run. Thus this fix removes the pathological Tree kernel, but Tree is +still about 2.18x slower for this prefix. The saved profile showed 276 edge +compression events totaling 2.40 s inside 48 update envelopes totaling 6.18 s; +the next optimization target is the gate update/threading/contraction path, +especially the central edges, rather than further route-length tuning. + ## Range / subtree canonicalisation The single orthogonality centre generalises to a connected **canonical region** From 23a54492b68af8b15ea2469e721623ae9084b548 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Sun, 2 Aug 2026 23:10:36 -0600 Subject: [PATCH 57/70] Optimize and clarify Tree replay defaults --- .github/skills/tree-optimizer/SKILL.md | 13 +- .../references/performance-layout.md | 48 +++ docs/api/optimizers/tree.md | 65 +++- src/pepsy/optimizers/tree/optimizer.py | 284 +++++++++++++----- src/pepsy/optimizers/tree/ttn.py | 186 ++++++++++-- tests/test_optimize_tree.py | 200 ++++++++++++ 6 files changed, 698 insertions(+), 98 deletions(-) diff --git a/.github/skills/tree-optimizer/SKILL.md b/.github/skills/tree-optimizer/SKILL.md index d2966a8..c93c7a3 100644 --- a/.github/skills/tree-optimizer/SKILL.md +++ b/.github/skills/tree-optimizer/SKILL.md @@ -410,7 +410,18 @@ for trajectory simulation without forming a density matrix: ## Performance and layout -Keep the thread cap, self-healing tensor-id cache, copy semantics, and +The public performance defaults are ``mode="auto"`` (direct two-site +threading), ``threads=1`` and ``subtree_workers=1`` (small-TTN operations avoid +oversubscription), ``profile=False``, and ``track_truncation=False`` (no +diagnostic spectrum SVDs). The low-level +``TreeTensorNetwork.compress_edge_`` default is the same ``cutoff_mode="rsum2"`` +used by ``TreeOptimizer``. A ``track_truncation=True`` warning is intentional: +it identifies the extra diagnostic work; backend conversion warnings and +legacy-mode deprecations are the other actionable warning classes. + +Dense and native trees share the direct one-edge contraction, immutable path +cache, routed-isometry reuse, and proof-forwarding optimizations. Keep the +thread cap, self-healing tensor-id cache, copy semantics, and TreeLayoutFinder objective plumbing intact. The detailed performance and non-binary layout contract is in [`references/performance-layout.md`](references/performance-layout.md); read diff --git a/.github/skills/tree-optimizer/references/performance-layout.md b/.github/skills/tree-optimizer/references/performance-layout.md index 0540c82..b7fd700 100644 --- a/.github/skills/tree-optimizer/references/performance-layout.md +++ b/.github/skills/tree-optimizer/references/performance-layout.md @@ -57,6 +57,54 @@ Tree Optimizer skill so the upload-facing `SKILL.md` stays concise. native compression remains an explicit graded SVD. The network derives orientation views directly from live tensors; do not cache a duplicate map in the optimizer. +- **Update-path bookkeeping.** A compression sweep now computes the live + isometry proof once per truncating edge and passes that proof into the native + compressor; lossless QR edges do not perform a reduction-proof lookup at + all. Private two-site gate-factor output labels are stable within the local + update, avoiding per-factor UUID/reindex work while the routed operator bond + remains fresh. Subtree QR messages are merged by destination, so sibling + messages landing at one hub use one multi-tensor contraction instead of + rebuilding that hub once per child. Dense message waves reuse one worker + pool; native fermionic waves remain serial, but use the same grouped merge + without changing graded block semantics. +- **Two-site routing prefactors.** The immutable geodesic for a repeated + qubit support is cached and reversed when the current centre chooses the + opposite endpoint as the source, so the cache never assumes a gauge + location. Ordinary adjacent two-tensor merges in gate threading and local + operator absorption use a direct backend ``tensordot``; Symmray retains its + graded fermionic contraction semantics and unsupported hyperedges fall back + to Quimb's general contraction path. +- **Public performance defaults.** ``mode="auto"`` selects the direct + two-site kernel, ``threads=1`` and ``subtree_workers=1`` avoid oversubscription + on small tree tensors, ``profile=False`` avoids timing overhead, and + ``track_truncation=False`` avoids diagnostic spectrum probes. The low-level + ``TreeTensorNetwork.compress_edge_`` API uses the same ``rsum2`` cutoff-mode + default as ``TreeOptimizer``. Dense and native trees share these routing, + contraction, path-cache, and proof-reuse optimizations; native-only code is + limited to graded QR/SVD semantics and the complex64 zero-sector safeguard. +- **Native complex64 QR stability.** Torch can return NaNs for a finite, + rank-deficient Symmray charge block whose norm has decayed into the + ``1e-9`` range. Native QR therefore uses a block-local power-of-two scale + for small complex64 blocks and divides the triangular factor by the same + scale. This leaves ``Q @ R`` unchanged, keeps zero-charge sectors finite, + and avoids promoting the complete replay to complex128. The native + ``reduced="right"`` branch now mirrors Quimb's one-sided path: it SVDs only + the active endpoint and contracts that factor into the already-left- + isometric endpoint; unproven ``right``, ``False``, and ``lazy`` hints retain + the conservative full-SVD fallback. +- On the same 6x6 χ=64 complex64 Torch-CPU harness, the post-compression-fix + Tree evolution was 6.21 s in the saved baseline; the update-path pass ran + in 4.90--5.15 s on repeated warm-cache runs. Threaded BLAS and cache warmth + affect absolute wall time, so use the profile envelopes for comparisons; + the optimization preserves the 48-gate state and χ/cutoff semantics. +- The full 468-gate ``6x6_nsteps=0`` replay then completed with + ``track_truncation=False`` and profiling enabled: one run measured 20.39 s + for MPS evolution, 137.11 s for Tree evolution, and 25.79 s for Tree layout + planning (Tree/MPS evolution ratio 6.72x). This is the stability baseline; + layout planning is reported separately and the remaining update envelopes + are the next prefactor target. The normalized Tree/MPS state fidelity in the + same χ=64 run was 0.590; because the two geometries truncate different + bonds, this is an accuracy diagnostic rather than an exact-gauge check. - `copy()` shares the immutable `TreePlan`, owns `self.tn.copy()`, resets the tid cache, and derives a deterministic child seed for an independent RNG. diff --git a/docs/api/optimizers/tree.md b/docs/api/optimizers/tree.md index 4ece1d5..7de72d6 100644 --- a/docs/api/optimizers/tree.md +++ b/docs/api/optimizers/tree.md @@ -194,8 +194,46 @@ Tree evolution improved from 127.77 s to 6.21 s; MPS took 2.85 s in the same post-fix run. Thus this fix removes the pathological Tree kernel, but Tree is still about 2.18x slower for this prefix. The saved profile showed 276 edge compression events totaling 2.40 s inside 48 update envelopes totaling 6.18 s; -the next optimization target is the gate update/threading/contraction path, -especially the central edges, rather than further route-length tuning. +it identified the gate update/threading/contraction path, especially the +central edges, rather than route-length tuning, as the next target. + +The update path now avoids repeating the same native isometry proof: each +truncating edge validates its `left_inds`/charge-map proof once, while a +lossless QR edge skips the reduction lookup entirely. Two-site factors use +private stable output labels, removing per-factor UUID allocation and one +metadata reindex. In multi-site Tree/MPO updates, independent QR messages +landing at the same node are contracted as one batch; dense waves reuse their +worker pool, while native fermionic routing remains serial for Symmray safety. +These changes preserve the complete-gate-before-truncation rule and the +configured χ/cutoff semantics. On repeated warm-cache runs of the same +harness, Tree evolution was 4.90--5.15 s; absolute timings vary with BLAS +thread state, so profile envelopes remain the authoritative comparison. + +The remaining two-site update bookkeeping is also shared across arbitrary +gate streams: immutable qubit-support geodesics are cached and re-oriented +against the live centre for each gate, while ordinary one-edge tensor merges +use a direct backend ``tensordot``. Symmray dispatches through its graded +fermionic contraction implementation; unusual hyperedges still use Quimb's +general contraction path. This removes repeated path construction and +contraction-expression setup without changing the routed QR/SVD sequence. + +Native complex64 QR also applies a reversible power-of-two scale separately +to small Symmray charge blocks. This avoids a Torch QR failure on finite, +rank-deficient blocks around ``1e-9`` without changing ``Q @ R`` or promoting +the replay to complex128. Native ``reduced="right"`` now has the matching +one-sided endpoint-SVD path; unproven ``right``, ``False``, and ``lazy`` modes +continue to use the conservative complete SVD. + +The exact 6x6 ``nsteps=0`` stream (468 gates, χ=64, complex64 Torch CPU, +12 threads, ``track_truncation=False``) subsequently completed without the +previous gate-235 NaN. In one profiled run, MPS evolution took 20.39 s and +Tree evolution 137.11 s; Tree layout planning was a separate 25.79 s, giving +an evolution ratio of 6.72x. This confirms stability, not parity: the remaining +Tree cost is concentrated in the per-gate update envelopes and needs further +backend/profile-guided reduction. The same run's normalized Tree/MPS state +fidelity was 0.590; at χ=64 this measures different truncation histories on +the two geometries, not a QR gauge error. Compare observables or increase χ +when using this number as an accuracy diagnostic. ## Range / subtree canonicalisation @@ -294,6 +332,29 @@ for shared frontends. Tree-edge SVD. Its defaults, `cutoff=1e-10` and `cutoff_mode="rsum2"`, match Quimb's open-boundary `MatrixProductState.gate_with_submpo` compression path. `"rel"` remains available as a relative largest-singular-value threshold. +The same defaults are used by the lower-level +`TreeTensorNetwork.compress_edge_` API, so constructing the state directly and +replaying it through `TreeOptimizer` does not silently change the truncation +criterion. + +### Performance-oriented defaults and warnings + +The default replay configuration is intended for production evolution: +`mode="auto"` uses the direct routed two-site kernel, `threads=1` avoids +oversubscribing the small tree contractions, `subtree_workers=1` keeps the +serial path allocation-free, `profile=False` avoids timing overhead, and +`track_truncation=False` avoids full-spectrum diagnostic SVDs. `record_history` +and `track_infidelity` retain the established API defaults; the latter only +adds norm-based progress readouts when a progress bar is requested. + +Warnings are reserved for an actionable behavior change: enabling +`track_truncation=True` emits one diagnostic-performance warning, explicit +backend/dtype/device conversion emits one compatibility warning per conversion +signature, and legacy mode selectors emit deprecation warnings. Dense and +native paths share the direct one-edge contraction, path-cache, routing, and +proof-reuse optimizations. Only the QR phase safeguard and graded reduced-core +SVD are native Symmray specializations; dense arrays continue through Quimb's +ordinary QR/SVD with the same cutoff, path, and truncation semantics. `TreeOptimizer.apply_submpo(...)` is the public form for an explicit MPO of arbitrary support. It losslessly QR-routes its virtual bonds, then uses its diff --git a/src/pepsy/optimizers/tree/optimizer.py b/src/pepsy/optimizers/tree/optimizer.py index 63c3aa5..f59deae 100644 --- a/src/pepsy/optimizers/tree/optimizer.py +++ b/src/pepsy/optimizers/tree/optimizer.py @@ -73,7 +73,7 @@ _normalize_time_window, _submpo_schmidt_rank_bound, ) -from .ttn import TreeTensorNetwork +from .ttn import TreeTensorNetwork, _contract_two_tensors __all__ = ["TreeOptimizer"] @@ -292,6 +292,14 @@ def _normalize_measure_axes(pauli, where): class TreeOptimizer: """Replay a bundled gate stream on a rooted tree tensor network. + The constructor defaults are performance-oriented: ordinary two-site + gates use the direct routed kernel (``mode="auto"``), small tree + contractions are capped to one BLAS/OpenMP thread (``threads=1``), and + full singular-spectrum diagnostics are disabled + (``track_truncation=False``). Enable those diagnostics explicitly when + collecting truncation reports; the one-time warning in that case is + intentional because it describes the additional SVD work. + Parameters ---------- gates : bundled gate stream, optional @@ -691,6 +699,11 @@ def __init__(self, gates=None, n=None, *, chi=64, # reuse cannot return a stale factorization. self._gate_factor_cache = {} self._gate_factor_cache_limit = 64 + # Gate supports recur frequently in circuit streams. Cache only the + # immutable geometry path; centre-dependent source orientation is + # selected on every update below. + self._two_site_path_cache = {} + self._two_site_path_cache_limit = 256 self._active_update = None self._truncation_log_survival = 0.0 @@ -1030,10 +1043,10 @@ def _warn_track_truncation_slow(self): if self.track_truncation and not self._track_warning_emitted: warnings.warn( "TreeOptimizer track_truncation=True enables complete " - "singular-spectrum probes and can add extra SVDs for each " - "compressed edge; this diagnostic mode can substantially " - "slow replay. Use track_truncation=False for performance " - "runs.", + "singular-spectrum probes for edges that may truncate. " + "Lossless zero-cutoff edges remain QR-only, but this " + "diagnostic mode can substantially slow replay. Use " + "track_truncation=False for performance runs.", UserWarning, stacklevel=3, ) @@ -1272,6 +1285,7 @@ def _install_tn(self, tn): tn.validate() self.tn = tn.copy() self.plan = self.tn.plan + self._two_site_path_cache.clear() self.tn.validate() self.n = self.tn.nqubits self._logical_qubits = list(range(self.n)) @@ -1471,6 +1485,34 @@ def _nearest_anchor(self, nodes): ) return nodes[0] + def _cached_two_site_path(self, qa, qb): + """Return ``(leaf_a, leaf_b, path_a_to_b)`` for a gate support. + + The path depends only on the immutable tree plan, while the direction + in which it is traversed depends on the current orthogonality centre. + Cache the undirected support path and orient it for the caller so + repeated gate supports avoid rebuilding the same geodesic without + making a stale centre assumption. + """ + key = (qa, qb) if qa < qb else (qb, qa) + cached = self._two_site_path_cache.get(key) + if cached is None: + qlo, qhi = key + leaf_lo = self.plan.node_of_qubit[qlo] + leaf_hi = self.plan.node_of_qubit[qhi] + path = tuple(self.plan.node_path(leaf_lo, leaf_hi)) + cached = (leaf_lo, leaf_hi, path) + if len(self._two_site_path_cache) >= self._two_site_path_cache_limit: + self._two_site_path_cache.pop( + next(iter(self._two_site_path_cache)) + ) + self._two_site_path_cache[key] = cached + + leaf_lo, leaf_hi, path = cached + if qa < qb: + return leaf_lo, leaf_hi, path + return leaf_hi, leaf_lo, path[::-1] + def shift_orthogonality_center(self, node): """Move the orthogonality centre to ``node`` along the tree geodesic. @@ -1902,6 +1944,7 @@ def run_pilot(job): selected_plan = final_candidates[final_selected_name]["plan"] if install: self.plan = selected_plan + self._two_site_path_cache.clear() self.tn = self._remount_product_state(self.tn) self.center = self.plan.root self.layout_finder = final_finder @@ -2853,7 +2896,12 @@ def _two_site_mpo_factors(self, submpo, qa, qb, *, site_where=None): outputs = {} for qubit, (factor_template, upper, lower) in raw_factors.items(): factor = factor_template.copy() - output = qtn.rand_uuid() + # These factors are private to this update and are consumed before + # they enter the live tree. A stable, private output label avoids + # two UUID allocations and one metadata reindex per local factor; + # the shared routed bond remains fresh because it enters the live + # tree during threading. + output = f"_pepsy_mpo_out_{qubit}" factor.reindex_({ upper: output, lower: self._phys(qubit), @@ -2935,16 +2983,15 @@ def _cached_direct_gate_factors( ) thread_ind = qtn.rand_uuid() - output_a, output_b = qtn.rand_uuid(), qtn.rand_uuid() left = left_template.copy() right = right_template.copy() + output_a = "_pepsy_gate_out_a" + output_b = "_pepsy_gate_out_b" left.reindex_({ - "_pepsy_gate_out_a": output_a, "_pepsy_gate_in_a": pa, "_pepsy_gate_thread": thread_ind, }) right.reindex_({ - "_pepsy_gate_out_b": output_b, "_pepsy_gate_in_b": pb, "_pepsy_gate_thread": thread_ind, }) @@ -2967,8 +3014,7 @@ def _apply_2q_factors_impl( one canonical compression sweep. """ plan = self.plan - la = plan.node_of_qubit[qa] - lb = plan.node_of_qubit[qb] + la, lb, path = self._cached_two_site_path(qa, qb) parent = plan.parent.get(la) if ( plan.is_leaf(la) @@ -2990,24 +3036,29 @@ def _apply_2q_factors_impl( ) source_node = plan.node_of_qubit[source] destination_node = plan.node_of_qubit[destination] + if source_node != path[0]: + path = path[::-1] self._move_center(source_node) self._thread_ind = thread_ind try: source_tensor = self.tn.tensor_map[self._tid(source_node)] - merged_source = qtn.tensor_contract( - source_tensor, factors[source] + merged_source = _contract_two_tensors( + source_tensor, + factors[source], + shared_ind=self._phys(source), ).reindex_({outputs[source]: self._phys(source)}) source_tensor.modify( data=merged_source.data, inds=merged_source.inds, ) - path = plan.node_path(source_node, destination_node) for u, v in zip(path, path[1:]): self._thread_hop(u, v) destination_tensor = self.tn.tensor_map[self._tid(destination_node)] - merged_destination = qtn.tensor_contract( - factors[destination], destination_tensor, + merged_destination = _contract_two_tensors( + factors[destination], + destination_tensor, + shared_ind=self._phys(destination), ).reindex_({outputs[destination]: self._phys(destination)}) destination_tensor.modify( data=merged_destination.data, inds=merged_destination.inds, @@ -3043,10 +3094,14 @@ def _apply_2q_sibling_factors( e_la = self._bond_name(la, parent) e_lb = self._bond_name(lb, parent) - merged_a = qtn.tensor_contract(tla, factors[qa]).reindex_( + merged_a = _contract_two_tensors( + tla, factors[qa], shared_ind=pa, + ).reindex_( {outputs[qa]: pa} ) - merged_b = qtn.tensor_contract(tlb, factors[qb]).reindex_( + merged_b = _contract_two_tensors( + tlb, factors[qb], shared_ind=pb, + ).reindex_( {outputs[qb]: pb} ) blob = qtn.tensor_contract(merged_a, tp, merged_b) @@ -3105,7 +3160,7 @@ def _thread_hop(self, u, v): absorb="right", get="tensors", ) - merged_v = qtn.tensor_contract(carry, tv) + merged_v = _contract_two_tensors(carry, tv, shared_ind=edge) # ``keep`` is the exact Q factor pointing toward ``v``. Preserve that # isometry metadata so a later canonical walk can recognize the tensor # without repeating the same QR decomposition. @@ -3143,7 +3198,7 @@ def _fermionic_thread_hop(self, u, v): cutoff=0.0, get="tensors", ) - merged_v = qtn.tensor_contract(carry, tv) + merged_v = _contract_two_tensors(carry, tv, shared_ind=edge) tu.modify( data=keep.data, inds=keep.inds, @@ -3379,6 +3434,7 @@ def _subtree_cutoff_for_size(self, before_bond, *, max_bond, cutoff): def _compress_edge_with_diagnostics( self, u, v, *, max_bond=None, cutoff=None, reduced=True, + reduction_proven=False, ): """Compress one live tree edge and record its truncation diagnostics.""" profile_started = time.perf_counter() if self.profile else None @@ -3430,6 +3486,7 @@ def _compress_edge_with_diagnostics( self.tn.compress_edge_( u, v, max_bond=max_bond, cutoff=cutoff, absorb="right", cutoff_mode=self.cutoff_mode, reduced=reduced, + _reduction_proven=reduction_proven, ) bond_after = self.tn.bond(u, v) after_bond = int(self.tn.ind_size(bond_after)) @@ -3463,6 +3520,31 @@ def _metadata_aware_reduction(self, u, v): return "left" return True + def _edge_reduction(self, u, v, *, max_bond, cutoff): + """Return ``(reduction, proof)`` for one compression edge. + + A lossless edge only needs a QR move, so asking the native tensor to + validate a truncating reduction hint would be wasted work. For a + truncating native edge, the proof is passed through to the low-level + compressor so the same expensive charge-map check is not repeated. + """ + requested_cutoff = self.cutoff if cutoff is None else float(cutoff) + effective_max_bond = ( + self.chi + if max_bond is None + else self._normalize_max_bond(max_bond) + ) + if requested_cutoff == 0.0 and ( + effective_max_bond is None + or int(self.tn.ind_size(self.tn.bond(u, v))) <= effective_max_bond + ): + # The low-level edge routine will take its lossless QR branch and + # never consume this hint. Keep the one-sided marker for + # diagnostics/observers while avoiding the proof lookup entirely. + return "left", False + reduced = self._metadata_aware_reduction(u, v) + return reduced, reduced == "left" + def _compress_path( self, path, *, max_bond=None, cutoff=None, preserve_subcap=False, ): @@ -3482,9 +3564,12 @@ def _compress_path( max_bond=max_bond, cutoff=cutoff, ) + reduced, reduction_proven = self._edge_reduction( + v, u, max_bond=max_bond, cutoff=edge_cutoff, + ) self._compress_edge_with_diagnostics( v, u, max_bond=max_bond, cutoff=edge_cutoff, - reduced=self._metadata_aware_reduction(v, u), + reduced=reduced, reduction_proven=reduction_proven, ) self.center = path[0] @@ -3524,12 +3609,17 @@ def descend(node, parent): if neighbor in snodes and neighbor != parent ) for child in children: + child_cutoff = edge_cutoff(node, child) + reduced, reduction_proven = self._edge_reduction( + node, child, max_bond=max_bond, cutoff=child_cutoff, + ) self._compress_edge_with_diagnostics( node, child, max_bond=max_bond, - cutoff=edge_cutoff(node, child), - reduced=self._metadata_aware_reduction(node, child), + cutoff=child_cutoff, + reduced=reduced, + reduction_proven=reduction_proven, ) descend(child, node) self.tn.canonize_edge_(child, node, absorb="right") @@ -3563,59 +3653,91 @@ def _route_subtree_messages( workers = 1 pending = list(order) - while pending: - ready = [ - (index, u, v) - for index, (u, v) in enumerate(pending) - if u not in {dst for _, dst in pending} - ] - if not ready: - raise RuntimeError( - "subtree peel order contains a cyclic message dependency." - ) - if workers == 1: - ready = ready[:1] - - def split_message(item): - index, u, v = item - state_bond = self.tn.bond(u, v) - left_inds = [ - ix for ix in local[u].inds - if ix != state_bond and ix not in operator_inds[u] - ] - new_bond = f"_ttn_mpo_route_{token}_{u}_{v}" - kept, message = self._qr_route_message( - local[u], left_inds, bond_ind=new_bond, - ) - return index, u, v, state_bond, new_bond, kept, message + pool = None + if workers > 1 and not self.tn.fermionic and len(order) > 1: + from concurrent.futures import ThreadPoolExecutor - if workers > 1 and len(ready) > 1: - from concurrent.futures import ThreadPoolExecutor + pool = ThreadPoolExecutor( + max_workers=min(workers, len(order)), + thread_name_prefix="pepsy-ttn-qr", + ) - with ThreadPoolExecutor( - max_workers=min(workers, len(ready)), - thread_name_prefix="pepsy-ttn-qr", - ) as pool: - results = list(pool.map(split_message, ready)) - else: - results = [split_message(item) for item in ready] + try: + while pending: + pending_destinations = {v for _, v in pending} + ready = [ + (index, u, v) + for index, (u, v) in enumerate(pending) + if u not in pending_destinations + ] + if not ready: + raise RuntimeError( + "subtree peel order contains a cyclic message dependency." + ) + if workers == 1: + ready = ready[:1] + + def split_message(item): + index, u, v = item + state_bond = self.tn.bond(u, v) + left_inds = [ + ix for ix in local[u].inds + if ix != state_bond and ix not in operator_inds[u] + ] + new_bond = f"_ttn_mpo_route_{token}_{u}_{v}" + kept, message = self._qr_route_message( + local[u], left_inds, bond_ind=new_bond, + ) + return index, u, v, state_bond, new_bond, kept, message - # Merge in peel-order order. At this point worker tensors are - # private and no destination tensor has been modified yet, so - # equal-destination messages remain deterministic. - for index, u, v, state_bond, new_bond, kept, message in sorted( - results - ): - local[u] = kept - local[v] = qtn.tensor_contract(local[v], message) - state_inds[v].discard(state_bond) - state_inds[v].add(new_bond) - operator_inds[v] = set(local[v].inds) - state_inds[v] - removed = {index for index, *_ in results} - pending = [ - edge for index, edge in enumerate(pending) - if index not in removed - ] + if pool is not None and len(ready) > 1: + results = list(pool.map(split_message, ready)) + else: + results = [split_message(item) for item in ready] + + # Keep worker results deterministic, then merge all messages + # landing on the same destination in one contraction. The + # old serial loop contracted a busy hub once per child, which + # repeatedly rebuilt the same destination tensor and paid the + # contraction dispatch/reindex bookkeeping for every message. + # Messages in one ready wave have disjoint source edges and + # are therefore independent until this grouped merge. + results = sorted(results) + by_destination = {} + for result in results: + by_destination.setdefault(result[2], []).append(result) + + for destination, destination_results in by_destination.items(): + messages = [result[-1] for result in destination_results] + if len(messages) == 1: + local[destination] = qtn.tensor_contract( + local[destination], messages[0] + ) + else: + local[destination] = qtn.tensor_contract( + local[destination], *messages + ) + for ( + _, source, _, state_bond, new_bond, kept, _ + ) in destination_results: + # The source tensors are private QR outputs and can + # be installed independently of the destination + # contraction. + local[source] = kept + state_inds[destination].discard(state_bond) + state_inds[destination].add(new_bond) + operator_inds[destination] = ( + set(local[destination].inds) + - state_inds[destination] + ) + removed = {index for index, *_ in results} + pending = [ + edge for index, edge in enumerate(pending) + if index not in removed + ] + finally: + if pool is not None: + pool.shutdown() def _install_routed_subtree(self, local, snodes, hub): """Install routed tensors and recover their proven hub centre. @@ -4128,7 +4250,9 @@ def _try_apply_native_submpo( operator_inds[nid] = set(op_t.inds) - { self._phys(q) + "*", self._phys(q) } - local[nid] = qtn.tensor_contract(state_t, op_t).reindex_( + local[nid] = _contract_two_tensors( + state_t, op_t, shared_ind=self._phys(q), + ).reindex_( {self._phys(q) + "*": self._phys(q)} ) except (KeyError, TypeError, ValueError): @@ -4176,7 +4300,12 @@ def _apply_factorized_subtree_operator_impl( op_bonds["physical"][q], self._phys(q), ) - local[nid] = qtn.tensor_contract(state_t, op_t) + if q is not None and q in where: + local[nid] = _contract_two_tensors( + state_t, op_t, shared_ind=self._phys(q), + ) + else: + local[nid] = qtn.tensor_contract(state_t, op_t) if q is not None and q in where: local[nid].reindex_({f"{self._phys(q)}*": self._phys(q)}) operator_inds[nid] = set(local[nid].inds) - state_inds[nid] @@ -4280,7 +4409,9 @@ def _apply_product_pauli_projector_impl( self._as_state_backend(operators, warn=False), inds=(p + "*", p, branch), ) - local[nid] = qtn.tensor_contract(state_t, op_t).reindex_( + local[nid] = _contract_two_tensors( + state_t, op_t, shared_ind=p, + ).reindex_( {p + "*": p} ) branch_index[nid] = branch @@ -5182,6 +5313,7 @@ def cap(self, q, vec, *, absorb="left", compact_labels=True, try: self.tn.cap_qubit_(q, self._as_state_backend(vec)) self.plan = self.tn.plan + self._two_site_path_cache.clear() self.n = self.tn.nqubits remaining = [label for label in self._logical_qubits if label != logical_q] if compact_labels: diff --git a/src/pepsy/optimizers/tree/ttn.py b/src/pepsy/optimizers/tree/ttn.py index fad8e02..d4e327b 100644 --- a/src/pepsy/optimizers/tree/ttn.py +++ b/src/pepsy/optimizers/tree/ttn.py @@ -44,6 +44,7 @@ import autoray as ar import numpy as np import quimb.tensor as qtn +from quimb.tensor.decomp import qr_stabilized as _quimb_qr_stabilized from quimb.tensor.tensor_core import TensorNetwork from numbers import Integral @@ -53,6 +54,50 @@ __all__ = ["TreeTensorNetwork"] +def _native_qr_block_scaled(array, **kwargs): + """QR one native charge block after a reversible power-of-two scaling. + + Torch's complex64 QR can return NaNs for a rank-deficient block whose + entries are small (around ``1e-9``) even though the block is finite. Native + Symmray QR is blockwise, so scaling each block independently is exact: the + isometric factor is unchanged and the triangular factor is divided by the + same scalar afterwards. Power-of-two scaling avoids introducing an extra + rounding step into the block values. + """ + opts = dict(kwargs) + opts.pop("method", None) + opts.pop("fn", None) + if ar.get_dtype_name(array) != "complex64": + return _quimb_qr_stabilized(array, **opts) + + block_max = to_float(ar.do("max", ar.do("abs", array))) + if not np.isfinite(block_max) or block_max == 0.0: + # Preserve the original failure behaviour for non-finite input, while + # allowing genuinely empty structural sectors through unchanged. + return _quimb_qr_stabilized(array, **opts) + + # Values above this scale are not affected by the low-norm complex64 QR + # failure and avoid an unnecessary multiply/divide pair. + if block_max >= 2.0**-8: + return _quimb_qr_stabilized(array, **opts) + + _, exponent = np.frexp(block_max) + scale = float(np.ldexp(1.0, -int(exponent))) + left, singular_values, right = _quimb_qr_stabilized( + array * scale, **opts, + ) + + # ``absorb='left'`` is the LQ orientation: the left factor carries the + # scale. All other QR orientations carry it in the right factor. + absorb = opts.get("absorb", "right") + if absorb in {-1, "left", "Us,VH", "lfactor", "Us"}: + if left is not None: + left = left / scale + elif right is not None: + right = right / scale + return left, singular_values, right + + try: # quimb renamed/removed this generic-vector base across releases. from quimb.tensor import TensorNetworkGenVector except ImportError: # pragma: no cover - exercised with older quimb releases @@ -117,6 +162,52 @@ def _is_symmray_array(value): return hasattr(value, "blocks") and hasattr(value, "indices") +def _contract_two_tensors(left, right, *, shared_ind=None): + """Contract two tensors along one ordinary shared index cheaply. + + Tree routing and native reduced compression repeatedly contract adjacent + tensors along one virtual edge. Dispatching that ordinary operation + directly to the active backend avoids rebuilding a generic Quimb/Cotengra + expression while retaining Symmray's graded fermionic ``tensordot``. + Unusual hyperedges fall back to Quimb's general contraction path. + """ + if shared_ind is None: + shared = tuple(qtn.bonds(left, right)) + if len(shared) != 1: + return qtn.tensor_contract(left, right) + shared_ind = shared[0] + elif shared_ind not in left.inds or shared_ind not in right.inds: + return qtn.tensor_contract(left, right) + + left_rest = tuple(ind for ind in left.inds if ind != shared_ind) + right_rest = tuple(ind for ind in right.inds if ind != shared_ind) + if set(left_rest) & set(right_rest): + return qtn.tensor_contract(left, right) + + axes = ((left.inds.index(shared_ind),), (right.inds.index(shared_ind),)) + try: + if _is_symmray_array(left.data): + data = ar.do( + "tensordot", + left.data, + right.data, + axes=axes, + preserve_array=True, + ) + else: + data = ar.do( + "tensordot", left.data, right.data, axes=axes, + ) + except (AttributeError, NotImplementedError, TypeError): + return qtn.tensor_contract(left, right) + + return qtn.Tensor( + data=data, + inds=left_rest + right_rest, + tags=left.tags | right.tags, + ) + + def _visible_len(s): """Length of ``s`` ignoring any ANSI colour escape sequences.""" return len(_ANSI_RE.sub("", s)) @@ -1185,6 +1276,8 @@ def _native_qr_options(self, tensor=None): def _native_qr_split(self, tensor, **kwargs): """Perform a QR split with the native graded zero-sector safeguard.""" kwargs.update(self._native_qr_options(tensor)) + if _is_symmray_array(tensor.data): + kwargs.setdefault("fn", _native_qr_block_scaled) kwargs.setdefault("method", "qr") return tensor.split(**kwargs) @@ -1211,7 +1304,7 @@ def _fermionic_canonize_edge_(self, a, b, absorb): cutoff=0.0, get="tensors", ) - merged = qtn.tensor_contract(carry, reduced) + merged = _contract_two_tensors(carry, reduced, shared_ind=bond) isometric.modify( data=kept.data, inds=kept.inds, @@ -1226,6 +1319,7 @@ def _fermionic_canonize_edge_(self, a, b, absorb): def _fermionic_compress_edge_( self, a, b, *, max_bond, cutoff, cutoff_mode, absorb, reduced=True, + reduction_proven=False, ): """Compress one native graded tree cut with a reduced graded SVD. @@ -1258,8 +1352,11 @@ def _fermionic_compress_edge_( # arbitrary reduction hints. if ( reduction_hint == "left" - and self.can_skip_canonize( - isometric_node, reduced_node, absorb="left" + and ( + reduction_proven + or self.can_skip_canonize( + isometric_node, reduced_node, absorb="left" + ) ) ): # The destination endpoint is already an isometry toward the @@ -1289,12 +1386,8 @@ def _fermionic_compress_edge_( get="tensors", bond_ind=compressed_bond, ) - kept = qtn.tensor_contract( - isometric_q, core_left, - ) - merged = qtn.tensor_contract( - core_right, reduced_tensor, - ) + kept = _contract_two_tensors(isometric_q, core_left) + merged = _contract_two_tensors(core_right, reduced_tensor) kept.reindex_({compressed_bond: bond}) merged.reindex_({compressed_bond: bond}) isometric.modify( @@ -1309,6 +1402,49 @@ def _fermionic_compress_edge_( ) return self + if ( + reduction_hint == "right" + and ( + reduction_proven + or self.can_skip_canonize( + isometric_node, reduced_node, absorb="right" + ) + ) + ): + # Mirror Quimb's ``reduced="right"`` branch: the active endpoint + # is split directly while the already-left-isometric endpoint is + # reused. Native SVD handles the charge blocks independently, so + # no complete two-node contraction is formed. + right_inds = [ + index for index in reduced_tensor.inds if index != bond + ] + compressed_bond = qtn.rand_uuid() + core_left, core_right = reduced_tensor.split( + left_inds=(bond,), + right_inds=right_inds, + method="svd", + max_bond=max_bond, + cutoff=cutoff, + cutoff_mode=cutoff_mode, + absorb="right", + get="tensors", + bond_ind=compressed_bond, + ) + kept = _contract_two_tensors(isometric, core_left) + kept.reindex_({compressed_bond: bond}) + core_right.reindex_({compressed_bond: bond}) + isometric.modify( + data=kept.data, + inds=kept.inds, + left_inds=left_inds, + ) + reduced_tensor.modify( + data=core_right.data, + inds=core_right.inds, + left_inds=None, + ) + return self + if reduction_hint is True: # Mirror ``qtn.tensor_compress_bond(reduced=True)`` while routing # both QR decompositions through the native policy above. This @@ -1337,7 +1473,7 @@ def _fermionic_compress_edge_( get="tensors", bond_ind=right_bond, ) - core = qtn.tensor_contract(isometric_r, reduced_l) + core = _contract_two_tensors(isometric_r, reduced_l) core_left, core_right = core.split( left_inds=(left_bond,), method="svd", @@ -1348,10 +1484,10 @@ def _fermionic_compress_edge_( get="tensors", bond_ind=bond, ) - isometric_compressed = qtn.tensor_contract( + isometric_compressed = _contract_two_tensors( isometric_q, core_left, ) - reduced_compressed = qtn.tensor_contract( + reduced_compressed = _contract_two_tensors( core_right, reduced_q, ) isometric.modify( @@ -1368,7 +1504,7 @@ def _fermionic_compress_edge_( # Keep the old complete graded split as a compatibility fallback for # direct callers that provide an unrecognised reduction hint. - theta = qtn.tensor_contract(isometric, reduced_tensor) + theta = _contract_two_tensors(isometric, reduced_tensor, shared_ind=bond) kept, remainder = theta.split( left_inds=left_inds, method="svd", @@ -1408,7 +1544,9 @@ def _track_edge_center(self, a, b, absorb, *, previous=None): else: self._canonical_region = None - def canonize_edge_(self, a, b, absorb="right"): + def canonize_edge_( + self, a, b, absorb="right", *, _isometry_proven=False, + ): """Canonicalise across the tree edge ``a -> b`` in place. Dense/nonfermionic trees delegate to Quimb's ``canonize_between``; @@ -1418,7 +1556,10 @@ def canonize_edge_(self, a, b, absorb="right"): the required isometry are metadata-only moves for both dense and native graded arrays. """ - if self.can_skip_canonize(a, b, absorb=absorb): + if ( + _isometry_proven + or self.can_skip_canonize(a, b, absorb=absorb) + ): # The local QR is already represented by the tensor's proven # ``left_inds``. Keep centre bookkeeping honest, but do not touch # tensor data or invalidate the norm cache. @@ -1453,9 +1594,10 @@ def compress_edge_( *, max_bond=None, cutoff=1e-10, - cutoff_mode="rel", + cutoff_mode="rsum2", absorb="right", reduced=True, + _reduction_proven=False, ): """Compress the tree edge ``a -> b`` in place. @@ -1465,7 +1607,10 @@ def compress_edge_( native QR helper. The tracked :attr:`orthogonality_center` advances as for :meth:`canonize_edge_`. - ``cutoff_mode`` selects Quimb's singular-value cutoff convention. + ``cutoff_mode`` selects Quimb's singular-value cutoff convention. The + default ``"rsum2"`` matches :class:`TreeOptimizer` and Quimb's + open-boundary MPS gate-application default; use ``"rel"`` explicitly + for a relative largest-singular-value threshold. ``reduced`` selects the dense Quimb reduction and the corresponding native graded reduction. Quimb's one-sided ``"left"`` mode is exact when node ``b`` is already isometric on its non-shared legs; native @@ -1492,6 +1637,7 @@ def compress_edge_( cutoff_mode=cutoff_mode, absorb=absorb, reduced=reduced, + reduction_proven=_reduction_proven, ) else: self.compress_between( @@ -1607,8 +1753,10 @@ def _recover_center_from_region(self, region, target, *, absorb="right"): a, b = node, neighbour else: a, b = neighbour, node - if not self.can_skip_canonize(a, b, absorb=absorb): - self.canonize_edge_(a, b, absorb=absorb) + proof = self.can_skip_canonize(a, b, absorb=absorb) + self.canonize_edge_( + a, b, absorb=absorb, _isometry_proven=proof, + ) remaining.remove(node) self._canonical_region = frozenset({target}) diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index dbb14f1..768bad5 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -1,5 +1,6 @@ """Tests for the tree-tensor-network gate simulator (:class:`TreeOptimizer`).""" +import inspect import sys import types @@ -14,6 +15,8 @@ TreePlan, TreeTensorNetwork, ) +from pepsy.optimizers.tree.optimizer import _contract_two_tensors +from pepsy.optimizers.tree.ttn import _native_qr_block_scaled # -- exact statevector reference ---------------------------------------------- @@ -235,6 +238,7 @@ def test_two_site_modes_reuse_path_isometries_for_compression( def traced_compress_edge( u, v, *, max_bond=None, cutoff=None, reduced=True, + reduction_proven=False, ): assert opt.tn.can_skip_canonize(u, v, absorb="left") reductions.append(reduced) @@ -244,6 +248,7 @@ def traced_compress_edge( max_bond=max_bond, cutoff=cutoff, reduced=reduced, + reduction_proven=reduction_proven, ) monkeypatch.setattr( @@ -271,6 +276,24 @@ def traced_compress_edge( assert opt.tn.validate(check_canonical=True) is opt.tn +def test_lossless_path_skips_reduction_proof_lookup(monkeypatch): + """A QR-only path does not query truncating-compression metadata.""" + x = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex) + opt = TreeOptimizer( + None, n=6, chi=64, cutoff=0.0, mode="direct", run=False, + ) + + def fail_reduction_lookup(*args, **kwargs): + raise AssertionError("lossless path queried truncation metadata") + + monkeypatch.setattr( + opt, "_metadata_aware_reduction", fail_reduction_lookup, + ) + opt.apply_2q(np.kron(x, x), 0, 5) + + assert opt.tn.validate(check_canonical=True) is opt.tn + + def test_dense_path_one_sided_compression_matches_full_reduction(monkeypatch): """Reusing routed Q tensors is exact even when the path truncates.""" rng = np.random.default_rng(921) @@ -550,6 +573,7 @@ def traced_compress_edge(tn, a, b, **kwargs): def force_full_reduction( u, v, *, max_bond=None, cutoff=None, reduced=True, + reduction_proven=False, ): del reduced return full_compress( @@ -558,6 +582,7 @@ def force_full_reduction( max_bond=max_bond, cutoff=cutoff, reduced=True, + reduction_proven=reduction_proven, ) monkeypatch.setattr( @@ -1465,7 +1490,22 @@ def test_optimizer_defaults_to_congestion_layout_for_replay_performance(): assert opt.layout_objective == "congestion" assert opt.layout_finder.objective == "congestion" + assert opt.mode == "auto" + assert opt.threads == 1 + assert opt.subtree_workers == 1 assert opt.track_truncation is False + assert opt.track_infidelity is True + assert opt.cutoff_mode == "rsum2" + assert opt.profile is False + + +def test_tree_state_compression_default_matches_optimizer(): + """The low-level TTN edge API uses the same cutoff convention.""" + parameter = inspect.signature( + TreeTensorNetwork.compress_edge_ + ).parameters["cutoff_mode"] + + assert parameter.default == "rsum2" def test_layout_recommends_arity_and_reports_tree_shape(): @@ -2521,6 +2561,58 @@ def traced_split(tensor, *args, **kwargs): assert sum(key[0] == "direct" for key in opt._gate_factor_cache) == 1 +def test_repeated_two_site_support_reuses_only_immutable_path(): + """Path caching never assumes the current centre or traversal direction.""" + opt = TreeOptimizer(None, n=8, chi=16, cutoff=0.0, run=False) + + leaf_a, leaf_b, path = opt._cached_two_site_path(0, 7) + reverse_a, reverse_b, reverse_path = opt._cached_two_site_path(7, 0) + cached_a, cached_b, cached_path = opt._cached_two_site_path(0, 7) + + assert (leaf_a, leaf_b) == (cached_a, cached_b) + assert (reverse_a, reverse_b) == (leaf_b, leaf_a) + assert reverse_path == path[::-1] + assert cached_path is path + + +def test_adjacent_two_tensor_contract_matches_quimb(): + """The direct one-edge backend contraction preserves Quimb ordering.""" + rng = np.random.default_rng(912) + left = qtn.Tensor( + rng.standard_normal((2, 4, 3)), inds=("a", "edge", "b"), + ) + right = qtn.Tensor( + rng.standard_normal((4, 5, 2)), inds=("edge", "c", "d"), + ) + + fast = _contract_two_tensors(left, right, shared_ind="edge") + reference = qtn.tensor_contract(left, right) + + assert fast.inds == reference.inds + assert np.allclose(fast.data, reference.data) + + +def test_adjacent_dense_contract_avoids_generic_quimb_dispatch(monkeypatch): + """The shared dense hot path does not rebuild a generic contraction.""" + left = qtn.Tensor( + np.arange(12.0).reshape(3, 4), inds=("a", "edge"), + ) + right = qtn.Tensor( + np.arange(20.0).reshape(4, 5), inds=("edge", "b"), + ) + + def unexpected_generic_contract(*args, **kwargs): + raise AssertionError("dense one-edge contraction used generic Quimb") + + monkeypatch.setattr(qtn, "tensor_contract", unexpected_generic_contract) + fast = _contract_two_tensors(left, right, shared_ind="edge") + + assert fast.inds == ("a", "b") + np.testing.assert_allclose( + fast.data, np.asarray(left.data) @ np.asarray(right.data) + ) + + def test_parallel_subtree_messages_match_serial(): """Independent dense QR message waves preserve the serial result.""" rng = np.random.default_rng(818) @@ -3146,6 +3238,32 @@ def test_tree_expectation_mpo_is_batched_and_non_mutating(): assert opt.tn.validate(check_canonical=True) is opt.tn +def test_tree_subtree_route_batches_sibling_messages(monkeypatch): + """Independent leaf messages landing at one node use one contraction.""" + n = 4 + plan = TreePlan.from_order(range(n), structure="balanced") + opt = TreeOptimizer( + None, n=n, tree=plan, chi=16, cutoff=0.0, + subtree_workers=2, run=False, + ) + identity = np.eye(2**n, dtype=complex) + before = opt.to_dense().copy() + original_contract = qtn.tensor_contract + grouped_calls = [] + + def traced_contract(*tensors, **kwargs): + if len(tensors) >= 3: + grouped_calls.append(len(tensors)) + return original_contract(*tensors, **kwargs) + + monkeypatch.setattr(qtn, "tensor_contract", traced_contract) + opt.apply_subtree_operator(identity, tuple(range(n))) + + assert grouped_calls + assert np.allclose(opt.to_dense(), before) + assert opt.tn.validate(check_canonical=True) is opt.tn + + def test_tree_profile_report_is_opt_in(): """Kernel timings are empty by default and available when requested.""" x = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex) @@ -4938,6 +5056,7 @@ def traced_install(local, snodes, hub): def traced_compress_edge( u, v, *, max_bond=None, cutoff=None, reduced=True, + reduction_proven=False, ): compressions.append(reduced) return compress_edge( @@ -4946,6 +5065,7 @@ def traced_compress_edge( max_bond=max_bond, cutoff=cutoff, reduced=reduced, + reduction_proven=reduction_proven, ) monkeypatch.setattr(candidate, "_install_routed_subtree", traced_install) @@ -5119,6 +5239,86 @@ def traced_split(self, *args, **kwargs): assert ttn.validate(check_canonical=True) is ttn +@pytest.mark.parametrize( + ("symmetry", "spinful", "occupations"), + [ + ("U1", False, (1, 0, 1, 0)), + ("U1U1", True, ((1, 0), (0, 1), (1, 0), (0, 1))), + ], +) +@pytest.mark.parametrize("state_dtype", ["complex64", "complex128"]) +def test_native_one_sided_and_two_sided_compression_fidelity( + symmetry, spinful, occupations, state_dtype, +): + """Native left/right/two-sided reductions preserve the state and gauge.""" + pytest.importorskip("symmray") + fermion = pepsy.Fermion( + spinful=spinful, + symmetry=symmetry, + dtype=state_dtype, + ) + plan = TreePlan.from_order(range(4), structure="balanced") + base = pepsy.ps_to_ttn( + 4, + tree=plan, + fermion=fermion, + occupations=occupations, + dtype=state_dtype, + ) + target = plan.leaf_of_qubit[0] + base.shift_orthogonality_center(target) + source = next( + nid for nid, toward in base.isometry_map().items() + if toward == target and base.can_skip_canonize(nid, target) + ) + cutoff = 1e-10 if state_dtype == "complex128" else 1e-7 + fidelity_floor = 1 - (1e-10 if state_dtype == "complex128" else 1e-5) + + for a, b, reduced in ( + (source, target, "right"), + (target, source, "left"), + (target, source, True), + ): + candidate = base.copy() + candidate._fermionic_compress_edge_( + a, + b, + max_bond=64, + cutoff=cutoff, + cutoff_mode="rel", + absorb="right", + reduced=reduced, + ) + assert float(pepsy.tn_fidelity(base, candidate)) > fidelity_floor + assert candidate.validate(check_canonical=True) is candidate + + +def test_native_complex64_qr_scales_low_norm_rank_deficient_block(): + """Native QR keeps tiny complex64 charge blocks finite and exact.""" + torch = pytest.importorskip("torch") + block = torch.zeros((16, 18), dtype=torch.complex64) + generator = torch.Generator().manual_seed(17) + left = torch.randn((16, 12), generator=generator) * 1e-9 + right = torch.randn((12, 18), generator=generator) + block = (left @ right).to(torch.complex64) + + q, _, r = _native_qr_block_scaled( + block, + method="qr", + absorb="right", + stabilized=False, + ) + + assert torch.isfinite(q).all() + assert torch.isfinite(r).all() + torch.testing.assert_close( + q @ r, + block, + rtol=2e-4, + atol=1e-12, + ) + + def test_tree_stable_labels_route_submpo_by_payload_sites(monkeypatch): """Stable logical labels do not disable native structured MPO routing.""" x = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex) From 0b5d24bd9e7c3d5b0bebbc21f608995a85599ffe Mon Sep 17 00:00:00 2001 From: rezaquant Date: Mon, 3 Aug 2026 10:16:17 -0700 Subject: [PATCH 58/70] Make CI lint checks compatible with lazy exports --- .github/workflows/ci.yml | 6 +++--- src/pepsy/optimizers/tree/layout.py | 3 ++- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d161f2a..d83e29c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -44,9 +44,9 @@ jobs: - name: Validate source syntax run: python -m compileall -q -f src tests - name: Check Python lint - run: | - python -m pyflakes src/pepsy tests - python -m ruff check src/pepsy tests + # Ruff's F rules cover the Pyflakes checks and honor noqa markers for + # intentional lazy, type-only, and compatibility exports. + run: python -m ruff check src/pepsy tests - name: Run tests with coverage if: matrix.python-version == '3.12' && matrix.profile == 'core' run: pytest -q --cov=pepsy --cov-report=term-missing diff --git a/src/pepsy/optimizers/tree/layout.py b/src/pepsy/optimizers/tree/layout.py index 6cdb162..8859868 100644 --- a/src/pepsy/optimizers/tree/layout.py +++ b/src/pepsy/optimizers/tree/layout.py @@ -25,6 +25,7 @@ from __future__ import annotations from collections.abc import Mapping +from importlib import import_module from numbers import Integral import autoray as ar @@ -202,7 +203,7 @@ def _normalize_layout_order(order): def _nevergrad_available(): """Return whether the optional Nevergrad dependency can be imported.""" try: - import nevergrad # pylint: disable=import-outside-toplevel,unused-import + import_module("nevergrad") except ImportError: return False return True From fbf7e6e3e3312e39ac1ad7b5b2048226587617c8 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Mon, 3 Aug 2026 10:16:53 -0700 Subject: [PATCH 59/70] Add Tree bond and QR profiling diagnostics * Add Tree bond and QR profiling diagnostics * Keep Tree bond diagnostics opt-in --- docs/api/optimizers/tree.md | 27 ++++- src/pepsy/optimizers/tree/optimizer.py | 149 ++++++++++++++++++++++++- tests/test_optimize_tree.py | 63 +++++++++++ 3 files changed, 231 insertions(+), 8 deletions(-) diff --git a/docs/api/optimizers/tree.md b/docs/api/optimizers/tree.md index 7de72d6..cb227f6 100644 --- a/docs/api/optimizers/tree.md +++ b/docs/api/optimizers/tree.md @@ -235,6 +235,28 @@ fidelity was 0.590; at χ=64 this measures different truncation histories on the two geometries, not a QR gauge error. Compare observables or increase χ when using this number as an accuracy diagnostic. +#### QR/hop and bond-growth diagnostics + +Construct `TreeOptimizer(..., profile=True)` to split the update envelope into +timed `thread_hop`, `edge_canonize`, and `edge_compress` events. The +`thread_hop` events are the exact, lossless QR carry moves; `edge_compress` +events are the truncating SVD work. These timings are nested inside the +per-update envelope and should not be added as independent wall-clock totals. + +For a dimension-level report, also pass +`track_bond_diagnostics=True`. `bond_diagnostic_report()` then records +`transient_max_bond` during routing/factorization and `live_max_bond_after` +after the compression sweep. The former may exceed `chi` by the gate's +operator-Schmidt rank; the latter is the enforced live-state cap. The extra +live maximum scans are opt-in so ordinary replay retains its default cost. + +For deterministic small-system fidelity checks, use `norm()` for the +canonical local norm and compare `to_dense()` with an independently replayed +NumPy statevector using a fixed gate stream. This avoids making a restricted +Cotengra overlap path the correctness oracle. A Tree/MPS overlap remains a +useful comparative accuracy diagnostic, but it includes both layouts' +different truncation histories. + ## Range / subtree canonicalisation The single orthogonality centre generalises to a connected **canonical region** @@ -343,8 +365,9 @@ The default replay configuration is intended for production evolution: `mode="auto"` uses the direct routed two-site kernel, `threads=1` avoids oversubscribing the small tree contractions, `subtree_workers=1` keeps the serial path allocation-free, `profile=False` avoids timing overhead, and -`track_truncation=False` avoids full-spectrum diagnostic SVDs. `record_history` -and `track_infidelity` retain the established API defaults; the latter only +`track_truncation=False` avoids full-spectrum diagnostic SVDs, while +`track_bond_diagnostics=False` avoids live-bond scans. `record_history` and +`track_infidelity` retain the established API defaults; the latter only adds norm-based progress readouts when a progress bar is requested. Warnings are reserved for an actionable behavior change: enabling diff --git a/src/pepsy/optimizers/tree/optimizer.py b/src/pepsy/optimizers/tree/optimizer.py index f59deae..3c0c772 100644 --- a/src/pepsy/optimizers/tree/optimizer.py +++ b/src/pepsy/optimizers/tree/optimizer.py @@ -416,6 +416,12 @@ class TreeOptimizer: Whether to collect opt-in kernel timing records in :meth:`profile_report`. Profiling is disabled by default and adds no synchronization or timing calls to the normal replay path. + track_bond_diagnostics : bool + Whether to record live and transient bond dimensions for each update. + This is disabled by default because determining the live maximum scans + the tree after QR hops. When enabled, :meth:`bond_diagnostic_report` + distinguishes temporary QR/gate growth from the post-compression live + state bonds. tn : TreeTensorNetwork or product MatrixProductState, optional Initial coefficient state. A tree state is copied and canonicalised if needed. Its plan must match ``tree``/``layout`` when it is entangled; @@ -479,7 +485,8 @@ def __init__(self, gates=None, n=None, *, chi=64, max_intermediate_bond=None, max_operator_qubits=_DEFAULT_MAX_OPERATOR_QUBITS, max_subtree_nodes=_DEFAULT_MAX_SUBTREE_NODES, - record_history=True, profile=False): + record_history=True, profile=False, + track_bond_diagnostics=False): # Preserve one-shot streams for both queue normalization and automatic # layout discovery. Materializing only inside # ``_normalize_gate_queue`` would leave the finder with an exhausted @@ -680,10 +687,12 @@ def __init__(self, gates=None, n=None, *, chi=64, ) self.record_history = bool(record_history) self.profile = bool(profile) + self.track_bond_diagnostics = bool(track_bond_diagnostics) self.profile_events = [] self.measurements = [] self.truncation_history = [] self.update_history = [] + self.bond_history = [] # Keep the same readout attributes as MpsOptimizer. Tree infidelity # samples are populated when ``track_truncation`` is enabled; without # the spectrum probes the only honest trace is the initial zero. @@ -1307,6 +1316,7 @@ def set_tn(self, tn): self._install_tn(tn) self.truncation_history.clear() self.update_history.clear() + self.bond_history.clear() self.infidelities[:] = [0.0] self.infidelity_samples.clear() self.normalizations.clear() @@ -2167,14 +2177,38 @@ def _begin_update(self, kind, where): """Start aggregating edge truncations for one state update.""" if self._active_update is not None: return False + live_before = ( + int(self.tn.max_bond()) + if self.track_bond_diagnostics else None + ) self._active_update = { "kind": str(kind), "support": tuple(int(q) for q in where), "edge_start": len(self.truncation_history), "started_at": time.perf_counter(), + "live_max_bond_before": live_before, + "transient_max_bond": live_before, + "bond_trace": [], } return True + def _record_transient_bond(self, dimension, *, phase, edge=None): + """Record one potentially transient bond dimension for the update.""" + active = self._active_update + if active is None or not self.track_bond_diagnostics: + return + dimension = int(dimension) + previous = active["transient_max_bond"] + active["transient_max_bond"] = ( + dimension if previous is None else max(previous, dimension) + ) + if self.track_bond_diagnostics: + active["bond_trace"].append({ + "phase": str(phase), + "edge": None if edge is None else tuple(int(x) for x in edge), + "bond": dimension, + }) + def _abort_update(self): """Discard a partial aggregation after a failed state update.""" if self._active_update is None: @@ -2189,12 +2223,39 @@ def _finish_update(self): if active is None: return elapsed = time.perf_counter() - active["started_at"] + live_after = ( + int(self.tn.max_bond()) + if self.track_bond_diagnostics else None + ) + transient_max = active.get("transient_max_bond") + if transient_max is None: + transient_max = live_after + transient_over_chi = ( + None + if self.chi is None or transient_max is None + else bool(transient_max > self.chi) + ) + bond_record = { + "update": len(self.update_history), + "kind": active["kind"], + "support": active["support"], + "live_max_bond_before": active.get("live_max_bond_before"), + "transient_max_bond": transient_max, + "live_max_bond_after": live_after, + "transient_exceeds_chi": transient_over_chi, + "bond_trace": deepcopy(active.get("bond_trace", [])), + } + if self.track_bond_diagnostics: + self.bond_history.append(deepcopy(bond_record)) if not self.record_history: if self.profile: self.profile_events.append({ "kind": "update", "support": active["support"], "seconds": elapsed, + "live_max_bond_before": active.get("live_max_bond_before"), + "transient_max_bond": transient_max, + "live_max_bond_after": live_after, }) self._active_update = None return @@ -2264,6 +2325,7 @@ def _finish_update(self): "cumulative_relative_discarded_weight": cumulative_loss, "max_edge_discarded_weight": max_edge_loss, "max_edge_discarded_fraction": max_edge_fraction, + **bond_record, }) if self.profile: self.profile_events.append({ @@ -3147,8 +3209,38 @@ def _thread_hop(self, u, v): above its pre-gate size, and that growth is undone by :meth:`_compress_path`. """ - if self.tn.fermionic: - return self._fermionic_thread_hop(u, v) + profile_started = time.perf_counter() if self.profile else None + before_bond = self.tn.bond(u, v) + before_dim = int(self.tn.ind_size(before_bond)) + try: + if self.tn.fermionic: + result = self._fermionic_thread_hop(u, v) + else: + result = self._dense_thread_hop(u, v) + finally: + after_bond = self.tn.bond(u, v) + after_dim = int(self.tn.ind_size(after_bond)) + self._record_transient_bond( + after_dim, phase="thread_hop", edge=(u, v), + ) + if self.profile: + try: + thread_dim = int(self.tn.ind_size(self._thread_ind)) + except (KeyError, TypeError, ValueError): + thread_dim = None + self.profile_events.append({ + "kind": "thread_hop", + "edge": (u, v), + "before_bond": before_dim, + "after_bond": after_dim, + "thread_bond": thread_dim, + "native": bool(self.tn.fermionic), + "seconds": time.perf_counter() - profile_started, + }) + return result + + def _dense_thread_hop(self, u, v): + """Perform one dense QR thread hop without timing/diagnostic wrapping.""" tu = self.tn.tensor_map[self._tid(u)] tv = self.tn.tensor_map[self._tid(v)] @@ -3317,6 +3409,9 @@ def _record_truncation( full_spectrum=None, max_bond=None, cutoff=None, ): """Record one edge split/compression and optional discarded weight.""" + self._record_transient_bond( + before_bond, phase=str(kind), edge=edge, + ) if not self.record_history: return discarded_weight = None @@ -5371,6 +5466,7 @@ def copy(self): max_subtree_nodes=self.max_subtree_nodes, record_history=self.record_history, profile=self.profile, + track_bond_diagnostics=self.track_bond_diagnostics, seed=child_seed, run=False, tn=self.tn, @@ -5382,6 +5478,7 @@ def copy(self): other.measurements = deepcopy(self.measurements) other.truncation_history = deepcopy(self.truncation_history) other.update_history = deepcopy(self.update_history) + other.bond_history = deepcopy(self.bond_history) other.infidelities = list(self.infidelities) other.infidelity_samples = deepcopy(self.infidelity_samples) other.normalizations = deepcopy(self.normalizations) @@ -5411,14 +5508,54 @@ def get_infidelity_samples(self): """Return detailed cumulative tree-truncation sample records.""" return self.infidelity_samples + def bond_diagnostic_report(self): + """Return live-versus-transient bond diagnostics collected per update. + + The transient maximum includes the bond dimensions presented to the + final compression SVD and the dimensions observed after exact QR + thread hops. Consequently it can exceed ``chi`` even when every + ``live_max_bond_after`` is capped by ``chi``. The report is populated + only when ``track_bond_diagnostics=True``; otherwise the update list + remains available but its live/transient fields are ``None``. + """ + updates = deepcopy( + self.bond_history if self.track_bond_diagnostics else self.update_history + ) + measured = [ + update for update in updates + if update.get("transient_max_bond") is not None + ] + live_after = [ + update["live_max_bond_after"] + for update in measured + if update.get("live_max_bond_after") is not None + ] + transient = [ + update["transient_max_bond"] + for update in measured + if update.get("transient_max_bond") is not None + ] + return { + "enabled": self.track_bond_diagnostics, + "chi": self.chi, + "updates": updates, + "max_live_bond_after": max(live_after) if live_after else None, + "max_transient_bond": max(transient) if transient else None, + "n_transient_exceeds_chi": sum( + bool(update.get("transient_exceeds_chi")) + for update in measured + ), + } + def profile_report(self): """Return opt-in tree kernel timings grouped by operation kind. Construct the optimizer with ``profile=True`` to collect records. Timing is deliberately kept separate from truncation history so the - normal replay and diagnostic APIs remain unchanged. The returned - ``events`` list is a deep copy and can safely be serialized alongside - a benchmark result. + normal replay and diagnostic APIs remain unchanged. In addition to + update and compression events, two-site direct routing reports each + exact QR ``thread_hop`` separately. The returned ``events`` list is a + deep copy and can safely be serialized alongside a benchmark result. """ events = deepcopy(self.profile_events) grouped = {} diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index 768bad5..923fdda 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -3289,6 +3289,69 @@ def test_tree_profile_report_is_opt_in(): assert report["total_seconds"] > 0.0 +def test_tree_profile_separates_qr_thread_hops_from_compression(): + """Profiled direct routing exposes exact QR hops as separate events.""" + cnot = np.array( + [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]], + dtype=complex, + ) + opt = TreeOptimizer( + [(cnot, (0, 3))], n=4, chi=2, cutoff=0.0, + profile=True, track_bond_diagnostics=True, + ) + + report = opt.profile_report() + hops = [event for event in report["events"] if event["kind"] == "thread_hop"] + assert hops + assert all(event["seconds"] >= 0.0 for event in hops) + assert report["by_kind"]["thread_hop"]["count"] == len(hops) + assert report["by_kind"]["edge_canonize"]["count"] >= 1 + + +def test_tree_bond_diagnostics_distinguish_transient_qr_growth(): + """Temporary gate/QR growth may exceed chi while live bonds do not.""" + cnot = np.array( + [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]], + dtype=complex, + ) + opt = TreeOptimizer( + [(cnot, (0, 3))], n=4, chi=1, cutoff=0.0, + record_history=False, track_bond_diagnostics=True, + ) + + report = opt.bond_diagnostic_report() + assert report["enabled"] is True + assert report["max_transient_bond"] >= 2 + assert report["max_live_bond_after"] <= 1 + assert report["n_transient_exceeds_chi"] >= 1 + update = report["updates"][0] + assert update["transient_max_bond"] > update["live_max_bond_after"] + assert update["transient_exceeds_chi"] is True + assert update["bond_trace"] + + +def test_tree_norm_and_fidelity_check_is_deterministic_without_network_fidelity(): + """Small exact replay uses local norm plus a deterministic statevector oracle.""" + rng = np.random.default_rng(417) + stream = _random_stream(4, 10, rng, two_qubit_frac=0.8) + exact = _exact_state(stream, 4) + + first = TreeOptimizer( + stream, n=4, chi=64, cutoff=0.0, threads=1, + ) + second = TreeOptimizer( + stream, n=4, chi=64, cutoff=0.0, threads=2, + ) + first_dense = first.to_dense() + second_dense = second.to_dense() + + assert first.norm() == pytest.approx(np.linalg.norm(first_dense)) + assert second.norm() == pytest.approx(np.linalg.norm(second_dense)) + assert _fidelity(exact, first_dense) > 1.0 - 1e-10 + assert _fidelity(exact, second_dense) > 1.0 - 1e-10 + assert _fidelity(first_dense, second_dense) > 1.0 - 1e-12 + + def test_tree_pauli_expectation_and_projection_are_public(): """Pauli expectation/projection share the measurement backend semantics.""" h = np.array([[1.0, 1.0], [1.0, -1.0]], dtype=complex) / np.sqrt(2.0) From a0d9c510bcce7881063d0a0201219ab19e5927f8 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Mon, 3 Aug 2026 11:29:01 -0600 Subject: [PATCH 60/70] Add native compression route diagnostics --- docs/api/optimizers/tree.md | 5 +++ src/pepsy/optimizers/tree/optimizer.py | 11 +++++- src/pepsy/optimizers/tree/ttn.py | 53 ++++++++++++++++++++++++++ tests/test_optimize_tree.py | 44 +++++++++++++++++++++ 4 files changed, 111 insertions(+), 2 deletions(-) diff --git a/docs/api/optimizers/tree.md b/docs/api/optimizers/tree.md index cb227f6..7b741f7 100644 --- a/docs/api/optimizers/tree.md +++ b/docs/api/optimizers/tree.md @@ -242,6 +242,11 @@ timed `thread_hop`, `edge_canonize`, and `edge_compress` events. The `thread_hop` events are the exact, lossless QR carry moves; `edge_compress` events are the truncating SVD work. These timings are nested inside the per-update envelope and should not be added as independent wall-clock totals. +For native Symmray compression, `profile_report()` also returns +`native_compression_routes`: counts of `one_sided_left`, `one_sided_right`, and +`two_sided_reduced` show that the graded reduced-core paths are active, while +`full_svd_fallback` identifies a conservative complete two-node SVD. Route +records have zero duration; use the enclosing `edge_compress` event for timing. For a dimension-level report, also pass `track_bond_diagnostics=True`. `bond_diagnostic_report()` then records diff --git a/src/pepsy/optimizers/tree/optimizer.py b/src/pepsy/optimizers/tree/optimizer.py index 3c0c772..5c954cb 100644 --- a/src/pepsy/optimizers/tree/optimizer.py +++ b/src/pepsy/optimizers/tree/optimizer.py @@ -5554,11 +5554,14 @@ def profile_report(self): Timing is deliberately kept separate from truncation history so the normal replay and diagnostic APIs remain unchanged. In addition to update and compression events, two-site direct routing reports each - exact QR ``thread_hop`` separately. The returned ``events`` list is a - deep copy and can safely be serialized alongside a benchmark result. + exact QR ``thread_hop`` separately. Native compression events also + identify whether the reduced graded QR/SVD route or the conservative + full two-node SVD was selected. The returned ``events`` list is a deep + copy and can safely be serialized alongside a benchmark result. """ events = deepcopy(self.profile_events) grouped = {} + native_routes = {} for event in events: kind = str(event.get("kind", "unknown")) summary = grouped.setdefault( @@ -5566,10 +5569,14 @@ def profile_report(self): ) summary["count"] += 1 summary["seconds"] += float(event.get("seconds", 0.0)) + if kind == "native_compression_route": + route = str(event.get("route", "unknown")) + native_routes[route] = native_routes.get(route, 0) + 1 return { "enabled": self.profile, "events": events, "by_kind": grouped, + "native_compression_routes": native_routes, "total_seconds": float( sum(float(event.get("seconds", 0.0)) for event in events) ), diff --git a/src/pepsy/optimizers/tree/ttn.py b/src/pepsy/optimizers/tree/ttn.py index d4e327b..1016cf5 100644 --- a/src/pepsy/optimizers/tree/ttn.py +++ b/src/pepsy/optimizers/tree/ttn.py @@ -1281,6 +1281,30 @@ def _native_qr_split(self, tensor, **kwargs): kwargs.setdefault("method", "qr") return tensor.split(**kwargs) + def _record_native_compression_route( + self, route, *, edge, before_bond, reduction_hint, reduction_proven, + ): + """Record which native compression decomposition was used. + + Route records are attached to the existing opt-in profile sink, so + ordinary replay does not allocate diagnostic dictionaries or perform + any timing work. The records deliberately describe decomposition + selection rather than timing: the surrounding ``edge_compress`` event + remains the authoritative duration measurement. + """ + profile_sink = getattr(self, "_profile_sink", None) + if profile_sink is None: + return + profile_sink.append({ + "kind": "native_compression_route", + "route": route, + "edge": tuple(edge), + "before_bond": int(before_bond), + "reduced": reduction_hint, + "reduction_proven": bool(reduction_proven), + "seconds": 0.0, + }) + # -- edge-level canonical / compression helpers --------------------------- def _fermionic_canonize_edge_(self, a, b, absorb): @@ -1342,6 +1366,7 @@ def _fermionic_compress_edge_( isometric = self.node_tensor(isometric_node) reduced_tensor = self.node_tensor(reduced_node) bond = self.bond(isometric_node, reduced_node) + before_bond = int(self.ind_size(bond)) left_inds = [index for index in isometric.inds if index != bond] # ``reduced="left"`` is the metadata value emitted by the optimizer @@ -1400,6 +1425,13 @@ def _fermionic_compress_edge_( inds=merged.inds, left_inds=None, ) + self._record_native_compression_route( + "one_sided_left", + edge=(a, b), + before_bond=before_bond, + reduction_hint=reduction_hint, + reduction_proven=reduction_proven, + ) return self if ( @@ -1443,6 +1475,13 @@ def _fermionic_compress_edge_( inds=core_right.inds, left_inds=None, ) + self._record_native_compression_route( + "one_sided_right", + edge=(a, b), + before_bond=before_bond, + reduction_hint=reduction_hint, + reduction_proven=reduction_proven, + ) return self if reduction_hint is True: @@ -1500,6 +1539,13 @@ def _fermionic_compress_edge_( inds=reduced_compressed.inds, left_inds=None, ) + self._record_native_compression_route( + "two_sided_reduced", + edge=(a, b), + before_bond=before_bond, + reduction_hint=reduction_hint, + reduction_proven=reduction_proven, + ) return self # Keep the old complete graded split as a compatibility fallback for @@ -1525,6 +1571,13 @@ def _fermionic_compress_edge_( inds=remainder.inds, left_inds=None, ) + self._record_native_compression_route( + "full_svd_fallback", + edge=(a, b), + before_bond=before_bond, + reduction_hint=reduction_hint, + reduction_proven=reduction_proven, + ) return self def _track_edge_center(self, a, b, absorb, *, previous=None): diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index 923fdda..97b5521 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -3276,6 +3276,7 @@ def test_tree_profile_report_is_opt_in(): "enabled": False, "events": [], "by_kind": {}, + "native_compression_routes": {}, "total_seconds": 0.0, } @@ -3286,6 +3287,7 @@ def test_tree_profile_report_is_opt_in(): assert report["enabled"] is True assert report["events"] assert report["by_kind"]["update"]["count"] == 2 + assert report["native_compression_routes"] == {} assert report["total_seconds"] > 0.0 @@ -5302,6 +5304,48 @@ def traced_split(self, *args, **kwargs): assert ttn.validate(check_canonical=True) is ttn +def test_native_profile_reports_reduced_compression_routes(): + """Native profiling exposes reduced compression and no hidden fallback.""" + pytest.importorskip("symmray") + fermion = pepsy.Fermion( + spinful=True, + symmetry="U1U1", + dtype="complex128", + ) + plan = TreePlan.from_order(range(4), structure="balanced") + state = pepsy.ps_to_ttn( + 4, + tree=plan, + fermion=fermion, + occupations=((1, 0), (0, 1), (1, 0), (0, 1)), + dtype="complex128", + ) + hopping = fermion.hopping_gate(0.05, t=1.0, imaginary=False) + optimizer = TreeOptimizer( + None, + n=4, + tree=plan, + state=state, + chi=1, + cutoff=0.0, + mode="direct", + profile=True, + run=False, + ) + + optimizer.apply_2q(hopping, 0, 2) + report = optimizer.profile_report() + routes = report["native_compression_routes"] + + assert routes + assert routes.get("full_svd_fallback", 0) == 0 + assert sum( + count for route, count in routes.items() + if route != "full_svd_fallback" + ) == report["by_kind"]["native_compression_route"]["count"] + assert optimizer.tn.validate(check_canonical=True) is optimizer.tn + + @pytest.mark.parametrize( ("symmetry", "spinful", "occupations"), [ From 031391c798b32e6acdc240dd5037214227863ba4 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Mon, 3 Aug 2026 14:47:55 -0600 Subject: [PATCH 61/70] Optimize Tree native update paths --- docs/api/optimizers/tree.md | 29 +- src/pepsy/optimizers/tree/optimizer.py | 375 +++++++++++++++++++++---- src/pepsy/optimizers/tree/ttn.py | 116 +++++++- tests/test_optimize_tree.py | 62 +++- 4 files changed, 496 insertions(+), 86 deletions(-) diff --git a/docs/api/optimizers/tree.md b/docs/api/optimizers/tree.md index 7b741f7..26de114 100644 --- a/docs/api/optimizers/tree.md +++ b/docs/api/optimizers/tree.md @@ -197,11 +197,14 @@ compression events totaling 2.40 s inside 48 update envelopes totaling 6.18 s; it identified the gate update/threading/contraction path, especially the central edges, rather than route-length tuning, as the next target. -The update path now avoids repeating the same native isometry proof: each -truncating edge validates its `left_inds`/charge-map proof once, while a -lossless QR edge skips the reduction lookup entirely. Two-site factors use -private stable output labels, removing per-factor UUID allocation and one -metadata reindex. In multi-site Tree/MPO updates, independent QR messages +The update path carries the isometry proof produced by the lossless threading +sweep into the reverse compression sweep, so native edges do not revalidate +the same `left_inds`/charge-map proof at every central edge. Two-site factors +use state-owned unique work labels rather than per-factor UUID allocation; +live routed bonds remain collision-safe across copied states without random +label setup in the hot loop. Native Torch-CPU one-edge contractions use +Symmray's blockwise mode, while CUDA and other backends retain the fused mode. +In multi-site Tree/MPO updates, independent QR messages landing at the same node are contracted as one batch; dense waves reuse their worker pool, while native fermionic routing remains serial for Symmray safety. These changes preserve the complete-gate-before-truncation rule and the @@ -238,10 +241,17 @@ when using this number as an accuracy diagnostic. #### QR/hop and bond-growth diagnostics Construct `TreeOptimizer(..., profile=True)` to split the update envelope into -timed `thread_hop`, `edge_canonize`, and `edge_compress` events. The +timed `thread_hop`, `edge_canonize`, and `edge_compress` events. The profile +also records `gate_factorization`, `tensor_absorption`, `center_movement`, +`metadata_path`, and `subtree_hub_merge` phases when those routes are used. The `thread_hop` events are the exact, lossless QR carry moves; `edge_compress` events are the truncating SVD work. These timings are nested inside the -per-update envelope and should not be added as independent wall-clock totals. +per-update envelope and should not be added as independent wall-clock totals; +use `profile_report()["update_seconds"]` as the envelope total. The +`timing_semantics` field records this relationship explicitly. For asynchronous +CuPy or CUDA work, `profile_sync=True` synchronizes the active device at each +phase boundary so phase durations represent device execution; this is a +diagnostic mode and adds synchronization overhead. For native Symmray compression, `profile_report()` also returns `native_compression_routes`: counts of `one_sided_left`, `one_sided_right`, and `two_sided_reduced` show that the graded reduced-core paths are active, while @@ -250,8 +260,9 @@ records have zero duration; use the enclosing `edge_compress` event for timing. For a dimension-level report, also pass `track_bond_diagnostics=True`. `bond_diagnostic_report()` then records -`transient_max_bond` during routing/factorization and `live_max_bond_after` -after the compression sweep. The former may exceed `chi` by the gate's +the per-update `live_max_bond_before`, `transient_max_bond` during +routing/factorization, and `live_max_bond_after` after the compression sweep. +The former may exceed `chi` by the gate's operator-Schmidt rank; the latter is the enforced live-state cap. The extra live maximum scans are opt-in so ordinary replay retains its default cost. diff --git a/src/pepsy/optimizers/tree/optimizer.py b/src/pepsy/optimizers/tree/optimizer.py index 5c954cb..2cf22e4 100644 --- a/src/pepsy/optimizers/tree/optimizer.py +++ b/src/pepsy/optimizers/tree/optimizer.py @@ -416,6 +416,11 @@ class TreeOptimizer: Whether to collect opt-in kernel timing records in :meth:`profile_report`. Profiling is disabled by default and adds no synchronization or timing calls to the normal replay path. + profile_sync : bool + Whether profiled phase boundaries should synchronize the active device + before taking timestamps. This is useful for asynchronous CuPy and + CUDA backends, but adds a synchronization at every recorded phase and + is therefore disabled by default. track_bond_diagnostics : bool Whether to record live and transient bond dimensions for each update. This is disabled by default because determining the live maximum scans @@ -485,7 +490,7 @@ def __init__(self, gates=None, n=None, *, chi=64, max_intermediate_bond=None, max_operator_qubits=_DEFAULT_MAX_OPERATOR_QUBITS, max_subtree_nodes=_DEFAULT_MAX_SUBTREE_NODES, - record_history=True, profile=False, + record_history=True, profile=False, profile_sync=False, track_bond_diagnostics=False): # Preserve one-shot streams for both queue normalization and automatic # layout discovery. Materializing only inside @@ -687,6 +692,7 @@ def __init__(self, gates=None, n=None, *, chi=64, ) self.record_history = bool(record_history) self.profile = bool(profile) + self.profile_sync = bool(profile_sync) self.track_bond_diagnostics = bool(track_bond_diagnostics) self.profile_events = [] self.measurements = [] @@ -714,6 +720,7 @@ def __init__(self, gates=None, n=None, *, chi=64, self._two_site_path_cache = {} self._two_site_path_cache_limit = 256 self._active_update = None + self._update_counter = 0 self._truncation_log_survival = 0.0 if tree is None: @@ -1309,6 +1316,42 @@ def _attach_profile_sink(self): """Attach the optimizer's optional timing sink to the live TTN.""" self.tn._profile_sink = self.profile_events if self.profile else None + def _profile_synchronize(self): + """Synchronize an asynchronous backend for opt-in phase timing.""" + if not self.profile_sync: + return + backend = getattr(self, "backend", None) + array_backend = getattr(self, "array_backend", backend) + if array_backend == "cupy" or backend == "cupy": + import cupy as cp # pylint: disable=import-outside-toplevel + + cp.cuda.runtime.deviceSynchronize() + elif backend == "torch": + data = self._state_like() + if bool(getattr(data, "is_cuda", False)): + import torch # pylint: disable=import-outside-toplevel + + torch.cuda.synchronize(data.device) + + def _profile_phase_start(self): + """Return a phase timestamp, or ``None`` when profiling is disabled.""" + if not self.profile: + return None + self._profile_synchronize() + return time.perf_counter() + + def _profile_phase_event(self, kind, started, **payload): + """Append one opt-in phase event with optional device synchronization.""" + if started is None: + return + self._profile_synchronize() + event = { + "kind": str(kind), + **payload, + "seconds": time.perf_counter() - started, + } + self.profile_events.append(event) + def set_tn(self, tn): """Replace the live tree state with a canonical independent copy.""" if not isinstance(tn, TreeTensorNetwork): @@ -1322,6 +1365,7 @@ def set_tn(self, tn): self.normalizations.clear() self.projection_diagnostics.clear() self._truncation_log_survival = 0.0 + self._update_counter = 0 self._attach_profile_sink() return self @@ -1467,7 +1511,17 @@ def _move_center(self, target): a path walk when a multi-node canonical region is known. Only an otherwise uncatalogued state requires a full O(N) canonicalisation. """ - self.tn.shift_orthogonality_center(target) + started = self._profile_phase_start() + previous = self.center + try: + return self.tn.shift_orthogonality_center(target) + finally: + self._profile_phase_event( + "center_movement", + started, + source=previous, + target=target, + ) def _nearest_anchor(self, nodes): """Choose the closest node to the current centre or canonical region. @@ -2181,9 +2235,12 @@ def _begin_update(self, kind, where): int(self.tn.max_bond()) if self.track_bond_diagnostics else None ) + update_index = self._update_counter + self._update_counter += 1 self._active_update = { "kind": str(kind), "support": tuple(int(q) for q in where), + "update": update_index, "edge_start": len(self.truncation_history), "started_at": time.perf_counter(), "live_max_bond_before": live_before, @@ -2230,13 +2287,19 @@ def _finish_update(self): transient_max = active.get("transient_max_bond") if transient_max is None: transient_max = live_after + update_index = active.get( + "update", + len(self.bond_history) + if self.track_bond_diagnostics + else len(self.update_history), + ) transient_over_chi = ( None if self.chi is None or transient_max is None else bool(transient_max > self.chi) ) bond_record = { - "update": len(self.update_history), + "update": update_index, "kind": active["kind"], "support": active["support"], "live_max_bond_before": active.get("live_max_bond_before"), @@ -2251,6 +2314,7 @@ def _finish_update(self): if self.profile: self.profile_events.append({ "kind": "update", + "update": update_index, "support": active["support"], "seconds": elapsed, "live_max_bond_before": active.get("live_max_bond_before"), @@ -2313,7 +2377,7 @@ def _finish_update(self): cumulative_loss = None self.update_history.append({ - "update": len(self.update_history), + "update": update_index, "kind": active["kind"], "support": active["support"], "elapsed_seconds": float(elapsed), @@ -2330,6 +2394,7 @@ def _finish_update(self): if self.profile: self.profile_events.append({ "kind": "update", + "update": update_index, "support": active["support"], "seconds": elapsed, }) @@ -2731,7 +2796,16 @@ def _apply_1q_impl(self, gate, q, *, renormalize=False): self._move_center(site_node) region = self.tn.canonical_region left_inds = self.tn.node_tensor(site_node).left_inds - self.tn.gate_inds_(gate, [self._phys(q)], contract=True) + absorb_started = self._profile_phase_start() + try: + self.tn.gate_inds_(gate, [self._phys(q)], contract=True) + finally: + self._profile_phase_event( + "tensor_absorption", + absorb_started, + support=(q,), + route="one_site", + ) if unitary: # A physical unitary preserves the isometric exterior, but the # state-owned gate mutator deliberately invalidates metadata for @@ -2833,6 +2907,7 @@ def _apply_2q_mpo_impl( db = int(self.tn.ind_size(self._phys(qb))) cache_source = gate gate = self._as_gate_tensor4(gate, da, db) + factor_started = self._profile_phase_start() # ``MatrixProductOperator.from_dense`` is backend-generic: for a # Symmray FermionicArray its block-aware SVD returns native fermionic @@ -2852,6 +2927,7 @@ def _apply_2q_mpo_impl( self.n, ) cached = self._gate_factor_cache.get(cache_key) + cache_hit = cached is not None and cached[0] is cache_source if cached is not None and cached[0] is cache_source: submpo = cached[1] else: @@ -2866,6 +2942,14 @@ def _apply_2q_mpo_impl( "two-site gate could not be represented as a Quimb sub-MPO " "with one local factor per requested site." ) + self._profile_phase_event( + "gate_factorization", + factor_started, + route="mpo", + cache_hit=cache_hit, + support=(qa, qb), + input_shape=(da, db, da, db), + ) return self._apply_2q_factors_impl( *factors, qa, @@ -2918,6 +3002,12 @@ def _two_site_mpo_factors(self, submpo, qa, qb, *, site_where=None): if len(tids) != 1: return None factor = tensor_map[tids[0]].copy() + # The two-site fast path consumes these private factors + # directly, before the structured sub-MPO route gets a + # chance to coerce its payload tensors. Keep the caller's + # MPO untouched while matching each factor to the live TTN + # backend, dtype, and device. + factor.modify(data=self._as_state_backend(factor.data)) upper = upper_id.format(site) lower = lower_id.format(site) if upper not in factor.inds or lower not in factor.inds: @@ -2953,7 +3043,7 @@ def _two_site_mpo_factors(self, submpo, qa, qb, *, site_where=None): cache_key, submpo, (raw_factors, shared_bond) ) - thread_ind = qtn.rand_uuid() + thread_ind = self.tn._new_work_bond("mpo_thread", qa, qb) factors = {} outputs = {} for qubit, (factor_template, upper, lower) in raw_factors.items(): @@ -3004,6 +3094,7 @@ def _cached_direct_gate_factors( self, gate, source, qa, qb, pa, pb, da, db, ): """Return fresh-index copies of a cached direct gate factorization.""" + factor_started = self._profile_phase_start() key = ( "direct", id(source), @@ -3015,6 +3106,7 @@ def _cached_direct_gate_factors( db, ) cached = self._gate_factor_cache.get(key) + cache_hit = cached is not None and cached[0] is source if cached is not None and cached[0] is source: left_template, right_template = cached[1] else: @@ -3044,7 +3136,16 @@ def _cached_direct_gate_factors( key, source, (left_template, right_template) ) - thread_ind = qtn.rand_uuid() + self._profile_phase_event( + "gate_factorization", + factor_started, + route="direct", + cache_hit=cache_hit, + support=(qa, qb), + input_shape=(da, db, da, db), + ) + + thread_ind = self.tn._new_work_bond("gate_thread", qa, qb) left = left_template.copy() right = right_template.copy() output_a = "_pepsy_gate_out_a" @@ -3076,6 +3177,7 @@ def _apply_2q_factors_impl( one canonical compression sweep. """ plan = self.plan + path_started = self._profile_phase_start() la, lb, path = self._cached_two_site_path(qa, qb) parent = plan.parent.get(la) if ( @@ -3084,6 +3186,13 @@ def _apply_2q_factors_impl( and parent is not None and plan.parent.get(lb) == parent ): + self._profile_phase_event( + "metadata_path", + path_started, + support=(qa, qb), + route="sibling", + path_length=len(path), + ) return self._apply_2q_sibling_factors( factors, outputs, qa, qb, la, lb, parent, max_bond=max_bond, @@ -3100,15 +3209,33 @@ def _apply_2q_factors_impl( destination_node = plan.node_of_qubit[destination] if source_node != path[0]: path = path[::-1] + self._profile_phase_event( + "metadata_path", + path_started, + support=(qa, qb), + route="threaded", + path_length=len(path), + source=source_node, + destination=destination_node, + ) self._move_center(source_node) self._thread_ind = thread_ind try: source_tensor = self.tn.tensor_map[self._tid(source_node)] - merged_source = _contract_two_tensors( - source_tensor, - factors[source], - shared_ind=self._phys(source), - ).reindex_({outputs[source]: self._phys(source)}) + absorb_started = self._profile_phase_start() + try: + merged_source = _contract_two_tensors( + source_tensor, + factors[source], + shared_ind=self._phys(source), + ).reindex_({outputs[source]: self._phys(source)}) + finally: + self._profile_phase_event( + "tensor_absorption", + absorb_started, + support=(source,), + route="threaded_source", + ) source_tensor.modify( data=merged_source.data, inds=merged_source.inds, ) @@ -3117,11 +3244,20 @@ def _apply_2q_factors_impl( self._thread_hop(u, v) destination_tensor = self.tn.tensor_map[self._tid(destination_node)] - merged_destination = _contract_two_tensors( - factors[destination], - destination_tensor, - shared_ind=self._phys(destination), - ).reindex_({outputs[destination]: self._phys(destination)}) + absorb_started = self._profile_phase_start() + try: + merged_destination = _contract_two_tensors( + factors[destination], + destination_tensor, + shared_ind=self._phys(destination), + ).reindex_({outputs[destination]: self._phys(destination)}) + finally: + self._profile_phase_event( + "tensor_absorption", + absorb_started, + support=(destination,), + route="threaded_destination", + ) destination_tensor.modify( data=merged_destination.data, inds=merged_destination.inds, ) @@ -3156,17 +3292,26 @@ def _apply_2q_sibling_factors( e_la = self._bond_name(la, parent) e_lb = self._bond_name(lb, parent) - merged_a = _contract_two_tensors( - tla, factors[qa], shared_ind=pa, - ).reindex_( - {outputs[qa]: pa} - ) - merged_b = _contract_two_tensors( - tlb, factors[qb], shared_ind=pb, - ).reindex_( - {outputs[qb]: pb} - ) - blob = qtn.tensor_contract(merged_a, tp, merged_b) + absorb_started = self._profile_phase_start() + try: + merged_a = _contract_two_tensors( + tla, factors[qa], shared_ind=pa, + ).reindex_( + {outputs[qa]: pa} + ) + merged_b = _contract_two_tensors( + tlb, factors[qb], shared_ind=pb, + ).reindex_( + {outputs[qb]: pb} + ) + blob = qtn.tensor_contract(merged_a, tp, merged_b) + finally: + self._profile_phase_event( + "tensor_absorption", + absorb_started, + support=(qa, qb), + route="sibling_blob", + ) # Split off leaf a (isometric), then leaf b, leaving the centre at the # parent; both new bonds keep their canonical tree-edge names. @@ -3230,6 +3375,10 @@ def _thread_hop(self, u, v): thread_dim = None self.profile_events.append({ "kind": "thread_hop", + "update": ( + None if self._active_update is None + else self._active_update.get("update") + ), "edge": (u, v), "before_bond": before_dim, "after_bond": after_dim, @@ -3289,6 +3438,7 @@ def _fermionic_thread_hop(self, u, v): absorb="right", cutoff=0.0, get="tensors", + bond_ind=self.tn._new_work_bond("thread_hop", u, v), ) merged_v = _contract_two_tensors(carry, tv, shared_ind=edge) tu.modify( @@ -3557,6 +3707,10 @@ def _compress_edge_with_diagnostics( if profile_started is not None: self.profile_events.append({ "kind": "edge_canonize", + "update": ( + None if self._active_update is None + else self._active_update.get("update") + ), "edge": (u, v), "before_bond": before_bond, "after_bond": int(self.tn.ind_size(bond_after)), @@ -3593,6 +3747,10 @@ def _compress_edge_with_diagnostics( if profile_started is not None: self.profile_events.append({ "kind": "edge_compress", + "update": ( + None if self._active_update is None + else self._active_update.get("update") + ), "edge": (u, v), "before_bond": before_bond, "after_bond": after_bond, @@ -3651,6 +3809,12 @@ def _compress_path( leaves the centre at ``path[0]``. This is the re-orthonormalisation sweep of Seitz et al. (Fig. 6) applied along the gate geodesic. """ + # Every node before the destination was produced by the lossless QR + # threading sweep. Its ``left_inds`` therefore prove that it is + # isometric toward the destination side of the next compression edge. + # Carry this proof through the reverse sweep instead of asking every + # native edge to revalidate the same charge maps. The proof is local + # to this update and is not used by public arbitrary edge callers. for v, u in zip(path[::-1], path[-2::-1]): edge_cutoff = cutoff if preserve_subcap: @@ -3659,9 +3823,7 @@ def _compress_path( max_bond=max_bond, cutoff=cutoff, ) - reduced, reduction_proven = self._edge_reduction( - v, u, max_bond=max_bond, cutoff=edge_cutoff, - ) + reduced, reduction_proven = "left", True self._compress_edge_with_diagnostics( v, u, max_bond=max_bond, cutoff=edge_cutoff, reduced=reduced, reduction_proven=reduction_proven, @@ -3717,7 +3879,22 @@ def descend(node, parent): reduction_proven=reduction_proven, ) descend(child, node) - self.tn.canonize_edge_(child, node, absorb="right") + canonize_bond = int( + self.tn.ind_size(self.tn.bond(child, node)) + ) + canonize_started = self._profile_phase_start() + try: + self.tn.canonize_edge_(child, node, absorb="right") + finally: + self._profile_phase_event( + "edge_canonize", + canonize_started, + edge=(child, node), + before_bond=canonize_bond, + after_bond=int( + self.tn.ind_size(self.tn.bond(child, node)) + ), + ) descend(hub, None) self.center = hub @@ -3804,13 +3981,22 @@ def split_message(item): for destination, destination_results in by_destination.items(): messages = [result[-1] for result in destination_results] - if len(messages) == 1: - local[destination] = qtn.tensor_contract( - local[destination], messages[0] - ) - else: - local[destination] = qtn.tensor_contract( - local[destination], *messages + merge_started = self._profile_phase_start() + try: + if len(messages) == 1: + local[destination] = qtn.tensor_contract( + local[destination], messages[0] + ) + else: + local[destination] = qtn.tensor_contract( + local[destination], *messages + ) + finally: + self._profile_phase_event( + "subtree_hub_merge", + merge_started, + destination=destination, + message_count=len(messages), ) for ( _, source, _, state_bond, new_bond, kept, _ @@ -3985,10 +4171,20 @@ def _apply_submpo_resolved(self, submpo, where, *, max_bond=None, with self._thread_ctx(): applied = None if len(where) == 2: - factors = self._two_site_mpo_factors( - submpo, where[0], where[1], - site_where=(logical_where[0], logical_where[1]), - ) + factor_started = self._profile_phase_start() + try: + factors = self._two_site_mpo_factors( + submpo, where[0], where[1], + site_where=(logical_where[0], logical_where[1]), + ) + finally: + self._profile_phase_event( + "gate_factorization", + factor_started, + route="submpo", + cache_hit=False, + support=tuple(logical_where), + ) if factors is not None: self._apply_2q_factors_impl( *factors, where[0], where[1], @@ -4256,9 +4452,17 @@ def _apply_subtree_operator_impl(self, op, where, *, max_bond=None, # with any other state tensor: each edge creates one local message, # which is immediately absorbed by its parent and split again. order, hub = self._peel_order(snodes) + factor_started = self._profile_phase_start() op_factors, op_bonds = self._decompose_tree_operator( op_arr, where, snodes, order, hub, ) + self._profile_phase_event( + "gate_factorization", + factor_started, + route="tree_operator", + support=tuple(logical_where), + subtree_nodes=len(snodes), + ) self._apply_factorized_subtree_operator_impl( op_factors, op_bonds, where, snodes, order, hub, max_bond=max_bond, cutoff=cutoff, @@ -4313,7 +4517,17 @@ def _try_apply_native_submpo( site_nodes = [self.plan.node_of_qubit[q] for q in where] snodes = self._steiner_nodes(site_nodes) self._move_center(self._nearest_anchor(site_nodes)) + path_started = self._profile_phase_start() order, hub = self._peel_order(snodes) + self._profile_phase_event( + "metadata_path", + path_started, + support=tuple(payload_where), + route="submpo_subtree", + subtree_nodes=len(snodes), + message_edges=len(order), + hub=hub, + ) local = {} state_inds = {} operator_inds = {} @@ -4345,11 +4559,20 @@ def _try_apply_native_submpo( operator_inds[nid] = set(op_t.inds) - { self._phys(q) + "*", self._phys(q) } - local[nid] = _contract_two_tensors( - state_t, op_t, shared_ind=self._phys(q), - ).reindex_( - {self._phys(q) + "*": self._phys(q)} - ) + absorb_started = self._profile_phase_start() + try: + local[nid] = _contract_two_tensors( + state_t, op_t, shared_ind=self._phys(q), + ).reindex_( + {self._phys(q) + "*": self._phys(q)} + ) + finally: + self._profile_phase_event( + "tensor_absorption", + absorb_started, + support=(q,), + route="submpo_site", + ) except (KeyError, TypeError, ValueError): return None @@ -4390,17 +4613,34 @@ def _apply_factorized_subtree_operator_impl( # Operator sites are packed into one dimension-four leg. Split # that leg only at physical sites, then contract its input leg # with the live state physical index. - op_t = self._expand_tree_operator_leaf( - op_t, - op_bonds["physical"][q], - self._phys(q), - ) - if q is not None and q in where: - local[nid] = _contract_two_tensors( - state_t, op_t, shared_ind=self._phys(q), - ) + absorb_started = self._profile_phase_start() + try: + op_t = self._expand_tree_operator_leaf( + op_t, + op_bonds["physical"][q], + self._phys(q), + ) + local[nid] = _contract_two_tensors( + state_t, op_t, shared_ind=self._phys(q), + ) + finally: + self._profile_phase_event( + "tensor_absorption", + absorb_started, + support=(q,), + route="tree_operator_site", + ) else: - local[nid] = qtn.tensor_contract(state_t, op_t) + absorb_started = self._profile_phase_start() + try: + local[nid] = qtn.tensor_contract(state_t, op_t) + finally: + self._profile_phase_event( + "tensor_absorption", + absorb_started, + support=(), + route="tree_operator_internal", + ) if q is not None and q in where: local[nid].reindex_({f"{self._phys(q)}*": self._phys(q)}) operator_inds[nid] = set(local[nid].inds) - state_inds[nid] @@ -5466,6 +5706,7 @@ def copy(self): max_subtree_nodes=self.max_subtree_nodes, record_history=self.record_history, profile=self.profile, + profile_sync=self.profile_sync, track_bond_diagnostics=self.track_bond_diagnostics, seed=child_seed, run=False, @@ -5489,6 +5730,7 @@ def copy(self): other._track_warning_emitted = self._track_warning_emitted other._logical_qubits = list(self._logical_qubits) other._logical_positions = dict(self._logical_positions) + other._update_counter = self._update_counter other._truncation_log_survival = self._truncation_log_survival other.profile_events = deepcopy(self.profile_events) other._attach_profile_sink() @@ -5572,11 +5814,28 @@ def profile_report(self): if kind == "native_compression_route": route = str(event.get("route", "unknown")) native_routes[route] = native_routes.get(route, 0) + 1 + update_seconds = float(grouped.get("update", {}).get("seconds", 0.0)) return { "enabled": self.profile, "events": events, "by_kind": grouped, "native_compression_routes": native_routes, + "update_seconds": update_seconds, + "timing_semantics": { + "wall_envelope": "update", + "nested_event_kinds": [ + "gate_factorization", + "center_movement", + "metadata_path", + "thread_hop", + "tensor_absorption", + "edge_canonize", + "edge_compress", + "subtree_hub_merge", + "native_compression_route", + ], + "total_seconds_is_sum_of_events_not_wall_time": True, + }, "total_seconds": float( sum(float(event.get("seconds", 0.0)) for event in events) ), diff --git a/src/pepsy/optimizers/tree/ttn.py b/src/pepsy/optimizers/tree/ttn.py index 1016cf5..65bc578 100644 --- a/src/pepsy/optimizers/tree/ttn.py +++ b/src/pepsy/optimizers/tree/ttn.py @@ -67,25 +67,73 @@ def _native_qr_block_scaled(array, **kwargs): opts = dict(kwargs) opts.pop("method", None) opts.pop("fn", None) + + def torch_qr(x, qr_opts): + """Run the common native Torch QR block without composed dispatch.""" + absorb = qr_opts.get("absorb", "right") + left_like = absorb in { + -1, "left", "Us,VH", "lfactor", "Us", + } + qr_kwargs = { + key: value for key, value in qr_opts.items() + if key not in {"absorb", "stabilized"} + } + if left_like: + x = ar.do("transpose", x, (1, 0)) + q, r = ar.do("linalg.qr", x, **qr_kwargs) + if left_like: + left = ar.do("transpose", r, (1, 0)) + right = ar.do("transpose", q, (1, 0)) + if absorb in {-1, "left", "Us,VH"}: + return left, None, right + return left, None, None + if absorb in {"lorthog", "U", 10}: + return q, None, None + if absorb in {"rfactor", "sVH", 11}: + return None, None, r + return q, None, r + + try: + backend = ar.infer_backend(array) + except (AttributeError, TypeError): + backend = None + + # ``array_split`` calls this once per native charge block. For the + # Torch path, bypassing quimb's composed linalg wrapper removes a Python + # dispatch layer from every block while retaining exactly the same + # reduced QR and the same ``stabilized=False`` policy. + use_torch_qr = backend == "torch" if ar.get_dtype_name(array) != "complex64": + if use_torch_qr: + return torch_qr(array, opts) return _quimb_qr_stabilized(array, **opts) - block_max = to_float(ar.do("max", ar.do("abs", array))) + if backend == "torch": + block_max = float(array.detach().abs().amax().item()) + else: + block_max = to_float(ar.do("max", ar.do("abs", array))) if not np.isfinite(block_max) or block_max == 0.0: # Preserve the original failure behaviour for non-finite input, while # allowing genuinely empty structural sectors through unchanged. + if use_torch_qr: + return torch_qr(array, opts) return _quimb_qr_stabilized(array, **opts) # Values above this scale are not affected by the low-norm complex64 QR # failure and avoid an unnecessary multiply/divide pair. if block_max >= 2.0**-8: + if use_torch_qr: + return torch_qr(array, opts) return _quimb_qr_stabilized(array, **opts) _, exponent = np.frexp(block_max) scale = float(np.ldexp(1.0, -int(exponent))) - left, singular_values, right = _quimb_qr_stabilized( - array * scale, **opts, - ) + if use_torch_qr: + left, singular_values, right = torch_qr(array * scale, opts) + else: + left, singular_values, right = _quimb_qr_stabilized( + array * scale, **opts, + ) # ``absorb='left'`` is the LQ orientation: the left factor carries the # scale. All other QR orientations carry it in the right factor. @@ -187,11 +235,23 @@ def _contract_two_tensors(left, right, *, shared_ind=None): axes = ((left.inds.index(shared_ind),), (right.inds.index(shared_ind),)) try: if _is_symmray_array(left.data): - data = ar.do( - "tensordot", - left.data, + # A single shared tree bond is the hot native operation. Fused + # Symmray contraction is a good general default, but on Torch CPU + # it first builds larger fused block views for a contraction that + # is already one-leg local. Blockwise dispatch keeps the charge + # sectors small and avoids that temporary workspace. CUDA/Torch + # and other backends retain the fused path, which usually wins by + # reducing the number of small kernel launches. + mode = "fused" + if getattr(left.data, "backend", None) == "torch": + blocks = getattr(left.data, "blocks", None) + sample = next(iter(blocks.values()), None) if blocks else None + if not bool(getattr(sample, "is_cuda", False)): + mode = "blockwise" + data = left.data.tensordot( right.data, axes=axes, + mode=mode, preserve_array=True, ) else: @@ -286,6 +346,7 @@ class TreeTensorNetwork(TensorNetworkGenVector): "_symmetry", "_fermionic", "_physical_sectors", + "_work_bond_counter", ) def __init__(self, ts=(), *, plan=None, sites=None, site_tag_id="I{}", @@ -309,6 +370,7 @@ def __init__(self, ts=(), *, plan=None, sites=None, site_tag_id="I{}", self._fermionic_norm_cache = None self._fermionic_norm_cache_version = 0 self._fermionic_norm_cache_value_version = None + self._work_bond_counter = getattr(ts, "_work_bond_counter", 0) return if plan is None: raise TypeError( @@ -327,6 +389,7 @@ def __init__(self, ts=(), *, plan=None, sites=None, site_tag_id="I{}", self._fermionic_norm_cache = None self._fermionic_norm_cache_version = 0 self._fermionic_norm_cache_value_version = None + self._work_bond_counter = 0 self.validate() return super().__init__(ts, **tn_opts) @@ -350,6 +413,22 @@ def __init__(self, ts=(), *, plan=None, sites=None, site_tag_id="I{}", # Tracked here -- surviving ``.copy()`` via ``_EXTRA_PROPS`` -- so the # canonical form is a property of the *state*, not of any one driver. self._canonical_region = None + self._work_bond_counter = 0 + + def _new_work_bond(self, kind, *nodes): + """Return a unique private bond label for a native decomposition. + + Native QR/SVD factors are short-lived, but a routed operator bond can + become a live tree edge before the next hop. A state-owned counter is + cheaper than generating a UUID for every factor and, because it is an + extra copied property, remains collision-free across TTN branches. + """ + counter = int(getattr(self, "_work_bond_counter", 0)) + self._work_bond_counter = counter + 1 + suffix = "_".join(str(node) for node in nodes) + if suffix: + suffix = "_" + suffix + return f"_pepsy_{kind}_{counter}{suffix}" # -- mutation / canonical metadata -------------------------------------- @@ -1327,6 +1406,9 @@ def _fermionic_canonize_edge_(self, a, b, absorb): absorb="right", cutoff=0.0, get="tensors", + bond_ind=self._new_work_bond( + "canon", isometric_node, reduced_node, + ), ) merged = _contract_two_tensors(carry, reduced, shared_ind=bond) isometric.modify( @@ -1390,7 +1472,9 @@ def _fermionic_compress_edge_( # tensor can be thousands by thousands even though its shared # bond is only O(chi). The QR leaves a core whose right dimension # is the live bond, avoiding a full SVD of that large matrix. - reduced_bond = qtn.rand_uuid() + reduced_bond = self._new_work_bond( + "compress_qr", isometric_node, reduced_node, + ) isometric_q, isometric_r = self._native_qr_split( isometric, left_inds=left_inds, @@ -1400,7 +1484,9 @@ def _fermionic_compress_edge_( get="tensors", bond_ind=reduced_bond, ) - compressed_bond = qtn.rand_uuid() + compressed_bond = self._new_work_bond( + "compress_svd", isometric_node, reduced_node, + ) core_left, core_right = isometric_r.split( left_inds=(reduced_bond,), method="svd", @@ -1450,7 +1536,9 @@ def _fermionic_compress_edge_( right_inds = [ index for index in reduced_tensor.inds if index != bond ] - compressed_bond = qtn.rand_uuid() + compressed_bond = self._new_work_bond( + "compress_right", isometric_node, reduced_node, + ) core_left, core_right = reduced_tensor.split( left_inds=(bond,), right_inds=right_inds, @@ -1492,8 +1580,12 @@ def _fermionic_compress_edge_( right_inds = [ index for index in reduced_tensor.inds if index != bond ] - left_bond = qtn.rand_uuid() - right_bond = qtn.rand_uuid() + left_bond = self._new_work_bond( + "compress_left", isometric_node, reduced_node, + ) + right_bond = self._new_work_bond( + "compress_right_qr", isometric_node, reduced_node, + ) isometric_q, isometric_r = self._native_qr_split( isometric, left_inds=left_inds, diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index 97b5521..e9881c2 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -3217,6 +3217,46 @@ def test_tree_public_submpo_and_pauli_backend_operations(): ) +def test_tree_two_site_numpy_mpo_is_coerced_to_cupy_state_backend(): + """Two-site MPO factors follow a CuPy TTN without mutating the MPO.""" + cupy = pytest.importorskip("cupy") + try: + if cupy.cuda.runtime.getDeviceCount() < 1: + pytest.skip("CuPy is installed without a CUDA device.") + except cupy.cuda.runtime.CUDARuntimeError as exc: + pytest.skip(f"CuPy CUDA runtime unavailable: {exc}") + + n = 4 + plan = TreePlan.from_order(range(n), structure="balanced") + state = TreeTensorNetwork.from_plan(plan) + state.apply_to_arrays( + lambda array: cupy.asarray(array, dtype=cupy.complex64) + ) + x = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex) + mpo = qtn.MatrixProductOperator.from_dense( + np.kron(x, x), dims=(2, 2), sites=(0, 1), L=n, + ) + before = [tensor.data.copy() for tensor in mpo.tensors] + opt = TreeOptimizer( + None, state=state, tree=plan, chi=8, cutoff=0.0, run=False, + ) + + with pytest.warns(UserWarning, match="converting a gate/operator payload"): + opt.apply_submpo(mpo, (0, 1)) + + assert opt.backend_info() == { + "backend": "cupy", + "dtype": "complex64", + "device": str(cupy.cuda.Device()), + } + expected = np.zeros(2**n, dtype=np.complex64) + expected[12] = 1.0 # X_0 X_1 |0000> = |1100> + np.testing.assert_allclose(opt.to_dense(), expected, atol=1e-5) + assert opt.tn.validate() is opt.tn + for tensor, original in zip(mpo.tensors, before): + np.testing.assert_array_equal(tensor.data, original) + + def test_tree_expectation_mpo_is_batched_and_non_mutating(): """A structured MPO expectation uses one tree pass and preserves state.""" h = np.array([[1.0, 1.0], [1.0, -1.0]], dtype=complex) / np.sqrt(2.0) @@ -3272,13 +3312,16 @@ def test_tree_profile_report_is_opt_in(): dtype=complex, ) quiet = TreeOptimizer([(x, 0)], n=4, chi=4) - assert quiet.profile_report() == { - "enabled": False, - "events": [], - "by_kind": {}, - "native_compression_routes": {}, - "total_seconds": 0.0, - } + quiet_report = quiet.profile_report() + assert quiet_report["enabled"] is False + assert quiet_report["events"] == [] + assert quiet_report["by_kind"] == {} + assert quiet_report["native_compression_routes"] == {} + assert quiet_report["update_seconds"] == 0.0 + assert quiet_report["total_seconds"] == 0.0 + assert quiet_report["timing_semantics"][ + "total_seconds_is_sum_of_events_not_wall_time" + ] is True profiled = TreeOptimizer( [(x, 0), (cnot, (0, 3))], n=4, chi=4, profile=True, @@ -3288,6 +3331,11 @@ def test_tree_profile_report_is_opt_in(): assert report["events"] assert report["by_kind"]["update"]["count"] == 2 assert report["native_compression_routes"] == {} + assert report["by_kind"]["gate_factorization"]["count"] == 1 + assert report["by_kind"]["tensor_absorption"]["count"] >= 2 + assert report["by_kind"]["metadata_path"]["count"] == 1 + assert report["by_kind"]["center_movement"]["count"] >= 1 + assert report["update_seconds"] == report["by_kind"]["update"]["seconds"] assert report["total_seconds"] > 0.0 From c08fad5c2e3ee965da0fcf24b0ba3e99c1251bfb Mon Sep 17 00:00:00 2001 From: rezaquant Date: Mon, 3 Aug 2026 23:41:22 -0600 Subject: [PATCH 62/70] Add native tree-routed MPO API --- .github/skills/tree-optimizer/SKILL.md | 35 + AGENTS.md | 25 + docs/api/optimizers/tree.md | 107 ++ src/pepsy/__init__.py | 4 +- src/pepsy/optimizers/__init__.py | 2 + src/pepsy/optimizers/tree/__init__.py | 3 + src/pepsy/optimizers/tree/layout.py | 56 + src/pepsy/optimizers/tree/operators.py | 2030 ++++++++++++++++++++++++ src/pepsy/optimizers/tree/optimizer.py | 154 +- src/pepsy/optimizers/tree/ttn.py | 243 ++- src/pepsy/tensors/symmetric.py | 78 + tests/test_optimize_tree.py | 66 + tests/test_public_api.py | 4 +- tests/test_tree_mpo.py | 333 ++++ 14 files changed, 3124 insertions(+), 16 deletions(-) create mode 100644 src/pepsy/optimizers/tree/operators.py create mode 100644 tests/test_tree_mpo.py diff --git a/.github/skills/tree-optimizer/SKILL.md b/.github/skills/tree-optimizer/SKILL.md index c93c7a3..7d1ae0e 100644 --- a/.github/skills/tree-optimizer/SKILL.md +++ b/.github/skills/tree-optimizer/SKILL.md @@ -386,6 +386,41 @@ interface use the dense `to_dense()` fallback and remain subject to `updates` group edge events by support and include cumulative relative loss, analogous to the MPS infidelity trace. +### Tree-native MPO API + +When the consumer is a `TreeTensorNetwork`, `TreeMPO` is the primary operator +API. Use `TreePlan.to_tree_mpo(...)` or +`Fermion.to_tree_mpo(..., tree=plan)`: + +```python +tree_operator = fermion.to_tree_mpo( + hamiltonian=hamiltonian, + tree=plan, + compress=True, +) +energy = tree_operator.expectation(tree) +# equivalent exact readout through the state API: +energy = tree.expectation_mpo_exact(tree_operator, range(plan.n)) +``` + +`tree_operator.chain_mpo` is optional compatibility data for ordinary MPS/MPO +workflows. `TreePlan.to_mpo(...)` and `tree_mpo(...)` return that regular chain +MPO and attach the `TreeMPO`; they do not change the tree contraction route. +The chain MPO must not be moved into the tree, densified, or compressed as a +state update for exact tree measurement. + +For native fermionic Hamiltonians, one-, two-, and higher-site neutral terms +are fused and factorized from their native Symmray operator tensor over the +TreePlan Steiner subtree, then amalgamated into one charge-aware direct-sum +TTNO. This is the normal general-term route and is canonicalizable/compressible; +it is not a list of ordinary hyperedges. Structured observables may select a +smaller dedicated TTNO, such as the four-state eta-pair endpoint automaton. +`TreeMPO.canonicalize()` performs lossless native QR gauge fixing and +`TreeMPO.compress(cutoff=..., max_bond=...)` performs native graded SVD +truncation. Native operator QR uses the same centralized +`_native_qr_split_tensor` policy as tree-state QR, including the +`stabilized=False` structural-zero safeguard for Symmray arrays. + ## Noisy trajectory replay `run_trajectory_shots` and `run_coalesced_trajectory_shots` support diff --git a/AGENTS.md b/AGENTS.md index d4d88f9..967b725 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -88,6 +88,31 @@ The native `TreeTensorNetwork` QR policy is centralized in not change the separate `MpsOptimizer` QR implementation or globally patch Quimb/Symmray. +## Native TreeMPO contract + +The tree-native operator API lives in `pepsy.optimizers.tree.operators`: + +- `TreeMPO` is the primary tree measurement object. Prefer + `TreePlan.to_tree_mpo(...)` or `Fermion.to_tree_mpo(..., tree=plan)` when the + consumer is a `TreeTensorNetwork`. +- `TreePlan.to_mpo(...)` and `tree_mpo(...)` remain compatibility constructors + for the ordinary low-bond chain MPO. They attach the `TreeMPO`, but the + chain MPO is never moved into the tree, densified, or used as the tree + operator during `expectation_mpo_exact`. +- Neutral native term sums are factorized from their native Symmray tensors on + each term's TreePlan Steiner subtree and amalgamated into one direct-sum + TTNO. Do not replace this with a Jordan--Wigner dense factorization or a + list of ordinary hyperedges. The compact eta-pair observable is an explicit + structured exception with its four-state TTNO automaton. +- `TreeMPO.expectation(...)` and `TreeTensorNetwork.expectation_mpo_exact(...)` + contract separate bra, operator, and ket networks. `TreeMPO.canonicalize()` + is lossless native QR gauge fixing; `TreeMPO.compress(...)` is the explicit + native graded SVD truncation stage. +- Native operator QR must use the shared `_native_qr_split_tensor` policy (and + therefore the same `stabilized=False` structural-zero safeguard as the + state). Do not add direct `tensor.split(method="qr")` calls to TreeMPO + canonicalization. + ## Dependency and backend rules - Prefer public `quimb`, `cotengra`, `cotengrust`, and `autoray` APIs over diff --git a/docs/api/optimizers/tree.md b/docs/api/optimizers/tree.md index 26de114..fe7540a 100644 --- a/docs/api/optimizers/tree.md +++ b/docs/api/optimizers/tree.md @@ -46,6 +46,113 @@ structured sub-MPOs may include it in their support. `TreeLayoutFinder` keeps the site fixed at the root while its path, Steiner, congestion, greedy, and Nevergrad objectives permute only the remaining leaf sites. +### Layout-aware native MPOs + +After selecting a plan, build a Hamiltonian MPO in the plan's logical tree +order with `plan.to_mpo(...)` or `tree_mpo(plan, hamiltonian)`. The default is +the native graded Symmray path, including `U1FermionicArray` and +`U1U1FermionicArray` tensors: + +```python +from pepsy import Fermion +from pepsy.optimizers.tree import TreeLayoutFinder, tree_mpo + +finder = TreeLayoutFinder(gates, n=8, max_arity=2) +plan = finder.run(order="quality") +fermion = Fermion(spinful=True, symmetry="U1U1") +hamiltonian = fermion.hamiltonian(edges, t=1.0, U=2.0, mu=0.1) + +mpo = tree_mpo(plan, hamiltonian) # fermionic=True by default +# equivalent: +# mpo = plan.to_mpo(hamiltonian) +energy = opt.tn.expectation_mpo_exact(mpo, range(plan.n)) +``` + +For the operator-level API, use `TreeMPO` (or +`fermion.to_tree_mpo(...)`). It exposes both representations without mixing +their tensor-network geometries: + +```python +from pepsy import Fermion + +tree_operator = fermion.to_tree_mpo( + hamiltonian=hamiltonian, + tree=plan, + compress=True, +) +chain_operator = tree_operator.chain_mpo # optional linear MPO +energy = tree_operator.expectation(opt.tn) # TreePlan-native readout +``` + +`TreeMPO.from_terms(plan, dense_terms)` provides the corresponding ordinary +dense backend. Native fermionic `TreeMPO` objects retain Symmray arrays for +U1, U1U1, and other supported symmetries; dense and native operators are not +silently mixed with incompatible tree states. + +For an exact native readout that keeps the chain MPO separate, use +`expectation_mpo_exact` as above. The general `expectation_mpo` API remains +available for an explicitly approximate structured-MPO application: it uses a +private transformed copy of the TTN. Routing itself is lossless, but the final +subtree sweep may compress that copy when an MPO increases a bond beyond +`max_bond`; `cutoff=0.0` does not disable a finite bond cap. The default +`warn_on_truncation=True` reports this approximation, and the diagnostic form +makes it easy to check a benchmark: + +```python +energy, report = opt.tn.expectation_mpo( + mpo, + range(plan.n), + max_bond=64, + cutoff=0.0, + return_diagnostics=True, +) +assert not report["truncated"] +``` + +Use a larger `max_bond` when an untruncated measurement is required. Native +fermionic trees also reject ordinary dense MPOs (and dense trees reject native +Symmray MPOs) instead of silently changing the fermionic interpretation. + +For an exact readout that does not move the MPO into the tree at all, use +`expectation_mpo_exact`: + +```python +energy = opt.tn.expectation_mpo_exact( + mpo, + range(plan.n), +) +``` + +The same method is available directly as `opt.expectation_mpo_exact(...)`. + +This keeps the bra, ket, and structured MPO as separate networks. Fresh ket +physical indices connect to the MPO input legs, the MPO output legs connect to +the bra, and Quimb contracts the complete doubled network. No state-bond +compression or `to_dense()` lowering occurs. Native Symmray MPOs retain their +graded contraction and fermionic sign rules. + +`TreePlan.mpo_order()` is the structural leaf-position order chosen by the +plan; a physical `root_qubit`, when present, is placed first. The helper +constructs the chain MPO in that one-dimensional order, then restores the +physical labels to their logical qubit numbers. It returns the ordinary Quimb +MPO for compatibility, while exact native readout uses the separate +TreePlan-aware embedding attached by the builder. + +`TreeMPO` is the tree-routed operator class. Its optional `chain_mpo` remains a +linear MPO for MPS workflows, while `tree_networks` contains the +TreePlan-labelled operator networks used by `expectation`. Native neutral +terms are factorized directly from their native Symmray operator tensor over +the term's TreePlan Steiner subtree, then amalgamated into one charge-aware +direct-sum TTNO. This applies to one-, two-, and higher-site native terms; it +does not create a hyperedge for the normal Hamiltonian path. The resulting +TTNO can be canonicalized and compressed with +`tree_operator.canonicalize()` and +`tree_operator.compress(cutoff=..., max_bond=...)`; no Jordan--Wigner +conversion is used. Structured observables can use a smaller compact TTNO. +Pass `fermionic=False` only for dense ordinary/Jordan--Wigner-compatible terms. +Existing `OneDMap` lattice maps remain unchanged and should continue to be used +for regular 2D/3D coordinate layouts. + The conventional binary TTN with a three-leg top tensor is the default when there are at least three leaves and no `root_qubit`. Pass `max_arity=2, top_arity=3` explicitly to `TreePlan.from_order`, diff --git a/src/pepsy/__init__.py b/src/pepsy/__init__.py index e69b031..4acfc2a 100644 --- a/src/pepsy/__init__.py +++ b/src/pepsy/__init__.py @@ -153,8 +153,10 @@ "SweepOptimizer": ".optimizers", "TreeEnergyOptimizer": ".optimizers", "TreeLayoutFinder": ".optimizers", + "TreeMPO": ".optimizers", "TreeOptimizer": ".optimizers", "TreePlan": ".optimizers", + "tree_mpo": ".optimizers", "square_lattice_zigzag": ".optimizers", "TreeStabOptimizer": ".optimizers", "TreeTensorNetwork": ".optimizers", @@ -360,7 +362,7 @@ def __getattr__(name): y, z, ) - from .optimizers import DeferredInjectionRecord, DeferredInjectionReport, DeferredProjectionRecord, GlobalOptimizer, ImmediateInjectionReport, ImmediateProjectionRecord, MeasurementRecord, MpoOptimizer, MpsEnergyOptimizer, MpsOptimizer, MpsStabOptimizer, NormEventRecord, PepsEnergyOptimizer, PepsOptimizer, STNState, StabilizerMpsSettingsAdvice, StabilizerMpsRunResult, StabilizerMpsSimulator, StabilizerTreeRunResult, StreamAnalysisRecord, SimpleUpdateGen, SymDMRG2, SweepOptimizer, TreeEnergyOptimizer, TreeLayoutFinder, TreeOptimizer, TreePlan, TreeStabOptimizer, run_stabilizer_mps_stream, run_stabilizer_tree_stream # noqa: F401 + from .optimizers import DeferredInjectionRecord, DeferredInjectionReport, DeferredProjectionRecord, GlobalOptimizer, ImmediateInjectionReport, ImmediateProjectionRecord, MeasurementRecord, MpoOptimizer, MpsEnergyOptimizer, MpsOptimizer, MpsStabOptimizer, NormEventRecord, PepsEnergyOptimizer, PepsOptimizer, STNState, StabilizerMpsSettingsAdvice, StabilizerMpsRunResult, StabilizerMpsSimulator, StabilizerTreeRunResult, StreamAnalysisRecord, SimpleUpdateGen, SymDMRG2, SweepOptimizer, TreeEnergyOptimizer, TreeLayoutFinder, TreeMPO, TreeOptimizer, TreePlan, TreeStabOptimizer, run_stabilizer_mps_stream, run_stabilizer_tree_stream # noqa: F401 from .sampling import FermionConfigurationEncoding, MpsDiagonalEstimate, MpsBatchSampleResult, MpsSampleResult, MpsSampler, PEPSSampleResult, PepsBpSampler, TreeBatchSampleResult, TreeSampleResult, TreeSampler, VecSampler # noqa: F401 from .solvers import FDSolver # noqa: F401 from .tensors import ( # noqa: F401 diff --git a/src/pepsy/optimizers/__init__.py b/src/pepsy/optimizers/__init__.py index 528baa4..8d92680 100644 --- a/src/pepsy/optimizers/__init__.py +++ b/src/pepsy/optimizers/__init__.py @@ -39,8 +39,10 @@ "SymDMRG2": ".sym_dmrg", "SweepOptimizer": ".sweep", "TreeLayoutFinder": ".tree", + "TreeMPO": ".tree", "TreeOptimizer": ".tree", "TreePlan": ".tree", + "tree_mpo": ".tree", "square_lattice_zigzag": "._layout_orders", "TreeStabOptimizer": ".tree_stabilizer", "TreeTensorNetwork": ".tree", diff --git a/src/pepsy/optimizers/tree/__init__.py b/src/pepsy/optimizers/tree/__init__.py index eeb30d8..33c8a66 100644 --- a/src/pepsy/optimizers/tree/__init__.py +++ b/src/pepsy/optimizers/tree/__init__.py @@ -8,6 +8,7 @@ from .layout import TreeLayoutFinder, TreePlan from .optimizer import TreeOptimizer +from .operators import TreeMPO, tree_mpo from .ttn import TreeTensorNetwork __all__ = [ @@ -15,4 +16,6 @@ "TreeTensorNetwork", "TreeLayoutFinder", "TreePlan", + "TreeMPO", + "tree_mpo", ] diff --git a/src/pepsy/optimizers/tree/layout.py b/src/pepsy/optimizers/tree/layout.py index 8859868..f765915 100644 --- a/src/pepsy/optimizers/tree/layout.py +++ b/src/pepsy/optimizers/tree/layout.py @@ -1041,6 +1041,62 @@ def leaves(self): """Return the leaf node ids.""" return list(self.qubit_of_leaf.keys()) + def mpo_order(self, *, include_root=True): + """Return the deterministic logical-site order for a tree MPO. + + The leaf positions are ordered by their structural node ids, matching + the order used by :class:`TreeLayoutFinder` when it refines a plan. + When ``root_qubit`` is present, that physical site is placed first by + default, followed by the ordinary leaf positions. The result is a + permutation of ``0 .. n - 1`` and is suitable for constructing a + layout-aware chain MPO whose sites are subsequently routed over this + tree. + + Parameters + ---------- + include_root : bool, optional + Include a physical site carried by the structural root. The + default is ``True``; pass ``False`` to obtain only the leaf order. + """ + order = tuple( + self.qubit_of_leaf[nid] + for nid in sorted(self.qubit_of_leaf) + ) + if include_root and self.root_qubit is not None: + return (self.root_qubit, *order) + return order + + def to_mpo(self, hamiltonian, **kwargs): + """Build a native chain MPO and its TreePlan embedding. + + This delegates to :func:`pepsy.optimizers.tree.tree_mpo`. The returned + object is the ordinary Quimb ``MatrixProductOperator`` with native + Symmray tensors when ``fermionic=True``. Its chain order follows + :meth:`mpo_order`; exact native tree readout uses the separate + TreePlan-routed operator attached by the builder and contracts the + doubled ``tree.H | operator | tree`` network. + """ + from .operators import tree_mpo + + return tree_mpo(self, hamiltonian, **kwargs) + + def to_tree_mpo(self, hamiltonian, **kwargs): + """Build the public :class:`TreeMPO` operator for this plan. + + The returned class keeps the optional chain MPO available as + ``.chain_mpo`` and exposes the TreePlan-routed representation through + ``.tree_networks`` and ``.expectation``. + """ + from .operators import tree_mpo + + chain_mpo = tree_mpo(self, hamiltonian, **kwargs) + if isinstance(chain_mpo, dict): + return { + charge: mpo.pepsy_tree_operator + for charge, mpo in chain_mpo.items() + } + return chain_mpo.pepsy_tree_operator + def is_leaf(self, nid): return len(self.children.get(nid, ())) == 0 diff --git a/src/pepsy/optimizers/tree/operators.py b/src/pepsy/optimizers/tree/operators.py new file mode 100644 index 0000000..efe53cc --- /dev/null +++ b/src/pepsy/optimizers/tree/operators.py @@ -0,0 +1,2030 @@ +"""Tree-plan-aware native fermionic operator construction. + +``SymHamiltonian.to_mpo`` is still used to build the ordinary, low-bond chain +MPO. A chain's Jordan--Wigner wire, however, is not an embedding of a +branched fermionic tree. This module therefore also constructs a native tree +operator, without applying it to the state. Exact tree readout contracts +that operator between a private bra and ket copy. + +Two routes are provided: + +* general neutral native term sums are decomposed term-by-term on the + TreePlan Steiner subtrees and amalgamated into one direct-sum TTNO; +* a rank-one pair correlator with separable coefficients is compiled into a + four-state endpoint automaton. This is the compact route for the full + staggered eta-pair observable and keeps its tree bond independent of the + lattice size. + +The returned public object remains a regular Quimb MPO for compatibility with +MPS/MPO APIs. It carries the tree operator as private metadata consumed by +``TreeTensorNetwork.expectation_mpo_exact``. The two networks remain +separate throughout the contraction. +""" + +from __future__ import annotations + +import heapq +from numbers import Integral +import warnings + +import numpy as np + +from .layout import TreePlan + +__all__ = ["TreeMPO", "tree_mpo"] + + +def _tree_plan_signature(plan): + """Return a stable structural signature for a tree-MPO annotation.""" + return ( + int(plan.root), + tuple( + (int(node), tuple(int(child) for child in children)) + for node, children in sorted(plan.children.items()) + ), + tuple( + (int(node), int(qubit)) + for node, qubit in sorted(plan.qubit_of_leaf.items()) + ), + None if plan.root_qubit is None else int(plan.root_qubit), + ) + + +class TreeMPO: + """TreePlan-aware operator with dense and native Symmray backends. + + ``TreeMPO`` is the operator-level API for measurements on a + :class:`TreeTensorNetwork`. It deliberately keeps the optional linear + chain MPO separate from the tree representation: + + ``chain_mpo`` + The ordinary Quimb ``MatrixProductOperator`` produced by + ``SymHamiltonian.to_mpo``. This is useful for MPS workflows. + + ``tree_networks`` + One or more operator tensor networks whose physical indices are + labelled by the logical qubits in ``plan``. General native terms are + combined into one Symmray TTNO whose graded source/target channels + preserve the fermionic contraction rules. + + General neutral native sums use one direct-sum TTNO. Each term is first + factorized on its native graded TreePlan subtree, then all term channels + are amalgamated on common charge-aware virtual bonds. The resulting + operator can be canonicalized and compressed with native graded QR/SVD. + Structured observables such as the eta-pair table may use a smaller + compact network instead. + """ + + def __init__( + self, + plan, + tree_networks, + *, + chain_mpo=None, + terms=None, + backend="dense", + fermionic=False, + symmetry=None, + cutoff=1e-12, + compressed=False, + ): + if not isinstance(plan, TreePlan): + raise TypeError("plan must be a TreePlan.") + if isinstance(tree_networks, (tuple, list)): + networks = tuple(tree_networks) + else: + networks = (tree_networks,) + if not networks or any(network is None for network in networks): + raise ValueError("TreeMPO requires at least one tree operator network.") + self.plan = plan + self.tree_networks = networks + self.chain_mpo = chain_mpo + self.terms = None if terms is None else dict(terms) + self.backend = str(backend) + self.fermionic = bool(fermionic) + self.symmetry = symmetry + self.cutoff = float(cutoff) + self.compressed = bool(compressed) + self.pepsy_tree_plan_signature = _tree_plan_signature(plan) + if chain_mpo is not None: + chain_mpo.pepsy_tree_plan_signature = self.pepsy_tree_plan_signature + chain_mpo.pepsy_tree_terms = ( + None if self.terms is None else dict(self.terms) + ) + chain_mpo.pepsy_tree_operator = self + chain_mpo.pepsy_tree_operator_networks = self.tree_networks + + @classmethod + def from_hamiltonian( + cls, + plan, + hamiltonian, + *, + chain_mpo=None, + cutoff=1e-12, + max_bond=None, + compress=True, + dtype=None, + fermionic=True, + ): + """Construct a ``TreeMPO`` from a ``SymHamiltonian``.""" + from ...tensors.symmetric import SymHamiltonian + + if not isinstance(hamiltonian, SymHamiltonian): + raise TypeError("hamiltonian must be a SymHamiltonian instance.") + networks = _build_tree_operator( + plan, + hamiltonian, + cutoff=cutoff, + max_bond=max_bond, + compress=compress, + dtype=dtype, + fermionic=fermionic, + ) + if isinstance(networks, (tuple, list)): + native_networks = tuple(networks) + else: + native_networks = (networks,) + backend = "symmray" if fermionic else "dense" + operator = cls( + plan, + native_networks, + chain_mpo=chain_mpo, + terms=hamiltonian.terms, + backend=backend, + fermionic=fermionic, + symmetry=hamiltonian.symmetry, + cutoff=cutoff, + compressed=compress, + ) + if compress: + operator.compress(max_bond=max_bond, cutoff=cutoff) + return operator + + @classmethod + def from_terms( + cls, + plan, + terms, + *, + chain_mpo=None, + cutoff=1e-12, + dtype=None, + max_bond=None, + compress=True, + ): + """Construct one ordinary dense TTNO from a term mapping. + + ``terms`` maps an integer site or support tuple to a dense operator + array. The dense route is useful for non-fermionic trees and for + callers that already have Jordan--Wigner-compatible local matrices. + """ + if not hasattr(terms, "items"): + raise TypeError("terms must be a mapping of supports to operators.") + network = _combined_tree_operator( + plan, + terms, + symmetry=None, + cutoff=cutoff, + dtype=dtype, + fermionic=False, + ) + operator = cls( + plan, + network, + chain_mpo=chain_mpo, + terms=terms, + backend="dense", + fermionic=False, + cutoff=cutoff, + compressed=compress, + ) + if compress: + operator.compress(max_bond=max_bond, cutoff=cutoff) + return operator + + @property + def tree_network(self): + """Return the sole tree network, or raise for a term sum.""" + if len(self.tree_networks) != 1: + raise AttributeError( + "this TreeMPO contains multiple internal networks; use " + "tree_networks or expectation()" + ) + return self.tree_networks[0] + + def max_bond(self): + """Return the largest virtual bond among the tree networks.""" + bonds = [] + for network in self.tree_networks: + for index in network.inner_inds(): + bonds.append(network.ind_size(index)) + return max(bonds, default=1) + + def canonicalize(self, center=None): + """Canonicalize every stored TTNO around one TreePlan node.""" + if center is None: + center = self.plan.root + for network in self.tree_networks: + _canonicalize_tree_operator(network, self.plan, center) + return self + + def compress(self, *, max_bond=None, cutoff=None): + """Compress the TTNO on every TreePlan edge with native SVD.""" + if cutoff is None: + cutoff = self.cutoff + cutoff = float(cutoff) + reports = [] + for network in self.tree_networks: + reports.append(_compress_tree_operator( + network, + self.plan, + max_bond=max_bond, + cutoff=cutoff, + )) + self.cutoff = cutoff + self.compressed = True + self.pepsy_compression_report = reports[0] if len(reports) == 1 else reports + return self + + def copy(self): + """Copy the operator and both of its optional representations.""" + chain_mpo = None if self.chain_mpo is None else self.chain_mpo.copy() + copied = type(self)( + self.plan, + tuple(network.copy() for network in self.tree_networks), + chain_mpo=chain_mpo, + terms=self.terms, + backend=self.backend, + fermionic=self.fermionic, + symmetry=self.symmetry, + cutoff=self.cutoff, + compressed=self.compressed, + ) + if chain_mpo is not None: + chain_mpo.pepsy_tree_plan_signature = copied.pepsy_tree_plan_signature + chain_mpo.pepsy_tree_terms = ( + None if copied.terms is None else dict(copied.terms) + ) + chain_mpo.pepsy_tree_operator = copied + chain_mpo.pepsy_tree_operator_networks = copied.tree_networks + return copied + + def expectation(self, state, *, normalized=True, optimize="auto"): + """Evaluate ```` in one public operation.""" + import quimb.tensor as qtn # pylint: disable=import-outside-toplevel + + tree = getattr(state, "tn", state) + if getattr(tree, "plan", None) is None: + raise TypeError("state must be a TreeTensorNetwork or TreeOptimizer.") + if _tree_plan_signature(tree.plan) != self.pepsy_tree_plan_signature: + raise ValueError("TreeMPO and state use different TreePlans.") + if self.fermionic and not getattr(tree, "fermionic", False): + raise TypeError("native TreeMPO requires a native fermionic tree state.") + if self.fermionic is False and getattr(tree, "fermionic", False): + raise TypeError("dense TreeMPO cannot be contracted with a native fermionic tree.") + + sites = tuple(sorted(tree.plan.node_of_qubit)) + numerator = 0.0 + for operator in self.tree_networks: + ket = tree.copy() + operator_work = operator.copy() + ket_reindex = {} + operator_reindex = {} + for site in sites: + physical = tree.site_ind(site) + upper = f"k{site}" + lower = f"b{site}" + if upper not in operator_work.ind_map or lower not in operator_work.ind_map: + raise ValueError(f"TreeMPO is missing physical site {site!r}.") + fresh = qtn.rand_uuid() + ket_reindex[physical] = fresh + operator_reindex[lower] = fresh + ket.reindex_(ket_reindex) + operator_work.reindex_(operator_reindex) + numerator = numerator + (tree.H | operator_work | ket).contract( + all, + optimize=optimize, + ) + if not normalized: + return numerator + denominator = (tree.H | tree).contract(all, optimize=optimize) + return numerator / denominator + + def __repr__(self): + return ( + f"TreeMPO(nsite={self.plan.n}, backend={self.backend!r}, " + f"networks={len(self.tree_networks)}, max_bond={self.max_bond()})" + ) + + def __getattr__(self, name): + """Preserve the old metadata attributes on compact tree operators.""" + if name.startswith("pepsy_tree_operator_"): + networks = self.__dict__.get("tree_networks", ()) + if len(networks) == 1: + return getattr(networks[0], name) + raise AttributeError(name) + + +def _term_support(where): + """Normalize one Hamiltonian key to an integer support tuple.""" + if isinstance(where, Integral): + support = (int(where),) + else: + try: + support = tuple(int(site) for site in where) + except (TypeError, ValueError) as exc: + raise TypeError( + "tree MPO Hamiltonian term locations must be integer sites " + "or tuples of integer sites." + ) from exc + if not support: + raise ValueError("tree MPO Hamiltonian term supports cannot be empty.") + if len(set(support)) != len(support): + raise ValueError( + f"tree MPO Hamiltonian term support {support!r} repeats a site." + ) + return support + + +def _expanded_index_charges(index): + """Expand a native block index into its dense charge ordering.""" + chargemap = getattr(index, "chargemap", None) + if chargemap is None: + raise TypeError("native operator factors must expose block charges.") + return [charge for charge, size in chargemap.items() for _ in range(size)] + + +def _operator_charge_neg(charge): + """Negate one Abelian charge used by a native operator channel.""" + if isinstance(charge, tuple): + return tuple(-value for value in charge) + return -charge + + +def _operator_charge_sub(left, right): + """Subtract two expanded physical charges componentwise.""" + if isinstance(left, tuple): + return tuple(a - b for a, b in zip(left, right)) + return left - right + + +def _operator_charge_from_matrix(data, physical_map, *, tol=1e-10): + """Infer the homogeneous local operator charge from a dense matrix.""" + values = { + _operator_charge_sub(physical_map[int(out)], physical_map[int(inp)]) + for out, inp in np.argwhere(np.abs(data) > tol) + } + if len(values) != 1: + raise ValueError( + "operator-Schmidt factors must have one homogeneous physical " + f"charge, got {sorted(values, key=repr)!r}." + ) + return values.pop() + + +def _operator_native_channels( + term, support, *, symmetry, dtype=None, cutoff=1e-12, +): + """Split one native two-site term into charged local operator channels.""" + original_support = tuple(int(site) for site in support) + support = tuple(sorted(original_support)) + if len(support) != 2: + raise ValueError("native operator channels require two sites.") + if tuple(_term_support(support)) != support: + raise ValueError("operator support must contain distinct sites.") + ordered_term = term + if original_support != support: + ordered_term = term.transpose((1, 0, 3, 2)) + + # ``cutoff=0`` retains structural zero sectors as separate channels. The + # small fixed threshold removes only those exact numerical zeros; the + # user-facing TreeMPO cutoff is applied later to the combined TTNO. + structural_cutoff = 64.0 * np.finfo(float).eps + fused = ordered_term.fuse((0, 2), (1, 3)) + left, _, right = fused.svd( + absorb="right", cutoff=structural_cutoff, + ) + if left is None or right is None: + raise ValueError("could not split a native two-site operator.") + left = left.unfuse(0).transpose((2, 0, 1)) + right = right.unfuse(1) + left_data = np.asarray(left.to_dense(), dtype=dtype) + right_data = np.asarray(right.to_dense(), dtype=dtype) + physical_map = _expanded_index_charges(left.indices[1]) + if _expanded_index_charges(left.indices[2]) != physical_map: + raise ValueError("native operator factors have mismatched physical maps.") + channels = [] + for channel in range(left_data.shape[0]): + source = left_data[channel] + target = right_data[channel] + if not np.any(np.abs(source) > 1e-10): + continue + if not np.any(np.abs(target) > 1e-10): + raise ValueError("native operator SVD produced an empty channel.") + source_charge = _operator_charge_from_matrix(source, physical_map) + target_charge = _operator_charge_from_matrix(target, physical_map) + if target_charge != _operator_charge_neg(source_charge): + raise ValueError( + "native operator channel charges do not cancel: " + f"{source_charge!r} and {target_charge!r}." + ) + channels.append((source, target, source_charge)) + if not channels: + raise ValueError("native two-site operator has no nonzero channels.") + return channels, physical_map + + +def _operator_dense_channels(operator, support, *, dtype=None, cutoff=1e-12): + """Split one ordinary dense two-site term into local channels.""" + support = tuple(sorted(int(site) for site in support)) + data = _dense_operator_array(operator, dtype=dtype) + if data.ndim != 4 or data.shape[0] != data.shape[1] or data.shape[0] != data.shape[2]: + raise ValueError("dense two-site operators must have shape (d, d, d, d).") + if data.shape[2] != data.shape[3]: + raise ValueError("dense two-site operators must have matching input legs.") + dim = data.shape[0] + matrix = data.transpose(0, 2, 1, 3).reshape(dim * dim, dim * dim) + left, singular, right = np.linalg.svd(matrix, full_matrices=False) + structural_cutoff = max(64.0 * np.finfo(float).eps, float(cutoff)) + channels = [] + for channel, value in enumerate(singular): + if float(value) <= structural_cutoff: + continue + scale = np.sqrt(value) + channels.append(( + (left[:, channel] * scale).reshape(dim, dim), + (scale * right[channel, :]).reshape(dim, dim), + 0, + )) + if not channels: + raise ValueError("dense two-site operator has no nonzero channels.") + return channels, [0] * dim + + +def _operator_valid_child_states(nchildren, nstate, nchannel, done): + """Generate only the valid sparse automaton child configurations.""" + if not nchildren: + return [()] + states = {(0,) * nchildren} + active_states = tuple(range(1, done)) + for child in range(nchildren): + for state in active_states: + values = [0] * nchildren + values[child] = state + states.add(tuple(values)) + values = [0] * nchildren + values[child] = done + states.add(tuple(values)) + nchannel = int(nchannel) + source = lambda channel: 1 + channel + target = lambda channel: 1 + nchannel + channel + for left in range(nchildren): + for right in range(left + 1, nchildren): + for channel in range(nchannel): + for first, second in ( + (source(channel), target(channel)), + (target(channel), source(channel)), + ): + values = [0] * nchildren + values[left] = first + values[right] = second + states.add(tuple(values)) + return tuple(sorted(states)) + + +def _combined_tree_operator( + plan, terms, *, symmetry=None, cutoff=1e-12, dtype=None, fermionic=True, +): + """Build one TreePlan TTNO for a neutral one-/two-site term mapping. + + Each two-site operator-Schmidt channel becomes a source/target charge + channel. At a branching node, the channel automaton can collect one source + and one target from different child subtrees before closing into the + neutral ``done`` sector. This is the tree analogue of the start/channel/ + done construction used by a native chain MPO, but the channels follow the + selected TreePlan rather than a Jordan--Wigner wire. + """ + import quimb.tensor as qtn # pylint: disable=import-outside-toplevel + + if not hasattr(terms, "items") or not terms: + raise ValueError("At least one operator term is required.") + + channels = [] + one_site = {} + physical_map = None + for where, term in terms.items(): + support = _term_support(where) + if any(site not in plan.node_of_qubit for site in support): + raise ValueError(f"operator support {support!r} is outside the TreePlan.") + if len(support) == 1: + data = ( + _dense_operator_array(term, dtype=dtype) + if not fermionic else np.asarray(term.to_dense(), dtype=dtype) + ) + if data.ndim != 2 or data.shape[0] != data.shape[1]: + raise ValueError("one-site operators must be square matrices.") + if physical_map is None: + physical_map = ( + _expanded_index_charges(term.indices[0]) + if fermionic else [0] * data.shape[0] + ) + if fermionic: + charge = _operator_charge_from_matrix(data, physical_map) + zero = ( + tuple(0 for _ in charge) if isinstance(charge, tuple) else 0 + ) + if charge != zero: + raise ValueError( + "combined native TreeMPO currently requires neutral " + "one-site terms." + ) + site = support[0] + one_site[site] = one_site.get(site, 0) + data + continue + if len(support) != 2: + raise NotImplementedError( + "combined TreeMPO currently supports one- and two-site terms; " + "use a precompiled structured TTNO for higher-rank terms." + ) + if fermionic: + term_channels, term_map = _operator_native_channels( + term, support, symmetry=symmetry, dtype=dtype, cutoff=cutoff, + ) + else: + term_channels, term_map = _operator_dense_channels( + term, support, dtype=dtype, cutoff=cutoff, + ) + if physical_map is None: + physical_map = list(term_map) + elif list(term_map) != list(physical_map): + raise ValueError("all TreeMPO terms must share one physical map.") + for source, target, charge in term_channels: + channels.append({ + "source": source, + "target": target, + "charge": charge if fermionic else 0, + "sites": tuple(sorted(support)), + }) + + if physical_map is None: + raise ValueError("At least one operator term is required.") + if fermionic: + first_charge = physical_map[0] + zero = tuple(0 for _ in first_charge) if isinstance(first_charge, tuple) else 0 + else: + zero = 0 + nchannel = len(channels) + if not nchannel and not one_site: + raise ValueError("operator terms produced no nonzero channels.") + + source_id = lambda channel: 1 + channel + target_id = lambda channel: 1 + nchannel + channel + done = 1 + 2 * nchannel + state_map = [zero] + state_map.extend( + _operator_charge_neg(channel["charge"]) for channel in channels + ) + state_map.extend(channel["charge"] for channel in channels) + state_map.append(zero) + physical_dim = len(physical_map) + tensors = [] + + for node in plan.nodes(): + children = tuple(plan.children[node]) + parent = plan.parent.get(node) + has_parent = parent is not None + neighbors = list(children) + ([parent] if has_parent else []) + maps = [state_map] * len(neighbors) + duals = [True] * len(children) + ([False] if has_parent else []) + inds = [ + f"_pepsy_tnno_{min(node, neighbor)}_{max(node, neighbor)}" + for neighbor in neighbors + ] + qubit = plan.qubit_of_node.get(node) + if qubit is not None: + # Native tree leaves conventionally expose physical legs before + # their single virtual parent. This ordering is not cosmetic: + # Symmray's graded contraction phase depends on the ordered leg + # exterior. Keep the TTNO leaf in the same convention as the + # native state and direct local-observable route. + maps = [physical_map, physical_map] + maps + duals = [False, True] + duals + inds = [f"k{qubit}", f"b{qubit}"] + inds + shape = [len(index_map) for index_map in maps] + data = np.zeros(shape, dtype=dtype or complex) + identity = ( + np.eye(physical_dim, dtype=data.dtype) + if qubit is not None else 1.0 + ) + endpoint = {} + for channel, info in enumerate(channels): + if qubit == info["sites"][0]: + endpoint.setdefault("source", []).append(channel) + if qubit == info["sites"][1]: + endpoint.setdefault("target", []).append(channel) + + valid_children = _operator_valid_child_states( + len(children), len(state_map), nchannel, done, + ) + for child_states in valid_children: + active = [] + completed = False + invalid = False + for state in child_states: + if state == 0: + continue + if state == done: + if completed or active: + invalid = True + completed = True + continue + if state < source_id(nchannel): + channel = state - 1 + flag = "source" + else: + channel = state - target_id(0) + flag = "target" + if completed or any( + old_channel == channel and old_flag == flag + for old_channel, old_flag in active + ): + invalid = True + active.append((channel, flag)) + if invalid or len({channel for channel, _ in active}) > 1: + continue + if completed: + base = done + elif active: + flags = {flag for _, flag in active} + channel = active[0][0] + base = ( + done + if flags == {"source", "target"} + else source_id(channel) + if "source" in flags + else target_id(channel) + ) + else: + base = 0 + + options = [(base, identity)] + if base == 0 and qubit in one_site: + options.append((done, one_site[qubit])) + if base != done: + for channel in endpoint.get("source", ()): + local = channels[channel]["source"] + if base == 0: + options.append((source_id(channel), local)) + elif base == target_id(channel): + options.append((done, local)) + for channel in endpoint.get("target", ()): + local = channels[channel]["target"] + if base == 0: + options.append((target_id(channel), local)) + elif base == source_id(channel): + options.append((done, local)) + + for output, local in options: + if not has_parent and output != done: + continue + index = child_states + (output,) if has_parent else child_states + if qubit is not None: + data[(slice(None), slice(None)) + index] += local + else: + data[index] += local + + if fermionic: + native = _native_from_dense( + data, + symmetry=symmetry, + index_maps=maps, + duals=duals, + charge=zero, + label=None, + ) + else: + native = data + tags = [f"N{node}"] + if qubit is not None: + tags.append(f"I{qubit}") + tensors.append(qtn.Tensor(native, inds=inds, tags=tags)) + + network = qtn.TensorNetwork(tensors) + network.pepsy_tree_operator_kind = ( + "native_tree_tnno" if fermionic else "dense_tree_tnno" + ) + network.pepsy_tree_operator_bond = len(state_map) + network.pepsy_tree_operator_raw_bond = len(state_map) + network.pepsy_tree_operator_is_ttno = True + return network + + +def _tree_plan_neighbors(plan, node): + """Return a plan node's children followed by its optional parent.""" + return tuple(plan.children[node]) + ( + (plan.parent[node],) if plan.parent.get(node) is not None else () + ) + + +def _tree_operator_peel_order(plan, nodes): + """Return a deterministic leaf-to-hub order for a connected node set.""" + remaining = set(nodes) + adjacency = { + node: tuple( + neighbor for neighbor in _tree_plan_neighbors(plan, node) + if neighbor in remaining + ) + for node in remaining + } + degree = { + node: sum(neighbor in remaining for neighbor in neighbors) + for node, neighbors in adjacency.items() + } + leaves = [node for node, value in degree.items() if value == 1] + heapq.heapify(leaves) + order = [] + while len(remaining) > 1: + while leaves and leaves[0] not in remaining: + heapq.heappop(leaves) + if not leaves: + raise ValueError("operator decomposition requires a connected tree") + leaf = heapq.heappop(leaves) + neighbor = next( + node for node in adjacency[leaf] if node in remaining + ) + order.append((leaf, neighbor)) + remaining.remove(leaf) + degree[leaf] = 0 + degree[neighbor] -= 1 + if degree[neighbor] == 1: + heapq.heappush(leaves, neighbor) + return tuple(order), next(iter(remaining)) + + +def _native_tree_term_network( + plan, term, support, *, symmetry, cutoff=1e-12, dtype=None, +): + """Decompose one native term into a graded TTNO on the selected tree. + + The decomposition is performed on the native operator tensor itself, not + on a dense Jordan--Wigner matrix. The physical upper/lower pair at every + supported site is fused into one packed leg and the resulting tensor is + peeled across the TreePlan Steiner subtree with native Symmray SVDs. This + is the important fermionic distinction from factorizing ordinary dense + local matrices: the native fuse/SVD retains the graded phases at every + branch of the tree. + """ + import quimb.tensor as qtn # pylint: disable=import-outside-toplevel + + support = _term_support(support) + if any(site not in plan.node_of_qubit for site in support): + raise ValueError(f"term support {support!r} is outside the TreePlan.") + + physical_map = _expanded_index_charges(term.indices[0]) + zero = ( + tuple(0 for _ in physical_map[0]) + if physical_map and isinstance(physical_map[0], tuple) + else 0 + ) + term_charge = getattr(term, "charge", zero) + if term_charge != zero: + raise ValueError( + "a single native TTNO must be neutral; use charge_sectors=True " + "for a charged operator sum." + ) + + endpoint_nodes = tuple(plan.node_of_qubit[site] for site in support) + if len(support) == 1: + factors = { + endpoint_nodes[0]: qtn.Tensor( + term, + inds=(f"k{support[0]}", f"b{support[0]}"), + ) + } + active_nodes = {endpoint_nodes[0]} + else: + ordered_support = tuple(sorted(support)) + if support == ordered_support: + ordered_term = term + else: + rank = len(support) + order = tuple(sorted(range(rank), key=support.__getitem__)) + ordered_term = term.transpose( + (*order, *(axis + rank for axis in order)) + ) + rank = len(ordered_support) + fused = ordered_term.fuse(*( + (axis, axis + rank) for axis in range(rank) + )) + packed_inds = tuple(f"_pepsy_op_packed_{site}" for site in ordered_support) + blob = qtn.Tensor(fused, inds=packed_inds) + ordered_nodes = tuple(plan.node_of_qubit[site] for site in ordered_support) + active_nodes = {ordered_nodes[0]} + for target in ordered_nodes[1:]: + anchor = min( + active_nodes, + key=lambda node: len(plan.node_path(node, target)), + ) + active_nodes.update(plan.node_path(anchor, target)) + peel_order, hub = _tree_operator_peel_order(plan, active_nodes) + owned = {node: set() for node in active_nodes} + for node, site in zip(ordered_nodes, ordered_support): + owned[node].add(f"_pepsy_op_packed_{site}") + factors = {} + for node, neighbor in peel_order: + left_inds = tuple( + index for index in blob.inds if index in owned[node] + ) + if not left_inds: + raise RuntimeError( + f"operator decomposition lost subtree payload at {node}." + ) + left, right = blob.split( + left_inds=left_inds, + method="svd", + absorb="right", + cutoff=max(64.0 * np.finfo(float).eps, float(cutoff)), + get="tensors", + bond_ind=f"_pepsy_op_bond_{node}_{neighbor}", + ) + factors[node] = left + blob = right + owned[neighbor].add(f"_pepsy_op_bond_{node}_{neighbor}") + factors[hub] = blob + + def edge_name(node, neighbor): + return f"_pepsy_tnno_{min(node, neighbor)}_{max(node, neighbor)}" + + def rebuild_with_axis(data, maps, duals, dense): + return _native_from_dense( + dense, + symmetry=symmetry, + index_maps=maps, + duals=duals, + charge=getattr(data, "charge", zero), + ) + + tensors = [] + for node in plan.nodes(): + qubit = plan.qubit_of_node.get(node) + neighbors = _tree_plan_neighbors(plan, node) + if node in factors: + factor = factors[node] + data = factor.data + inds = list(factor.inds) + if qubit in support: + packed = next( + ( + index for index in inds + if index.startswith("_pepsy_op_packed_") + ), + None, + ) + if packed is not None: + axis = inds.index(packed) + data = data.unfuse(axis) + inds[axis:axis + 1] = [f"k{qubit}", f"b{qubit}"] + elif qubit is not None: + # A physical TreePlan root can lie on the active Steiner + # subtree without being an endpoint. Its operator action is + # the identity, so add that even physical pair explicitly. + dense = np.asarray(data.to_dense(), dtype=dtype or complex) + dense = np.einsum( + "ab,...->ab...", + np.eye(len(physical_map), dtype=dense.dtype), + dense, + ) + maps = [physical_map, physical_map] + [ + _expanded_index_charges(index) for index in data.indices + ] + duals = [False, True] + [ + index.dual for index in data.indices + ] + data = rebuild_with_axis(data, maps, duals, dense) + inds = [f"k{qubit}", f"b{qubit}"] + inds + + # Rename the native decomposition's temporary bond labels before + # adding the trivial exterior bonds. + inds = [ + edge_name(*map(int, index.removeprefix("_pepsy_op_bond_").split("_"))) + if index.startswith("_pepsy_op_bond_") else index + for index in inds + ] + existing = set(inds) + for neighbor in neighbors: + index = edge_name(node, neighbor) + if index in existing: + continue + dense = np.expand_dims( + np.asarray(data.to_dense(), dtype=dtype or complex), + axis=-1, + ) + maps = [ + _expanded_index_charges(axis) + for axis in data.indices + ] + [[zero]] + duals = [axis.dual for axis in data.indices] + [ + neighbor in plan.children[node] + ] + data = rebuild_with_axis(data, maps, duals, dense) + inds.append(index) + existing.add(index) + + desired = [ + *((f"k{qubit}", f"b{qubit}") if qubit is not None else ()), + *(edge_name(node, neighbor) for neighbor in neighbors), + ] + tensor = qtn.Tensor(data, inds=inds).transpose(*desired) + tensor.add_tag(f"N{node}") + if qubit is not None: + tensor.add_tag(f"I{qubit}") + tensors.append(tensor) + continue + + # Nodes outside the term's Steiner subtree carry a neutral identity. + maps = [] + duals = [] + inds = [] + if qubit is not None: + maps.extend((physical_map, physical_map)) + duals.extend((False, True)) + inds.extend((f"k{qubit}", f"b{qubit}")) + for neighbor in neighbors: + maps.append([zero]) + duals.append(neighbor in plan.children[node]) + inds.append(edge_name(node, neighbor)) + shape = tuple(len(index_map) for index_map in maps) + data = np.zeros(shape, dtype=dtype or complex) + if qubit is None: + data[(0,) * len(neighbors)] = 1.0 + else: + data[(slice(None), slice(None)) + (0,) * len(neighbors)] = np.eye( + len(physical_map), dtype=data.dtype + ) + tensors.append(qtn.Tensor( + _native_from_dense( + data, + symmetry=symmetry, + index_maps=maps, + duals=duals, + charge=zero, + ), + inds=inds, + tags=[f"N{node}"] + ([f"I{qubit}"] if qubit is not None else []), + )) + + network = qtn.TensorNetwork(tensors) + network.pepsy_tree_operator_kind = "native_tree_term_tnno" + network.pepsy_tree_operator_is_ttno = True + return network + + +def _normalize_native_term_edge_orientation(network, plan, *, symmetry, dtype=None): + """Normalize native term-network virtual duals to the TreePlan orientation.""" + if not hasattr(symmetry, "parity"): + from symmray import get_symmetry # pylint: disable=import-outside-toplevel + + symmetry = get_symmetry(symmetry) + for node in plan.nodes(): + tensor = network[f"N{node}"] + for neighbor in _tree_plan_neighbors(plan, node): + edge = f"_pepsy_tnno_{min(node, neighbor)}_{max(node, neighbor)}" + axis = tensor.inds.index(edge) + index = tensor.data.indices[axis] + desired_dual = neighbor in plan.children[node] + if index.dual == desired_dual: + continue + old_charges = _expanded_index_charges(index) + relabelled = [_operator_charge_neg(charge) for charge in old_charges] + zero = ( + tuple(0 for _ in relabelled[0]) + if relabelled and isinstance(relabelled[0], tuple) + else 0 + ) + probe = _native_from_dense( + np.zeros((len(relabelled),), dtype=dtype or complex), + symmetry=symmetry, + index_maps=[relabelled], + duals=[desired_dual], + charge=zero, + ) + new_charges = _expanded_index_charges(probe.indices[0]) + old_positions = {} + for position, charge in enumerate(old_charges): + old_positions.setdefault(charge, []).append(position) + used = {charge: 0 for charge in old_positions} + permutation = [] + for charge in new_charges: + old_charge = _operator_charge_neg(charge) + position = used[old_charge] + permutation.append(old_positions[old_charge][position]) + used[old_charge] = position + 1 + dense = np.take( + np.asarray(tensor.data.to_dense(), dtype=dtype or complex), + permutation, + axis=axis, + ) + # Reversing a fermionic virtual edge is a graded dualization, not + # just a charge-label permutation. The plan-parent endpoint + # carries the parity gauge associated with that reversal. Apply + # it once per edge (the child endpoint gets the dual charge map, + # but not a second parity phase). + if desired_dual: + parity = np.asarray( + [ + -1 if symmetry.parity(charge) else 1 + for charge in new_charges + ], + dtype=dense.dtype, + ) + phase_shape = [1] * dense.ndim + phase_shape[axis] = len(parity) + dense = dense * parity.reshape(phase_shape) + maps = [ + new_charges if current_axis == axis else + _expanded_index_charges(current) + for current_axis, current in enumerate(tensor.data.indices) + ] + duals = [ + desired_dual if current_axis == axis else current.dual + for current_axis, current in enumerate(tensor.data.indices) + ] + rebuilt = _native_from_dense( + dense, + symmetry=symmetry, + index_maps=maps, + duals=duals, + charge=getattr(tensor.data, "charge", 0), + ) + tensor.modify(data=rebuilt) + return network + + +def _native_term_sum_tree_operator( + plan, terms, *, symmetry, cutoff=1e-12, dtype=None, +): + """Direct-sum exact native term TTNOs into one operator network.""" + import quimb.tensor as qtn # pylint: disable=import-outside-toplevel + + term_networks = [] + for where, term in terms.items(): + network = _native_tree_term_network( + plan, + term, + _term_support(where), + symmetry=symmetry, + cutoff=cutoff, + dtype=dtype, + ) + term_networks.append(_normalize_native_term_edge_orientation( + network, plan, symmetry=symmetry, dtype=dtype, + )) + term_networks = tuple(term_networks) + if not term_networks: + raise ValueError("At least one operator term is required.") + + def edge_name(node, neighbor): + return f"_pepsy_tnno_{min(node, neighbor)}_{max(node, neighbor)}" + + physical_map = None + for term in terms.values(): + physical_map = _expanded_index_charges(term.indices[0]) + break + zero = ( + tuple(0 for _ in physical_map[0]) + if physical_map and isinstance(physical_map[0], tuple) + else 0 + ) + edge_maps = {} + edge_positions = {index: {} for index in range(len(term_networks))} + for node in plan.nodes(): + for neighbor in plan.children[node]: + edge = (node, neighbor) + edge_index = edge_name(*edge) + charges = [] + for index, network in enumerate(term_networks): + tensor = network[f"N{neighbor}"] + local = tensor.data.indices[tensor.inds.index(edge_index)] + local_charges = _expanded_index_charges(local) + charges.extend(local_charges) + # Symmray groups an index's sectors by charge, so a raw + # concatenation of per-term charge lists is not a set of + # contiguous direct-sum slices when two terms share a charge. + # Build the actual expanded order through the same native index + # constructor and allocate duplicate charge sectors term by term. + probe = _native_from_dense( + np.zeros((len(charges),), dtype=dtype or complex), + symmetry=symmetry, + index_maps=[charges], + duals=[False], + charge=zero, + ) + global_charges = _expanded_index_charges(probe.indices[0]) + edge_maps[edge_index] = global_charges + positions_by_charge = {} + for position, charge in enumerate(global_charges): + positions_by_charge.setdefault(charge, []).append(position) + used_by_charge = {charge: 0 for charge in positions_by_charge} + for index, network in enumerate(term_networks): + tensor = network[f"N{neighbor}"] + local = tensor.data.indices[tensor.inds.index(edge_index)] + local_charges = _expanded_index_charges(local) + positions = [] + for charge in local_charges: + offset = used_by_charge[charge] + positions.append(positions_by_charge[charge][offset]) + used_by_charge[charge] = offset + 1 + edge_positions[index][edge_index] = tuple(positions) + + tensors = [] + for node in plan.nodes(): + neighbors = _tree_plan_neighbors(plan, node) + qubit = plan.qubit_of_node.get(node) + desired = [ + *((f"k{qubit}", f"b{qubit}") if qubit is not None else ()), + *(edge_name(node, neighbor) for neighbor in neighbors), + ] + global_maps = [] + global_duals = [] + if qubit is not None: + global_maps.extend((physical_map, physical_map)) + global_duals.extend((False, True)) + for neighbor in neighbors: + index = edge_name(node, neighbor) + global_maps.append(edge_maps[index]) + # Native tree decomposition can orient an odd operator bond + # differently at a hub than the ordinary state-tree convention. + # The two endpoint dual flags are part of the graded operator + # data; recomputing them from the plan would change valid local + # sectors and drop them during ``from_dense``. + reference = term_networks[0][f"N{node}"] + reference_axis = reference.data.indices[ + reference.inds.index(index) + ] + global_duals.append(reference_axis.dual) + shape = tuple(len(index_map) for index_map in global_maps) + data = np.zeros(shape, dtype=dtype or complex) + for term_index, network in enumerate(term_networks): + tensor = network[f"N{node}"].transpose(*desired) + local = np.asarray(tensor.data.to_dense(), dtype=data.dtype) + slices = [] + if qubit is not None: + slices.extend((slice(None), slice(None))) + for neighbor in neighbors: + slices.append(edge_positions[term_index][edge_name(node, neighbor)]) + # ``np.ix_`` is needed here because charge-grouped duplicate + # sectors are generally interleaved across the direct-sum axis. + data[np.ix_(*[ + np.arange(local.shape[axis]) if isinstance(selection, slice) + else np.asarray(selection) + for axis, selection in enumerate(slices) + ])] += local + tensors.append(qtn.Tensor( + _native_from_dense( + data, + symmetry=symmetry, + index_maps=global_maps, + duals=global_duals, + charge=zero, + ), + inds=desired, + tags=[f"N{node}"] + ([f"I{qubit}"] if qubit is not None else []), + )) + + network = qtn.TensorNetwork(tensors) + network.pepsy_tree_operator_kind = "native_tree_tnno" + network.pepsy_tree_operator_bond = max( + (network.ind_size(index) for index in network.inner_inds()), + default=1, + ) + network.pepsy_tree_operator_raw_bond = network.pepsy_tree_operator_bond + network.pepsy_tree_operator_is_ttno = True + return network + + +def _tree_operator_tensor(network, node): + """Fetch one operator tensor by its stable TreePlan node tag.""" + return network[f"N{node}"] + + +def _tree_operator_bond(network, plan, node, neighbor): + """Find the unique live operator bond for one TreePlan edge.""" + left = _tree_operator_tensor(network, node) + right = _tree_operator_tensor(network, neighbor) + shared = tuple(set(left.inds).intersection(right.inds)) + if len(shared) != 1: + raise ValueError( + f"operator TTNO edge {(node, neighbor)!r} has {len(shared)} bonds." + ) + return shared[0] + + +def _tree_operator_qr(tensor, *, left_inds, bond_ind): + """Run the lossless dense/native QR policy for one operator tensor.""" + import quimb.tensor as qtn # pylint: disable=import-outside-toplevel + from .ttn import _native_qr_split_tensor # pylint: disable=import-outside-toplevel + + options = { + "left_inds": tuple(left_inds), + "right_inds": (bond_ind,), + "method": "qr", + "absorb": "right", + "cutoff": 0.0, + "get": "tensors", + } + options["bond_ind"] = qtn.rand_uuid() + return _native_qr_split_tensor(tensor, **options) + + +def _canonicalize_tree_operator(network, plan, center): + """Canonicalize a tree operator by lossless QR from leaves to center.""" + import quimb.tensor as qtn # pylint: disable=import-outside-toplevel + + if center not in plan.children: + raise ValueError(f"operator canonicalization center {center!r} is invalid.") + order = sorted( + (node for node in plan.nodes() if node != center), + key=lambda node: len(plan.node_path(node, center)), + reverse=True, + ) + for node in order: + neighbor = plan.node_path(node, center)[1] + tensor = _tree_operator_tensor(network, node) + target = _tree_operator_tensor(network, neighbor) + bond = _tree_operator_bond(network, plan, node, neighbor) + kept, carry = _tree_operator_qr( + tensor, + left_inds=tuple(index for index in tensor.inds if index != bond), + bond_ind=bond, + ) + merged = qtn.tensor_contract(carry, target) + tensor.modify( + data=kept.data, + inds=kept.inds, + left_inds=kept.left_inds, + ) + target.modify( + data=merged.data, + inds=merged.inds, + left_inds=None, + ) + network.pepsy_tree_operator_center = int(center) + network.pepsy_tree_operator_canonical = True + return network + + +def _compress_tree_operator(network, plan, *, max_bond, cutoff): + """Compress one combined TTNO edge-by-edge with native graded SVD.""" + import quimb.tensor as qtn # pylint: disable=import-outside-toplevel + + raw_bond = max( + (network.ind_size(index) for index in network.inner_inds()), + default=1, + ) + if cutoff == 0.0 and max_bond is None: + _canonicalize_tree_operator(network, plan, plan.root) + final_bond = raw_bond + return { + "compressed": False, + "cutoff": cutoff, + "requested_max_bond": None, + "raw_max_bond": raw_bond, + "final_max_bond": final_bond, + "rank_reduced": False, + } + + order = sorted( + (node for node in plan.nodes() if node != plan.root), + key=lambda node: len(plan.node_path(node, plan.root)), + reverse=True, + ) + for node in order: + neighbor = plan.node_path(node, plan.root)[1] + tensor = _tree_operator_tensor(network, node) + target = _tree_operator_tensor(network, neighbor) + bond = _tree_operator_bond(network, plan, node, neighbor) + combined = qtn.tensor_contract(tensor, target) + left_inds = tuple(index for index in tensor.inds if index != bond) + right_inds = tuple(index for index in target.inds if index != bond) + options = { + "left_inds": left_inds, + "right_inds": right_inds, + "method": "svd", + "absorb": "right", + "cutoff": cutoff, + "cutoff_mode": "rsum2", + "get": "tensors", + "bond_ind": bond, + } + if max_bond is not None: + options["max_bond"] = int(max_bond) + left, right = combined.split(**options) + tensor.modify( + data=left.data, + inds=left.inds, + left_inds=left.left_inds, + ) + target.modify( + data=right.data, + inds=right.inds, + left_inds=right.left_inds, + ) + final_bond = max( + (network.ind_size(index) for index in network.inner_inds()), + default=1, + ) + network.pepsy_tree_operator_bond = final_bond + network.pepsy_tree_operator_canonical = False + return { + "compressed": True, + "cutoff": cutoff, + "requested_max_bond": None if max_bond is None else int(max_bond), + "raw_max_bond": raw_bond, + "final_max_bond": final_bond, + "rank_reduced": final_bond < raw_bond, + "max_bond_exceeded": ( + max_bond is not None and final_bond > int(max_bond) + ), + } + + +def _relocate_mpo(mpo, order, *, upper_ind_id, lower_ind_id): + """Relabel chain positions on ``mpo`` back to logical qubit labels.""" + tag_map = { + mpo.site_tag(position): mpo.site_tag(qubit) + for position, qubit in enumerate(order) + } + index_map = { + upper_ind_id.format(position): upper_ind_id.format(qubit) + for position, qubit in enumerate(order) + } + index_map.update({ + lower_ind_id.format(position): lower_ind_id.format(qubit) + for position, qubit in enumerate(order) + }) + mpo.retag_(tag_map) + mpo.reindex_(index_map) + mpo.pepsy_tree_order = tuple(order) + mpo.pepsy_tree_native = any( + type(tensor.data).__name__.endswith("FermionicArray") + for tensor in mpo + ) + return mpo + + +def _native_from_dense( + data, *, symmetry, index_maps, duals, charge, label=None, +): + """Create one native Symmray tensor lazily.""" + from symmray import utils as sr_utils # pylint: disable=import-outside-toplevel + + return sr_utils.from_dense( + data, + symmetry=symmetry, + index_maps=index_maps, + duals=duals, + fermionic=True, + charge=charge, + label=label, + ) + + +def _pair_coefficient_factors(terms, nsite): + """Factor an off-diagonal symmetric coefficient table, if possible.""" + first_support, first_term = next(iter(terms.items())) + first_matrix = np.asarray(first_term.to_dense()).reshape( + (first_term.shape[0] * first_term.shape[2],) * 2 + ) + table = np.zeros((nsite, nsite), dtype=complex) + for where, term in terms.items(): + support = _term_support(where) + if len(support) != 2 or support[0] >= support[1]: + return None + matrix = np.asarray(term.to_dense()).reshape( + (term.shape[0] * term.shape[2],) * 2 + ) + denominator = np.vdot(first_matrix, first_matrix) + ratio = np.vdot(first_matrix, matrix) / denominator + if not np.allclose(matrix, ratio * first_matrix, rtol=1e-10, atol=1e-12): + return None + table[support] = ratio / 2.0 + table[support[::-1]] = ratio / 2.0 + + nonzero = np.argwhere(np.abs(table) > 1e-14) + if len(nonzero) < 3: + return None + i0, j0 = map(int, nonzero[0]) + candidates = [ + index for index in range(nsite) + if index not in {i0, j0} and abs(table[i0, index]) > 1e-14 + ] + if not candidates or abs(table[i0, j0]) <= 1e-14: + return None + k = candidates[0] + a = np.zeros(nsite, dtype=complex) + b = np.zeros(nsite, dtype=complex) + a[i0] = 1.0 + b[j0] = table[i0, j0] + b[k] = table[i0, k] + a[k] = table[k, j0] / b[j0] + if abs(a[k]) <= 1e-14 or abs(b[k]) <= 1e-14: + return None + a[j0] = table[j0, k] / b[k] + b[i0] = table[k, i0] / a[k] + for index in range(nsite): + if index != j0 and abs(b[j0]) > 1e-14: + a[index] = table[index, j0] / b[j0] + if index != i0: + b[index] = table[i0, index] / a[i0] + + for i in range(nsite): + for j in range(nsite): + if i == j: + continue + if not np.allclose( + a[i] * b[j], table[i, j], rtol=1e-9, atol=1e-11, + ): + return None + return first_term, first_support, a, b + + +def _pair_endpoint_automaton( + plan, terms, *, symmetry, cutoff=1e-12, dtype=None, +): + """Compile a separable pair correlator into a four-state tree operator.""" + import quimb.tensor as qtn # pylint: disable=import-outside-toplevel + + factored = _pair_coefficient_factors(terms, plan.n) + if factored is None: + return None + first_term, first_support, source_weights, target_weights = factored + fused = first_term.fuse((0, 2), (1, 3)) + left, _, right = fused.svd(absorb="right", cutoff=cutoff) + if left is None or right is None: + return None + left = left.unfuse(0).transpose((2, 0, 1)) + right = right.unfuse(1) + left_data = np.asarray(left.to_dense(), dtype=dtype or complex)[0] + right_data = np.asarray(right.to_dense(), dtype=dtype or complex)[0] + physical_map = _expanded_index_charges(left.indices[1]) + physical_dim = len(physical_map) + zero = 0 if symmetry in {"U1", "Z2"} else (0, 0) + pair_charge = getattr(first_term, "pair_charge", None) + if pair_charge is None: + # The local factor's first physical charge is sufficient to infer the + # endpoint channel charge for the standard pair observable. + pair_charge = physical_map[-1] + opposite_pair = ( + tuple(-value for value in pair_charge) + if isinstance(pair_charge, tuple) else -pair_charge + ) + state_map = [zero, opposite_pair, pair_charge, zero] + tensors = [] + + for node in plan.nodes(): + children = tuple(plan.children[node]) + parent = plan.parent.get(node) + qubit = plan.qubit_of_node.get(node) + has_parent = parent is not None + edges = list(children) + ([parent] if has_parent else []) + shape = [4] * len(edges) + maps = [state_map] * len(edges) + duals = [True] * len(children) + ([False] if has_parent else []) + inds = [ + f"_to{min(node, neighbor)}_{max(node, neighbor)}" + for neighbor in edges + ] + if qubit is not None: + shape.extend((physical_dim, physical_dim)) + maps.extend((physical_map, physical_map)) + duals.extend((False, True)) + inds.extend((f"k{qubit}", f"b{qubit}")) + data = np.zeros(shape, dtype=dtype or complex) + + if qubit is not None: + source = source_weights[qubit] * left_data + target = target_weights[qubit] * right_data + identity = np.eye(physical_dim, dtype=data.dtype) + + for child_states in ( + np.ndindex(*(4 for _ in children)) if children else [()] + ): + source_count = sum(state & 1 for state in child_states) + target_count = sum((state >> 1) & 1 for state in child_states) + if source_count > 1 or target_count > 1: + continue + base = source_count | (target_count << 1) + options = [(base, identity if qubit is not None else 1.0)] + if qubit is not None: + if not source_count: + options.append((1 | (target_count << 1), source)) + if not target_count: + options.append((source_count | 2, target)) + for output, local_operator in options: + if has_parent: + index = child_states + (output,) + else: + if output != 3: + continue + index = child_states + data[index] += local_operator + + array = _native_from_dense( + data, + symmetry=symmetry, + index_maps=maps, + duals=duals, + charge=zero, + ) + tags = [f"N{node}"] + if qubit is not None: + tags.append(f"I{qubit}") + tensors.append(qtn.Tensor(array, inds=inds, tags=tags)) + + network = qtn.TensorNetwork(tensors) + network.pepsy_tree_operator_kind = "pair_endpoint_automaton" + network.pepsy_tree_operator_bond = 4 + return network + + +def _pair_chain_mpo( + terms, *, symmetry, nsite, cutoff=1e-12, dtype=None, + upper_ind_id="k{}", lower_ind_id="b{}", site_tag_id="I{}", +): + """Build the compact native chain MPO for a symmetric pair table. + + The two active bond sectors represent ``pair_create`` before + ``pair_annihilate`` and the reversed ordering. The two neutral sectors + are the open and closed boundaries, so the dense bond dimension is four + (and the native charge blocks remain explicit). This is the chain + counterpart of :func:`_pair_endpoint_automaton`. + """ + import quimb.tensor as qtn # pylint: disable=import-outside-toplevel + + factored = _pair_coefficient_factors(terms, nsite) + if factored is None: + return None + first_term, _, source_weights, target_weights = factored + fused = first_term.fuse((0, 2), (1, 3)) + left, _, right = fused.svd(absorb="right", cutoff=cutoff) + if left is None or right is None or left.shape[1] != 1: + return None + left = left.unfuse(0).transpose((2, 0, 1)) + right = right.unfuse(1) + left_data = np.asarray(left.to_dense(), dtype=dtype or complex)[0] + right_data = np.asarray(right.to_dense(), dtype=dtype or complex)[0] + physical_map = _expanded_index_charges(left.indices[1]) + physical_dim = len(physical_map) + zero = 0 if symmetry in {"U1", "Z2"} else (0, 0) + pair_charge = physical_map[-1] + opposite_pair = ( + tuple(-value for value in pair_charge) + if isinstance(pair_charge, tuple) else -pair_charge + ) + state_map = [zero, opposite_pair, pair_charge, zero] + identity = np.eye(physical_dim, dtype=dtype or complex) + + def make_array(data, maps, duals): + return _native_from_dense( + data, + symmetry=symmetry, + index_maps=maps, + duals=duals, + charge=zero, + ) + + arrays = [] + for position in range(nsite): + source = source_weights[position] * left_data + target = target_weights[position] * right_data + if position == 0: + data = np.zeros((4, physical_dim, physical_dim), dtype=identity.dtype) + data[0] = identity + data[1] = source + data[2] = target + arrays.append(make_array( + data, + [state_map, physical_map, physical_map], + [False, False, True], + )) + continue + if position == nsite - 1: + data = np.zeros((4, physical_dim, physical_dim), dtype=identity.dtype) + data[3] = identity + # Active sectors close into the neutral done boundary. + data[1] = target + data[2] = source + arrays.append(make_array( + data, + [state_map, physical_map, physical_map], + [True, False, True], + )) + continue + + data = np.zeros( + (4, 4, physical_dim, physical_dim), dtype=identity.dtype, + ) + for state in range(4): + data[state, state] = identity + data[0, 1] = source + data[0, 2] = target + data[1, 3] = target + data[2, 3] = source + arrays.append(make_array( + data, + [state_map, state_map, physical_map, physical_map], + [True, False, False, True], + )) + + return qtn.MatrixProductOperator( + arrays, + sites=range(nsite), + L=nsite, + shape="lrud", + upper_ind_id=upper_ind_id, + lower_ind_id=lower_ind_id, + site_tag_id=site_tag_id, + ) + + +def _tree_tensor_network_for_term( + plan, term, support, *, symmetry, cutoff=1e-12, dtype=None, +): + """Build an exact native fallback network for one higher-rank term. + + Higher-rank terms can still be kept as complete graded operator tensors + for callers that explicitly need this compatibility fallback. The normal + Hamiltonian path uses :func:`_native_tree_term_network` and amalgamates + the resulting factors into one canonicalizable TTNO rather than a list of + hyperedges. + """ + import quimb.tensor as qtn # pylint: disable=import-outside-toplevel + + support = tuple(int(site) for site in support) + if not support: + raise ValueError("native tree term support cannot be empty.") + if any(site not in plan.node_of_qubit for site in support): + raise ValueError(f"term support {support!r} is outside the TreePlan.") + expected_rank = 2 * len(support) + if len(term.indices) != expected_rank: + raise TypeError( + f"a {len(support)}-site native term must have rank " + f"{expected_rank}, got {len(term.indices)}." + ) + + physical_map = _expanded_index_charges(term.indices[0]) + physical_dim = len(physical_map) + zero = getattr(term, "zero_charge", None) + if zero is None: + zero = 0 if symmetry in {"U1", "Z2"} else (0, 0) + operator_dtype = np.dtype(dtype or np.asarray(term.to_dense()).dtype) + + # The term's native indices are ordered as all upper physical legs, + # followed by all lower physical legs. Keep that ordering intact while + # assigning the logical tree-site labels. + term_inds = [f"k{site}" for site in support] + term_inds.extend(f"b{site}" for site in support) + term_tags = [f"N{plan.node_of_qubit[site]}" for site in support] + term_tags.extend(f"I{site}" for site in support) + tensors = [qtn.Tensor(term, inds=term_inds, tags=term_tags)] + + support_set = set(support) + for site in sorted(plan.node_of_qubit): + if site in support_set: + continue + identity = _native_from_dense( + np.eye(physical_dim, dtype=operator_dtype), + symmetry=symmetry, + index_maps=[physical_map, physical_map], + duals=[False, True], + charge=zero, + ) + node = plan.node_of_qubit[site] + tensors.append(qtn.Tensor( + identity, + inds=[f"k{site}", f"b{site}"], + tags=[f"N{node}", f"I{site}"], + )) + + network = qtn.TensorNetwork(tensors) + network.pepsy_tree_operator_kind = "native_term_hyperedge" + span = {plan.node_of_qubit[support[0]]} + for site in support[1:]: + anchor = next(iter(span)) + span.update(plan.node_path(anchor, plan.node_of_qubit[site])) + network.pepsy_tree_operator_path = tuple(sorted(span)) + return network + + +def _dense_operator_array(operator, *, dtype=None): + """Extract one ordinary dense operator array.""" + if hasattr(operator, "to_dense"): + operator = operator.to_dense() + elif hasattr(operator, "data"): + operator = operator.data + return np.asarray(operator, dtype=dtype) + + +def _dense_tree_tensor_network_for_term(plan, operator, support, *, dtype=None): + """Build one exact ordinary dense tree operator hyperedge.""" + import quimb.tensor as qtn # pylint: disable=import-outside-toplevel + + support = tuple(int(site) for site in support) + data = _dense_operator_array(operator, dtype=dtype) + expected_rank = 2 * len(support) + if data.ndim != expected_rank: + raise ValueError( + f"a {len(support)}-site dense term must have rank {expected_rank}, " + f"got {data.ndim}." + ) + if any(site not in plan.node_of_qubit for site in support): + raise ValueError(f"term support {support!r} is outside the TreePlan.") + physical_dim = data.shape[0] + if any(size != physical_dim for size in data.shape): + raise ValueError("dense tree terms must have one physical dimension.") + tensors = [qtn.Tensor( + data, + inds=[f"k{site}" for site in support] + + [f"b{site}" for site in support], + tags=[f"N{plan.node_of_qubit[site]}" for site in support] + + [f"I{site}" for site in support], + )] + support_set = set(support) + identity = np.eye(physical_dim, dtype=data.dtype) + for site in sorted(plan.node_of_qubit): + if site in support_set: + continue + node = plan.node_of_qubit[site] + tensors.append(qtn.Tensor( + identity, + inds=[f"k{site}", f"b{site}"], + tags=[f"N{node}", f"I{site}"], + )) + network = qtn.TensorNetwork(tensors) + network.pepsy_tree_operator_kind = "dense_term_hyperedge" + span = {plan.node_of_qubit[support[0]]} + for site in support[1:]: + anchor = next(iter(span)) + span.update(plan.node_path(anchor, plan.node_of_qubit[site])) + network.pepsy_tree_operator_path = tuple(sorted(span)) + return network + + +def _build_tree_operator( + plan, + hamiltonian, + *, + cutoff=1e-12, + max_bond=None, + compress=True, + dtype=None, + fermionic=True, +): + """Build the backend-specific tree representation used by ``TreeMPO``.""" + symmetry = hamiltonian.symmetry + terms = hamiltonian.terms + + # The full staggered eta correlator is a symmetric rank-one pair table. + # Compile it before falling back to one actual tree contraction per term; + # this keeps p_eta_stag2 at a four-state tree bond for arbitrary N. + if fermionic: + # The factorization helper assumes every term is a two-site operator. + # In particular, an onsite-only Hamiltonian is a valid generic TTNO + # input but must go directly to the combined automaton. + is_pair_table = ( + len(terms) >= 3 + and all(len(_term_support(where)) == 2 for where in terms) + ) + pair_network = _pair_endpoint_automaton( + plan, terms, symmetry=symmetry, cutoff=cutoff, dtype=dtype, + ) if is_pair_table else None + if pair_network is not None: + return pair_network + return _native_term_sum_tree_operator( + plan, + terms, + symmetry=symmetry, + cutoff=cutoff, + dtype=dtype, + ) + + return _combined_tree_operator( + plan, + terms, + symmetry=symmetry, + cutoff=cutoff, + dtype=dtype, + fermionic=False, + ) + + +def _annotate_tree_mpo( + mpo, + plan, + terms, + tree_operator, + *, + symmetry=None, + compressed=False, + cutoff=1e-12, + max_bond=None, +): + """Attach a public :class:`TreeMPO` to a compatibility chain MPO.""" + mpo.pepsy_tree_plan_signature = _tree_plan_signature(plan) + mpo.pepsy_tree_terms = dict(terms) + if isinstance(tree_operator, TreeMPO): + operator = tree_operator + else: + networks = ( + tuple(tree_operator) + if isinstance(tree_operator, (tuple, list)) + else (tree_operator,) + ) + native = any( + type(tensor.data).__name__.endswith("FermionicArray") + for network in networks + for tensor in network + ) + operator = TreeMPO( + plan, + networks, + chain_mpo=mpo, + terms=terms, + backend="symmray" if native else "dense", + fermionic=native, + symmetry=symmetry, + compressed=compressed, + ) + mpo.pepsy_tree_operator = operator + mpo.pepsy_tree_operator_networks = operator.tree_networks + if compressed: + operator.compress(max_bond=max_bond, cutoff=cutoff) + return mpo + + +def tree_mpo( + plan, + hamiltonian, + *, + max_bond=None, + cutoff=1e-12, + compress=True, + upper_ind_id="k{}", + lower_ind_id="b{}", + site_tag_id="I{}", + dtype=None, + fermionic=True, + charge_sectors=False, + to_backend=None, +): + """Build a low-bond native chain MPO plus its TreePlan embedding. + + The returned value is the ordinary chain MPO produced by + :meth:`SymHamiltonian.to_mpo`, with logical site labels restored after the + selected ``TreePlan.mpo_order`` construction. On a native tree state, + :meth:`TreeTensorNetwork.expectation_mpo_exact` uses the attached native + tree embedding and contracts ``tree.H | tree_operator | tree``. The chain + MPO is never moved into the tree or compressed as part of that readout. + """ + if not isinstance(plan, TreePlan): + raise TypeError("plan must be a TreePlan.") + + from ...tensors.symmetric import SymHamiltonian + + if not isinstance(hamiltonian, SymHamiltonian): + raise TypeError("hamiltonian must be a SymHamiltonian instance.") + if not fermionic: + warnings.warn( + "tree_mpo(..., fermionic=False) selects the explicit " + "Jordan--Wigner compatibility path. For a native fermionic " + "TreeTensorNetwork, keep fermionic=True so graded Symmray " + "blocks and signs are preserved.", + UserWarning, + stacklevel=2, + ) + + order = plan.mpo_order() + if len(order) != plan.n or set(order) != set(range(plan.n)): + raise ValueError( + "TreePlan.mpo_order() must contain every logical qubit exactly once." + ) + position = {qubit: index for index, qubit in enumerate(order)} + mapped_terms = {} + for where, term in hamiltonian.terms.items(): + support = _term_support(where) + try: + mapped_support = tuple(position[qubit] for qubit in support) + except KeyError as exc: + raise ValueError( + f"Hamiltonian term support {support!r} is outside the tree " + f"qubits 0 .. {plan.n - 1}." + ) from exc + mapped_terms[mapped_support] = term + + if not mapped_terms: + raise ValueError("At least one Hamiltonian term is required.") + mapped_hamiltonian = SymHamiltonian.from_terms( + hamiltonian.model, + hamiltonian.symmetry, + mapped_terms, + parameters=hamiltonian.parameters, + ) + compact = None + if fermionic and not charge_sectors: + compact = _pair_chain_mpo( + mapped_terms, + symmetry=hamiltonian.symmetry, + nsite=plan.n, + cutoff=cutoff, + dtype=dtype, + upper_ind_id=upper_ind_id, + lower_ind_id=lower_ind_id, + site_tag_id=site_tag_id, + ) + if compact is not None: + built = compact + else: + built = mapped_hamiltonian.to_mpo( + L=plan.n, + max_bond=max_bond, + cutoff=cutoff, + compress=compress, + upper_ind_id=upper_ind_id, + lower_ind_id=lower_ind_id, + site_tag_id=site_tag_id, + dtype=dtype, + fermionic=fermionic, + charge_sectors=charge_sectors, + to_backend=to_backend, + ) + if charge_sectors: + terms_by_charge = {} + for where, term in hamiltonian.terms.items(): + terms_by_charge.setdefault(getattr(term, "charge", 0), {})[ + where + ] = term + result = {} + for charge, mpo in built.items(): + sector_terms = terms_by_charge.get(charge, {}) + # A charge-sector MPO must carry the matching native tree + # embedding. Reusing the full Hamiltonian embedding here would + # silently add terms from the other returned sectors. + sector_hamiltonian = SymHamiltonian.from_terms( + hamiltonian.model, + hamiltonian.symmetry, + sector_terms, + parameters=hamiltonian.parameters, + ) if sector_terms else None + tree_operator = ( + _build_tree_operator( + plan, + sector_hamiltonian, + cutoff=cutoff, + max_bond=max_bond, + compress=compress, + dtype=dtype, + fermionic=fermionic, + ) + if sector_hamiltonian is not None else [] + ) + result[charge] = _annotate_tree_mpo( + _relocate_mpo( + mpo, + order, + upper_ind_id=upper_ind_id, + lower_ind_id=lower_ind_id, + ), + plan, + sector_terms, + tree_operator, + symmetry=hamiltonian.symmetry, + compressed=compress, + cutoff=cutoff, + max_bond=max_bond, + ) + return result + tree_operator = _build_tree_operator( + plan, + hamiltonian, + cutoff=cutoff, + max_bond=max_bond, + compress=compress, + dtype=dtype, + fermionic=fermionic, + ) + return _annotate_tree_mpo( + _relocate_mpo( + built, + order, + upper_ind_id=upper_ind_id, + lower_ind_id=lower_ind_id, + ), + plan, + hamiltonian.terms, + tree_operator, + symmetry=hamiltonian.symmetry, + compressed=compress, + cutoff=cutoff, + max_bond=max_bond, + ) diff --git a/src/pepsy/optimizers/tree/optimizer.py b/src/pepsy/optimizers/tree/optimizer.py index 2cf22e4..59c2a18 100644 --- a/src/pepsy/optimizers/tree/optimizer.py +++ b/src/pepsy/optimizers/tree/optimizer.py @@ -40,6 +40,7 @@ import contextlib from copy import deepcopy import heapq +import inspect from numbers import Integral import time import warnings @@ -161,6 +162,25 @@ def _is_symmray_array(array): return False +def _submpo_is_native(submpo): + """Return whether an MPO visibly contains native Symmray tensors. + + ``tree_mpo`` records this explicitly, while ordinary Quimb MPOs are + inspected as a fallback. ``None`` means that the payload does not expose + enough information to classify it without materialising it. + """ + marker = getattr(submpo, "pepsy_tree_native", None) + if marker is not None: + return bool(marker) + tensors = getattr(submpo, "tensors", None) + if tensors is None: + return None + try: + return any(_is_symmray_array(tensor.data) for tensor in tensors) + except (AttributeError, TypeError): + return None + + def _array_backend_signature(array): """Return comparable backend / dtype / device metadata for an array.""" return infer_backend_signature(array) @@ -3759,6 +3779,35 @@ def _compress_edge_with_diagnostics( "seconds": time.perf_counter() - profile_started, }) + def _compress_edge_compat( + self, u, v, *, max_bond=None, cutoff=None, reduced=True, + reduction_proven=False, + ): + """Call the compression hook while supporting older overrides. + + Tree stabilizer and diagnostic integrations can wrap the private + compression hook. Keep wrappers with the older signature working + while the built-in hook receives the proof flag. + """ + method = self._compress_edge_with_diagnostics + try: + parameters = inspect.signature(method).parameters.values() + supports_proof = any( + parameter.name == "reduction_proven" + or parameter.kind is inspect.Parameter.VAR_KEYWORD + for parameter in parameters + ) + except (TypeError, ValueError): + supports_proof = True + kwargs = { + "max_bond": max_bond, + "cutoff": cutoff, + "reduced": reduced, + } + if supports_proof: + kwargs["reduction_proven"] = reduction_proven + return method(u, v, **kwargs) + def _metadata_aware_reduction(self, u, v): """Choose one-sided compression when ``v`` is proven isometric. @@ -3824,7 +3873,7 @@ def _compress_path( cutoff=cutoff, ) reduced, reduction_proven = "left", True - self._compress_edge_with_diagnostics( + self._compress_edge_compat( v, u, max_bond=max_bond, cutoff=edge_cutoff, reduced=reduced, reduction_proven=reduction_proven, ) @@ -3870,7 +3919,7 @@ def descend(node, parent): reduced, reduction_proven = self._edge_reduction( node, child, max_bond=max_bond, cutoff=child_cutoff, ) - self._compress_edge_with_diagnostics( + self._compress_edge_compat( node, child, max_bond=max_bond, @@ -4090,6 +4139,20 @@ def apply_submpo(self, submpo, where, *, max_bond=None, cutoff=None): self._invalidate_state_norm_cache() logical_where = _normalize_where(where) where = self._validate_support(logical_where) + payload_native = _submpo_is_native(submpo) + state_native = bool(getattr(self.tn, "fermionic", False)) + if payload_native is not None and payload_native != state_native: + if state_native: + raise TypeError( + "native fermionic TreeTensorNetwork requires a native " + "Symmray MPO. Build it with tree_mpo(..., fermionic=True) " + "or supply a model-native MPO." + ) + raise TypeError( + "a native Symmray MPO cannot be applied to an ordinary dense " + "TreeTensorNetwork. Use fermionic=False for the explicit " + "Jordan--Wigner compatibility MPO." + ) return self._apply_submpo_resolved( submpo, where, logical_where=logical_where, max_bond=max_bond, cutoff=cutoff, @@ -4097,7 +4160,8 @@ def apply_submpo(self, submpo, where, *, max_bond=None, cutoff=None): def expectation_mpo( self, submpo, where, *, max_bond=None, cutoff=0.0, - normalized=True, optimize="auto", + normalized=True, optimize="auto", warn_on_truncation=True, + return_diagnostics=False, ): """Evaluate ```` through one structured tree-MPO pass. @@ -4107,16 +4171,62 @@ def expectation_mpo( defaults to this optimizer's ``chi``; pass a larger cap when the operator application must retain more of the exact MPO-transformed state. ``cutoff=0.0`` is the default because this is a measurement, - not a variational update. + not a variational update. A finite ``max_bond`` can still truncate if + the transformed ket exceeds that cap. Such truncation emits a + ``UserWarning`` by default; set ``warn_on_truncation=False`` only when + that approximation is intentional. Set ``return_diagnostics=True`` to + receive ``(value, diagnostics)`` with the compression events from this + expectation only. """ + if not isinstance(warn_on_truncation, bool): + raise TypeError("warn_on_truncation must be a bool.") + if not isinstance(return_diagnostics, bool): + raise TypeError("return_diagnostics must be a bool.") + logical_where = _normalize_where(where) + effective_max_bond = ( + self.chi + if max_bond is None + else self._normalize_max_bond(max_bond) + ) + effective_cutoff = ( + self.cutoff if cutoff is None else float(cutoff) + ) + if effective_cutoff < 0.0: + raise ValueError("cutoff must be non-negative.") + plan_order = getattr(submpo, "pepsy_tree_order", None) + if plan_order is not None and tuple(plan_order) != self.plan.mpo_order(): + warnings.warn( + "the structured MPO was built for tree MPO order " + f"{tuple(plan_order)!r}, while this state uses " + f"{self.plan.mpo_order()!r}; the value remains logically " + "valid, but the MPO is not layout-optimized for this tree.", + UserWarning, + stacklevel=2, + ) event_start = len(self.profile_events) work = self.copy() + history_start = len(work.truncation_history) work.apply_submpo( submpo, - where, + logical_where, max_bond=max_bond, - cutoff=cutoff, + cutoff=effective_cutoff, ) + compression_events = work.truncation_history[history_start:] + truncated_events = [ + event for event in compression_events + if event.get("truncated", False) + ] + if truncated_events and warn_on_truncation: + warnings.warn( + "expectation_mpo compressed its private transformed ket on " + f"{len(truncated_events)} edge(s) with max_bond=" + f"{effective_max_bond!r}; the expectation is approximate. " + "Increase max_bond or inspect return_diagnostics=True if " + "an untruncated measurement is required.", + UserWarning, + stacklevel=2, + ) # The bra and ket are separate TTNs. Keep their physical indices shared # for the inner product, but rename the ket's virtual bonds so each @@ -4143,7 +4253,37 @@ def expectation_mpo( self.profile_events.extend( deepcopy(work.profile_events[event_start:]) ) - return result + if not return_diagnostics: + return result + diagnostics = { + "support": tuple(logical_where), + "max_bond": effective_max_bond, + "cutoff": effective_cutoff, + "n_events": len(compression_events), + "n_truncated": len(truncated_events), + "truncated": bool(truncated_events), + "events": deepcopy(compression_events), + } + return result, diagnostics + + def expectation_mpo_exact( + self, submpo, where, *, normalized=True, optimize="auto", + ): + """Evaluate an MPO by exact separate-network contraction. + + Unlike :meth:`expectation_mpo`, this method never applies the MPO to a + copied tree and never compresses a state bond. It delegates to + :meth:`TreeTensorNetwork.expectation_mpo_exact`, which connects the + MPO input legs to a private ket view and its output legs to the bra in + one complete doubled contraction. Native fermionic MPOs therefore + retain Symmray's graded contraction rules. + """ + return self.tn.expectation_mpo_exact( + submpo, + where, + normalized=normalized, + optimize=optimize, + ) def _apply_submpo_resolved(self, submpo, where, *, max_bond=None, cutoff=None, logical_where=None): diff --git a/src/pepsy/optimizers/tree/ttn.py b/src/pepsy/optimizers/tree/ttn.py index 65bc578..e2811fa 100644 --- a/src/pepsy/optimizers/tree/ttn.py +++ b/src/pepsy/optimizers/tree/ttn.py @@ -210,6 +210,50 @@ def _is_symmray_array(value): return hasattr(value, "blocks") and hasattr(value, "indices") +def _native_qr_options_for_tensor(tensor): + """Return the centralized lossless-QR options for one tensor.""" + return {"stabilized": False} if _is_symmray_array(tensor.data) else {} + + +def _native_qr_split_tensor(tensor, **kwargs): + """Split one tree tensor using the native graded QR policy.""" + kwargs.update(_native_qr_options_for_tensor(tensor)) + if _is_symmray_array(tensor.data): + kwargs.setdefault("fn", _native_qr_block_scaled) + kwargs.setdefault("method", "qr") + return tensor.split(**kwargs) + + +def _is_native_mpo(value): + """Return whether an MPO visibly contains native Symmray tensors.""" + marker = getattr(value, "pepsy_tree_native", None) + if marker is not None: + return bool(marker) + tensors = getattr(value, "tensors", None) + if tensors is None: + return None + try: + return any(_is_symmray_array(tensor.data) for tensor in tensors) + except (AttributeError, TypeError): + return None + + +def _tree_plan_signature(plan): + """Return the structural identity used by source-aware tree MPOs.""" + return ( + int(plan.root), + tuple( + (int(node), tuple(int(child) for child in children)) + for node, children in sorted(plan.children.items()) + ), + tuple( + (int(node), int(qubit)) + for node, qubit in sorted(plan.qubit_of_leaf.items()) + ), + None if plan.root_qubit is None else int(plan.root_qubit), + ) + + def _contract_two_tensors(left, right, *, shared_ind=None): """Contract two tensors along one ordinary shared index cheaply. @@ -793,9 +837,195 @@ def local_expectations(self, terms, *, optimize="auto", normalized=True): "seconds": time.perf_counter() - profile_started, }) + def expectation_mpo_exact( + self, mpo, where, *, normalized=True, optimize="auto", + ): + """Contract ```` without applying or compressing the TTN. + + The tree, a private ket view, and the MPO remain separate tensor + networks. A :class:`TreeMPO` or an MPO created by :func:`tree_mpo` + supplies a TreePlan-labelled operator representation, contracted as + ``tree.H | tree_operator | tree``; its optional chain MPO is never + moved into the tree, applied, densified, or compressed. For an + unannotated MPO, the lower physical legs are connected to fresh + copies of the ket physical legs, while its upper physical legs + connect to the bra, and the complete doubled network is contracted. + + ``mpo`` must expose Quimb's regular MPO site interface. Its active site + labels must match ``where``. For a native fermionic TTN the MPO must + contain native Symmray tensors, so the graded contraction rules remain + attached to the operator data. + """ + if isinstance(where, Integral): + where = (int(where),) + else: + where = tuple(int(site) for site in where) + if not where or len(set(where)) != len(where): + raise ValueError("where must contain distinct tree sites.") + if any(site not in self.plan.node_of_qubit for site in where): + raise ValueError(f"site(s) {where!r} are outside this tree state.") + + if hasattr(mpo, "expectation") and hasattr(mpo, "tree_networks"): + all_sites = tuple(sorted(self.plan.node_of_qubit)) + if tuple(sorted(where)) != all_sites: + raise ValueError( + "a TreeMPO must be evaluated on all tree sites so its " + "identity legs remain explicit." + ) + return mpo.expectation( + self, + normalized=normalized, + optimize=optimize, + ) + + tree_operator = getattr(mpo, "pepsy_tree_operator", None) + if tree_operator is not None: + if getattr(mpo, "pepsy_tree_plan_signature", None) != ( + _tree_plan_signature(self.plan) + ): + raise ValueError( + "tree MPO embedding was built for a different TreePlan." + ) + all_sites = tuple(sorted(self.plan.node_of_qubit)) + if tuple(sorted(where)) != all_sites: + raise ValueError( + "a TreePlan MPO embedding must be evaluated on all tree " + "sites so its identity legs remain explicit." + ) + if hasattr(tree_operator, "expectation"): + return tree_operator.expectation( + self, + normalized=normalized, + optimize=optimize, + ) + if not self.fermionic: + raise TypeError( + "native TreePlan MPO embeddings require a native " + "fermionic TreeTensorNetwork." + ) + + operators = ( + tree_operator + if isinstance(tree_operator, (tuple, list)) + else (tree_operator,) + ) + numerator = 0.0 + for operator in operators: + ket = self.copy() + operator_work = operator.copy() + ket_reindex = {} + operator_reindex = {} + for site in all_sites: + physical = self.site_ind(site) + upper = f"k{site}" + lower = f"b{site}" + if upper not in operator_work.ind_map: + raise ValueError( + "TreePlan MPO embedding is missing physical site " + f"{site!r}." + ) + if lower not in operator_work.ind_map: + raise ValueError( + "TreePlan MPO embedding is missing lower physical " + f"site {site!r}." + ) + fresh = qtn.rand_uuid() + ket_reindex[physical] = fresh + operator_reindex[lower] = fresh + ket.reindex_(ket_reindex) + operator_work.reindex_(operator_reindex) + numerator = numerator + (self.H | operator_work | ket).contract( + all, + optimize=optimize, + ) + if not normalized: + return numerator + denominator = (self.H | self).contract(all, optimize=optimize) + return numerator / denominator + + required = ( + "gen_sites_present", "site_tag", "upper_ind_id", "lower_ind_id", + "tag_map", "tensor_map", "copy", + ) + if not all(hasattr(mpo, name) for name in required): + raise TypeError( + "expectation_mpo_exact requires a regular Quimb MPO with " + "site, physical-index, tensor-map, and copy interfaces." + ) + try: + present = tuple(mpo.gen_sites_present()) + except Exception as exc: + raise TypeError( + "could not inspect the MPO's active site labels for an " + "exact tree contraction." + ) from exc + if set(present) != set(where): + raise ValueError( + "MPO active sites must match the declared support: " + f"MPO has {present!r}, where is {where!r}." + ) + + mpo_native = _is_native_mpo(mpo) + if mpo_native is not None and bool(mpo_native) != bool(self.fermionic): + if self.fermionic: + raise TypeError( + "native fermionic TreeTensorNetwork requires a native " + "Symmray MPO for exact graded contraction." + ) + raise TypeError( + "a native Symmray MPO cannot be exactly contracted with an " + "ordinary dense TreeTensorNetwork." + ) + + upper_id = mpo.upper_ind_id + lower_id = mpo.lower_ind_id + ket = self.copy() + mpo_work = mpo.copy() + ket_reindex = {} + mpo_reindex = {} + for site in where: + physical = self.site_ind(site) + upper = upper_id.format(site) + lower = lower_id.format(site) + try: + tids = tuple(mpo_work.tag_map[mpo_work.site_tag(site)]) + except (KeyError, TypeError) as exc: + raise ValueError( + f"MPO has no unique tensor for active site {site!r}." + ) from exc + if len(tids) != 1: + raise ValueError( + f"MPO site {site!r} must resolve to one tensor; " + f"got {len(tids)}." + ) + op_tensor = mpo_work.tensor_map[tids[0]] + if upper not in op_tensor.inds or lower not in op_tensor.inds: + raise ValueError( + f"MPO site {site!r} does not contain expected physical " + f"indices {upper!r} and {lower!r}." + ) + fresh = qtn.rand_uuid() + ket_reindex[physical] = fresh + mpo_reindex[lower] = fresh + + # This is the key orientation: bra <- MPO upper, MPO lower -> ket. + # Every physical index then appears exactly twice, while MPO virtual + # bonds stay internal to the separate structured MPO network. + ket.reindex_(ket_reindex) + mpo_work.reindex_(mpo_reindex) + numerator = (self.H | mpo_work | ket).contract( + all, + optimize=optimize, + ) + if not normalized: + return numerator + denominator = (self.H | self).contract(all, optimize=optimize) + return numerator / denominator + def expectation_mpo( self, mpo, where, *, max_bond=None, cutoff=0.0, - normalized=True, optimize="auto", + normalized=True, optimize="auto", warn_on_truncation=True, + return_diagnostics=False, ): """Evaluate a structured MPO expectation without changing this TTN. @@ -805,6 +1035,9 @@ def expectation_mpo( operator on the full support. The transformed-state bond cap defaults to this TTN's current maximum bond; pass ``max_bond`` explicitly when a larger measurement workspace is acceptable. + ``warn_on_truncation=True`` reports when that workspace actually + truncates the private transformed ket. ``return_diagnostics=True`` + returns the value together with the per-expectation compression report. """ from .optimizer import TreeOptimizer @@ -825,6 +1058,8 @@ def expectation_mpo( cutoff=cutoff, normalized=normalized, optimize=optimize, + warn_on_truncation=warn_on_truncation, + return_diagnostics=return_diagnostics, ) @property @@ -1354,11 +1589,7 @@ def _native_qr_options(self, tensor=None): def _native_qr_split(self, tensor, **kwargs): """Perform a QR split with the native graded zero-sector safeguard.""" - kwargs.update(self._native_qr_options(tensor)) - if _is_symmray_array(tensor.data): - kwargs.setdefault("fn", _native_qr_block_scaled) - kwargs.setdefault("method", "qr") - return tensor.split(**kwargs) + return _native_qr_split_tensor(tensor, **kwargs) def _record_native_compression_route( self, route, *, edge, before_bond, reduction_hint, reduction_proven, diff --git a/src/pepsy/tensors/symmetric.py b/src/pepsy/tensors/symmetric.py index e2dd14f..ea49e81 100644 --- a/src/pepsy/tensors/symmetric.py +++ b/src/pepsy/tensors/symmetric.py @@ -10404,6 +10404,84 @@ def to_mpo( to_backend=to_backend, ) + def to_tree_mpo( + self, + terms_or_edges=None, + *, + hamiltonian=None, + tree=None, + plan=None, + max_bond=None, + cutoff=1e-12, + compress=True, + dtype=None, + fermionic=True, + charge_sectors=False, + to_backend=None, + **params, + ): + """Build a :class:`pepsy.TreeMPO` for a selected ``TreePlan``. + + ``tree`` and ``plan`` are aliases. The returned object exposes the + optional linear representation as ``.chain_mpo`` and the TreePlan + representation through ``.tree_networks`` and ``.expectation``. + Native ``fermionic=True`` keeps Symmray's graded tensors intact for + U1, U1U1, and other supported symmetries. + """ + if tree is not None and plan is not None: + raise TypeError("pass only one of tree= or plan=") + plan = tree if tree is not None else plan + if plan is None: + raise TypeError("to_tree_mpo requires tree= or plan=.") + if hamiltonian is not None: + if terms_or_edges is not None: + raise TypeError( + "Pass either terms_or_edges or hamiltonian, not both." + ) + if not isinstance(hamiltonian, SymHamiltonian): + raise TypeError("hamiltonian must be a SymHamiltonian instance.") + target = hamiltonian + elif isinstance(terms_or_edges, SymHamiltonian): + target = terms_or_edges + else: + if terms_or_edges is None: + raise TypeError( + "to_tree_mpo requires terms_or_edges or hamiltonian." + ) + target = self.hamiltonian( + terms_or_edges, + to_backend=to_backend, + **params, + ) + params = {} + if params: + names = ", ".join(sorted(params)) + raise TypeError( + "Model parameters cannot be supplied with an existing " + f"SymHamiltonian: {names}." + ) + from ..optimizers.tree import tree_mpo + + built = tree_mpo( + plan, + target, + max_bond=max_bond, + cutoff=cutoff, + compress=compress, + dtype=dtype, + fermionic=fermionic, + charge_sectors=charge_sectors, + to_backend=to_backend, + ) + if isinstance(built, dict): + return { + charge: mpo.pepsy_tree_operator + for charge, mpo in built.items() + } + return built.pepsy_tree_operator + + build_tree_mpo = to_tree_mpo + def to_pepo( self, terms_or_edges=None, diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index e9881c2..a66b95d 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -3278,6 +3278,72 @@ def test_tree_expectation_mpo_is_batched_and_non_mutating(): assert opt.tn.validate(check_canonical=True) is opt.tn +def test_tree_expectation_mpo_reports_private_ket_truncation(): + """Expectation diagnostics expose accidental finite-cap truncation.""" + identity = np.eye(4, dtype=complex) + mpo = qtn.MatrixProductOperator.from_dense( + identity, dims=(2, 2), sites=(0, 1), L=4, + ) + h = np.array([[1.0, 1.0], [1.0, -1.0]], dtype=complex) / np.sqrt(2.0) + cnot = np.array( + [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]], + dtype=complex, + ) + opt = TreeOptimizer([(h, 0), (cnot, (0, 1))], n=4, chi=16) + + with pytest.warns(UserWarning, match="private transformed ket"): + value, diagnostics = opt.expectation_mpo( + mpo, (0, 1), max_bond=1, return_diagnostics=True, + ) + + assert diagnostics["truncated"] is True + assert diagnostics["n_truncated"] >= 1 + assert diagnostics["max_bond"] == 1 + assert value != pytest.approx(1.0) + + exact_value, exact_diagnostics = opt.expectation_mpo( + mpo, (0, 1), max_bond=16, return_diagnostics=True, + ) + assert exact_value == pytest.approx(1.0) + assert exact_diagnostics["truncated"] is False + + +def test_tree_exact_mpo_expectation_keeps_mpo_separate(monkeypatch): + """Exact MPO readout does not lower or compress the tree state.""" + identity = np.eye(4, dtype=complex) + mpo = qtn.MatrixProductOperator.from_dense( + identity, dims=(2, 2), sites=(0, 1), L=4, + ) + h = np.array([[1.0, 1.0], [1.0, -1.0]], dtype=complex) / np.sqrt(2.0) + cnot = np.array( + [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]], + dtype=complex, + ) + opt = TreeOptimizer([(h, 0), (cnot, (0, 1))], n=4, chi=1) + before = opt.to_dense().copy() + before_bond = opt.tn.max_bond() + + def forbid_dense(*args, **kwargs): + raise AssertionError("exact MPO readout must not call MPO.to_dense()") + + monkeypatch.setattr(qtn.MatrixProductOperator, "to_dense", forbid_dense) + value = opt.expectation_mpo_exact(mpo, (0, 1)) + + assert value == pytest.approx(1.0) + assert opt.tn.max_bond() == before_bond + assert np.allclose(opt.to_dense(), before) + + +def test_tree_rejects_native_mpo_on_dense_state(): + """Native and ordinary tensor backends cannot be mixed silently.""" + class NativeMarker: + pepsy_tree_native = True + + opt = TreeOptimizer(None, n=2, run=False) + with pytest.raises(TypeError, match="native Symmray MPO"): + opt.apply_submpo(NativeMarker(), (0, 1)) + + def test_tree_subtree_route_batches_sibling_messages(monkeypatch): """Independent leaf messages landing at one node use one contraction.""" n = 4 diff --git a/tests/test_public_api.py b/tests/test_public_api.py index a31c897..a6f9bfc 100644 --- a/tests/test_public_api.py +++ b/tests/test_public_api.py @@ -60,7 +60,7 @@ def test_tree_optimizers_are_available_from_high_level_api(): "SimulatorCandidate", "SimulatorPlan", "SimulatorPlanner", "recommend_simulator", "TreeEnergyOptimizer", "TreeLayoutFinder", - "TreeOptimizer", + "TreeMPO", "TreeOptimizer", "tree_mpo", "TreePlan", "TreeStabOptimizer", "TreeTensorNetwork", @@ -137,7 +137,7 @@ def test_internal_symbols_not_exported(): "PepsEnergyOptimizer", "PepsOptimizer", "SimpleUpdateGen", "SymDMRG2", "PEPSSampleResult", "PepsBpSampler", "CoherentCrosstalkModel", "compile_stim_circuit", "run_coalesced_noisy_shots", "run_coalesced_stim_shots", "run_coalesced_trajectory_shots", "run_noisy_shots", "run_stabilizer_mps_stream", "run_stabilizer_tree_stream", "run_stim_shots", "run_trajectory_shots", "sample_coalesced_bits", "sample_noisy_gate_stream", "sample_noisy_gate_streams", "sample_stim_circuit", "sample_stim_circuits", "sample_trajectory_stream", "TreeEnergyOptimizer", "TreeLayoutFinder", - "TreeOptimizer", + "TreeMPO", "TreeOptimizer", "tree_mpo", "TreePlan", "TreeStabOptimizer", "TreeTensorNetwork", diff --git a/tests/test_tree_mpo.py b/tests/test_tree_mpo.py new file mode 100644 index 0000000..b84135f --- /dev/null +++ b/tests/test_tree_mpo.py @@ -0,0 +1,333 @@ +"""Tests for tree-layout-aware native fermionic MPO construction.""" + +import numpy as np +import pytest + +import pepsy +from pepsy.optimizers.tree import TreeMPO, TreePlan, tree_mpo + + +pytest.importorskip("symmray") + + +@pytest.mark.parametrize("symmetry", ["U1", "U1U1"]) +def test_tree_mpo_preserves_native_fermionic_symmetry_and_expectation(symmetry): + """A tree-ordered MPO remains native and routes exactly over the TTN.""" + fermion = pepsy.Fermion(spinful=True, symmetry=symmetry) + hamiltonian = fermion.hamiltonian( + [(0, 1), (1, 2), (2, 3)], t=1.0, U=2.0, mu=0.1, + ) + plan = TreePlan.from_order( + (2, 0, 3, 1), structure="balanced", top_arity=2, + ) + + mpo = tree_mpo( + plan, + hamiltonian, + fermionic=True, + compress=False, + dtype="complex64", + ) + reference = fermion.to_mpo( + hamiltonian=hamiltonian, + L=4, + fermionic=True, + compress=False, + dtype="complex64", + ) + + assert pepsy.tree_mpo is tree_mpo + assert mpo.pepsy_tree_order == (2, 0, 3, 1) + assert mpo.pepsy_tree_native is True + assert all( + type(tensor.data).__name__ == f"{symmetry}FermionicArray" + for tensor in mpo + ) + assert all( + tuple(mpo[mpo.site_tag(qubit)].tags) == (f"I{qubit}",) + for qubit in range(4) + ) + + state = pepsy.ps_to_ttn( + 4, + tree=plan, + fermion=fermion, + occupations=[0, 1, 0, 1], + dtype="complex64", + ) + tree_value = state.expectation_mpo( + mpo, range(4), max_bond=64, + ) + exact_value = state.expectation_mpo_exact(mpo, range(4)) + reference_value = state.expectation_mpo( + reference, range(4), max_bond=64, + ) + np.testing.assert_allclose(exact_value, reference_value, rtol=3e-5, atol=3e-5) + np.testing.assert_allclose(tree_value, reference_value, rtol=3e-5, atol=3e-5) + + +@pytest.mark.parametrize("symmetry", ["U1", "U1U1"]) +def test_tree_mpo_nonlocal_native_sign_uses_tree_fermion_convention(symmetry): + """A cross-branch hopping keeps the native tree fermionic sign.""" + fermion = pepsy.Fermion( + spinful=True, symmetry=symmetry, dtype="complex128", + ) + plan = TreePlan.from_order(range(4), structure="balanced", top_arity=2) + if symmetry == "U1": + leaf_charges = { + q: fermion.local_fock_state((1, 0), site=q)[0] + for q in range(4) + } + else: + # This charge pattern leaves a non-zero cross-branch hopping sector + # while keeping the test deterministic in the two-component grading. + leaf_charges = { + 0: (0, 0), 1: (0, 0), 2: (0, 0), 3: (1, 0), + } + state = pepsy.TreeTensorNetwork.from_symmray_plan( + plan, + symmetry=symmetry, + physical_sectors=fermion.physical_sectors, + leaf_charges=leaf_charges, + bond_dim=4, + fermionic=True, + seed=10, + dtype="complex128", + ) + hamiltonian = fermion.hamiltonian( + [(0, 2)], t=1.0, U=0.0, mu=0.0, + ) + direct = state.local_expectation( + hamiltonian.terms[(0, 2)], (0, 2), + ) + mpo = tree_mpo(plan, hamiltonian, fermionic=True, compress=False) + + np.testing.assert_allclose( + state.expectation_mpo_exact(mpo, range(4)), + direct, + rtol=1e-12, + atol=1e-12, + ) + + +@pytest.mark.parametrize("symmetry", ["U1", "U1U1"]) +def test_tree_mpo_pair_observable_stays_compact_and_contracts_on_tree(symmetry): + """The full pair table uses the four-state native endpoint automaton.""" + fermion = pepsy.Fermion( + spinful=True, symmetry=symmetry, dtype="complex128", + ) + nsite = 4 + reference = fermion.hamiltonian([(0, 1)], t=0.0, U=0.0, mu=0.0) + terms = { + (left, right): fermion.operator_term( + [( + 2.0 / nsite * ((-1) ** (left + right)), + ((left, "pair_create"), (right, "pair_annihilate")), + )], + sites=(left, right), + ) + for left in range(nsite) + for right in range(left + 1, nsite) + } + hamiltonian = type(reference).from_terms( + reference.model, + reference.symmetry, + terms, + parameters=reference.parameters, + ) + plan = TreePlan.from_order( + range(nsite), structure="balanced", top_arity=2, + ) + if symmetry == "U1": + leaf_charges = { + q: fermion.local_fock_state((1, 0), site=q)[0] + for q in range(nsite) + } + else: + leaf_charges = { + 0: (1, 1), 1: (0, 0), 2: (0, 0), 3: (0, 0), + } + state = pepsy.TreeTensorNetwork.from_symmray_plan( + plan, + symmetry=symmetry, + physical_sectors=fermion.physical_sectors, + leaf_charges=leaf_charges, + bond_dim=4, + fermionic=True, + seed=10, + dtype="complex128", + ) + + mpo = tree_mpo(plan, hamiltonian, fermionic=True, compress=False) + assert mpo.max_bond() == 4 + assert mpo.pepsy_tree_operator.pepsy_tree_operator_kind == ( + "pair_endpoint_automaton" + ) + direct = sum( + state.local_expectation(term, support) + for support, term in terms.items() + ) + exact = state.expectation_mpo_exact(mpo, range(nsite)) + np.testing.assert_allclose(exact, direct.real, rtol=1e-12, atol=1e-12) + + +def test_tree_plan_mpo_order_handles_a_physical_root(): + """The root physical site is explicit and stable in the MPO order.""" + plan = TreePlan.from_order( + (1, 2, 3), root_qubit=0, structure="balanced", top_arity=2, + ) + + assert plan.mpo_order() == (0, 1, 2, 3) + assert plan.mpo_order(include_root=False) == (1, 2, 3) + + +def test_tree_mpo_class_dense_backend_and_direct_expectation(): + """The general TreeMPO class supports ordinary dense term mappings.""" + plan = TreePlan.from_order(range(4), structure="balanced", top_arity=2) + state = pepsy.TreeTensorNetwork.from_plan(plan, dtype=complex) + terms = { + (0,): np.diag([1.0, 2.0]), + (1, 2): np.arange(16, dtype=complex).reshape(2, 2, 2, 2), + } + + operator = TreeMPO.from_terms(plan, terms) + direct = sum( + state.local_expectation(term, support) + for support, term in terms.items() + ) + + assert operator.backend == "dense" + assert operator.chain_mpo is None + assert len(operator.tree_networks) == 1 + assert operator.tree_network.pepsy_tree_operator_kind == "dense_tree_tnno" + assert operator.compressed is True + np.testing.assert_allclose(operator.expectation(state), direct) + np.testing.assert_allclose( + state.expectation_mpo_exact(operator, range(4)), direct, + ) + operator.canonicalize() + operator.compress(max_bond=4) + np.testing.assert_allclose(operator.expectation(state), direct, atol=1e-10) + + +def test_native_tree_mpo_amalgamates_higher_rank_term(): + """A native three-site term is a TTNO, not an uncompressible hyperedge.""" + fermion = pepsy.Fermion( + spinful=True, symmetry="U1U1", dtype="complex128", + ) + term = fermion.operator_term([(0.7, ())], sites=(2, 0, 3)) + reference = fermion.hamiltonian( + [(0, 1)], t=0.0, U=0.0, mu=0.0, + ) + hamiltonian = type(reference).from_terms( + reference.model, + reference.symmetry, + {(2, 0, 3): term}, + parameters=reference.parameters, + ) + plan = TreePlan.from_order( + (3, 1, 0, 2), structure="balanced", top_arity=2, + ) + state = pepsy.TreeTensorNetwork.from_symmray_plan( + plan, + symmetry="U1U1", + physical_sectors=fermion.physical_sectors, + leaf_charges={q: (0, 0) for q in range(4)}, + bond_dim=3, + fermionic=True, + seed=13, + dtype="complex128", + ) + + operator = TreeMPO.from_hamiltonian( + plan, hamiltonian, compress=False, dtype="complex128", + ) + direct = state.local_expectation(term, (2, 0, 3)) + np.testing.assert_allclose(operator.expectation(state), direct) + assert operator.tree_network.pepsy_tree_operator_kind == "native_tree_tnno" + + operator.canonicalize().compress(cutoff=1e-12) + np.testing.assert_allclose(operator.expectation(state), direct) + + +@pytest.mark.parametrize("symmetry", ["U1", "U1U1"]) +def test_fermion_tree_mpo_class_is_native_and_keeps_chain_representation(symmetry): + """Fermion exposes the class API for both native Symmray symmetries.""" + fermion = pepsy.Fermion( + spinful=True, symmetry=symmetry, dtype="complex128", + ) + hamiltonian = fermion.hamiltonian( + [(0, 1), (1, 2)], t=1.0, U=2.0, mu=0.1, + ) + plan = TreePlan.from_order(range(4), structure="balanced", top_arity=2) + operator = fermion.to_tree_mpo( + hamiltonian=hamiltonian, + tree=plan, + compress=True, + ) + state = pepsy.TreeTensorNetwork.from_symmray_plan( + plan, + symmetry=symmetry, + physical_sectors=fermion.physical_sectors, + leaf_charges=( + {q: 1 for q in range(4)} + if symmetry == "U1" + else {q: (0, 0) for q in range(4)} + ), + bond_dim=2, + fermionic=True, + seed=17, + dtype="complex128", + ) + + assert isinstance(operator, TreeMPO) + assert operator.backend == "symmray" + assert operator.chain_mpo is not None + assert len(operator.tree_networks) == 1 + assert operator.tree_network.pepsy_tree_operator_kind == "native_tree_tnno" + assert operator.compressed is True + assert operator.pepsy_compression_report["compressed"] is True + assert all( + type(tensor.data).__name__.endswith("FermionicArray") + for network in operator.tree_networks + for tensor in network + ) + np.testing.assert_allclose( + operator.expectation(state), + state.expectation_mpo_exact(operator, range(4)), + rtol=1e-12, + atol=1e-12, + ) + copied = operator.copy().canonicalize() + np.testing.assert_allclose( + copied.expectation(state), + operator.expectation(state), + rtol=1e-12, + atol=1e-12, + ) + + +def test_sparse_pair_terms_do_not_invent_missing_pairs(): + """Only a complete separable pair table may use the compact automaton.""" + fermion = pepsy.Fermion( + spinful=True, symmetry="U1", dtype="complex128", + ) + hamiltonian = fermion.hamiltonian( + [(0, 1), (1, 2), (2, 3)], t=1.0, U=0.0, mu=0.0, + ) + plan = TreePlan.from_order(range(4), structure="balanced", top_arity=2) + + operator = fermion.to_tree_mpo( + hamiltonian=hamiltonian, + tree=plan, + compress=False, + ) + + assert len(operator.tree_networks) == 1 + assert operator.tree_network.pepsy_tree_operator_kind == "native_tree_tnno" + assert operator.max_bond() > 1 + raw_bond = operator.max_bond() + operator.canonicalize() + operator.compress(max_bond=16) + assert operator.max_bond() <= 16 + assert operator.max_bond() <= raw_bond From 060d5b79ae860b11f4028768675689461b1b2003 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Tue, 4 Aug 2026 09:19:24 -0600 Subject: [PATCH 63/70] Unify Fermion MPO and tree operator builders --- .github/skills/tree-optimizer/SKILL.md | 22 +- docs/api/operators/hamiltonians.md | 16 +- docs/api/optimizers/energy.md | 5 +- docs/api/optimizers/mpo.md | 23 +- docs/api/optimizers/tree.md | 39 +- docs/api/tensors/symmetric.md | 17 +- src/pepsy/__init__.py | 3 +- src/pepsy/operators/hamiltonians.py | 5 +- src/pepsy/optimizers/__init__.py | 1 + src/pepsy/optimizers/tree/__init__.py | 3 +- src/pepsy/optimizers/tree/layout.py | 23 +- src/pepsy/optimizers/tree/operators.py | 1430 +++++++++++++++++++++++- src/pepsy/optimizers/tree/optimizer.py | 3 +- src/pepsy/tensors/symmetric.py | 173 +-- tests/test_optimize_mpo.py | 28 +- tests/test_public_api.py | 4 +- tests/test_tree_mpo.py | 127 ++- 17 files changed, 1704 insertions(+), 218 deletions(-) diff --git a/.github/skills/tree-optimizer/SKILL.md b/.github/skills/tree-optimizer/SKILL.md index 7d1ae0e..1ec09dc 100644 --- a/.github/skills/tree-optimizer/SKILL.md +++ b/.github/skills/tree-optimizer/SKILL.md @@ -389,11 +389,11 @@ interface use the dense `to_dense()` fallback and remain subject to ### Tree-native MPO API When the consumer is a `TreeTensorNetwork`, `TreeMPO` is the primary operator -API. Use `TreePlan.to_tree_mpo(...)` or -`Fermion.to_tree_mpo(..., tree=plan)`: +API. Use `TreePlan.build_tree_operator(...)` or +`Fermion.build_tree_operator(..., tree=plan)`: ```python -tree_operator = fermion.to_tree_mpo( +tree_operator = fermion.build_tree_operator( hamiltonian=hamiltonian, tree=plan, compress=True, @@ -404,11 +404,23 @@ energy = tree.expectation_mpo_exact(tree_operator, range(plan.n)) ``` `tree_operator.chain_mpo` is optional compatibility data for ordinary MPS/MPO -workflows. `TreePlan.to_mpo(...)` and `tree_mpo(...)` return that regular chain -MPO and attach the `TreeMPO`; they do not change the tree contraction route. +workflows. `TreePlan.to_mpo(...)` and `tree_mpo(...)` remain compatibility +builders that return that regular chain MPO and attach the `TreeMPO`; they do +not change the tree contraction route. `to_tree_mpo(...)` remains a +compatibility alias for `build_tree_operator(...)`. The chain MPO must not be moved into the tree, densified, or compressed as a state update for exact tree measurement. +`TreeMPO` subclasses Quimb's `TensorNetworkGenOperator`, analogous to +`TreeTensorNetwork` subclassing `TensorNetworkGenVector`. It is the tree twin +of Quimb's `MatrixProductOperator`: its public operator surface includes +`sites`, `nsites`, `site_tag`, `upper_ind`, `lower_ind`, `to_dense`, `H`, +`copy`, `identity`, `from_dense`, `add_MPO`, `singular_values`, `amplitude`, +and canonicalize/compress helpers, while `plan`, `node_tensor`, `neighbors`, +and `bond` provide the branched geometry. It cannot inherit the chain-only +`MatrixProductOperator` implementation because a tree has no left/right +ordering; `chain_mpo` remains the separate chain-compatible representation. + For native fermionic Hamiltonians, one-, two-, and higher-site neutral terms are fused and factorized from their native Symmray operator tensor over the TreePlan Steiner subtree, then amalgamated into one charge-aware direct-sum diff --git a/docs/api/operators/hamiltonians.md b/docs/api/operators/hamiltonians.md index 8cc3bc2..603684a 100644 --- a/docs/api/operators/hamiltonians.md +++ b/docs/api/operators/hamiltonians.md @@ -16,22 +16,24 @@ mpo = builder.build_mpo( ``` The native model-facing shorthand is -`fermion.to_mpo(edges, L=3, t=..., U=..., mu=...)`. Couplings remain explicit; -they are not stored on the `Fermion` object. Use `fermion.build_mpo(...)` when -the Jordan-Wigner-compatible MPO convention is wanted. +`fermion.build_mpo(edges, L=3, t=..., U=..., mu=...)`. Couplings remain +explicit; they are not stored on the `Fermion` object. Pass `fermionic=False` +to the same builder when the Jordan-Wigner-compatible MPO convention is +wanted. -`Fermion.to_mpo(...)` and `SymHamiltonian.to_mpo(..., fermionic=True)` return +`Fermion.build_mpo(...)` and `SymHamiltonian.to_mpo(..., fermionic=True)` return native graded `FermionicArray` MPO tensors. Explicit mappings can contain arbitrary homogeneous-charge multi-site terms; non-contiguous supports are represented by charged virtual channels, and the open boundary carries the operator charge when it is nonzero. `ham_tn.build_mpo(..., fermionic=True)` -selects the same native path. Pass `to_backend=...` to map the stored Symmray -blocks to a selected array backend. +selects the same native path. `Fermion.to_mpo(...)` remains a compatibility +alias. Pass `to_backend=...` to map the stored Symmray blocks to a selected +array backend. For a mixed-charge operator, request an explicit charge-sector decomposition: ```python -sectors = fermion.to_mpo( +sectors = fermion.build_mpo( mixed_terms, L=4, fermionic=True, diff --git a/docs/api/optimizers/energy.md b/docs/api/optimizers/energy.md index 9ebabc7..3ce2a9e 100644 --- a/docs/api/optimizers/energy.md +++ b/docs/api/optimizers/energy.md @@ -30,11 +30,14 @@ the automatic re-encoding can create very large block-sparse intermediates. For a deliberately small or explicitly managed conversion, pass ``allow_encoding_conversion=True`` to ``MpsEnergyOptimizer``. -An MPO returned by ``Fermion.to_mpo(...)`` is native graded and can be measured +An MPO returned by ``Fermion.build_mpo(...)`` is native graded and can be measured directly with a native fermionic MPS. Pepsy applies that MPO sitewise as a factorized graded MPO-MPS network, preserving Symmray's contraction order without materializing an exponentially sized operator. +``Fermion.to_mpo(...)`` remains a compatibility alias of +``Fermion.build_mpo(...)``. + Repeated native-MPO evaluations reuse a per-optimizer cotengra path cache. The default is uncompressed and exact. For a controlled approximation, pass for example ``native_mpo_compression={"max_bond": 64, "cutoff": 1e-12, diff --git a/docs/api/optimizers/mpo.md b/docs/api/optimizers/mpo.md index 385c5ca..0f2795b 100644 --- a/docs/api/optimizers/mpo.md +++ b/docs/api/optimizers/mpo.md @@ -1,14 +1,15 @@ # `pepsy.optimizers.mpo.optimizer` `MpoOptimizer` accepts ordinary Quimb MPOs and Symmray block-sparse MPOs. For -a native graded fermion workflow, use `Fermion.to_mpo(...)` and replay the -matching native gate stream: +a native graded fermion workflow, use the canonical +`Fermion.build_mpo(...)` entry point and replay the matching native gate +stream: ```python fermion = pepsy.Fermion(spinful=True, symmetry="U1U1") edges = [(0, 1), (1, 2)] hamiltonian = fermion.hamiltonian(edges, t=1.0, U=2.0, mu=0.1) -mpo = fermion.to_mpo(hamiltonian=hamiltonian, L=3) +mpo = fermion.build_mpo(hamiltonian=hamiltonian, L=3) opt = pepsy.MpoOptimizer( mpo, @@ -28,18 +29,20 @@ term = fermion.operator_term( sites=(0, 2), add_hc=True, ) -mpo = fermion.to_mpo({(0, 2): term}, L=3) +mpo = fermion.build_mpo({(0, 2): term}, L=3) ``` -The Jordan-Wigner compatibility path remains available through -`Fermion.build_mpo(...)` and the matching +The Jordan-Wigner compatibility path remains available by passing +`fermionic=False` to the same builder, together with the matching `SymHamiltonian.jw_trotter_gates(...)` stream: ```python fermion = pepsy.Fermion(spinful=True, symmetry="U1U1") edges = [(0, 1), (1, 2)] hamiltonian = fermion.hamiltonian(edges, t=1.0, U=2.0, mu=0.1) -mpo = fermion.build_mpo(edges, L=3, t=1.0, U=2.0, mu=0.1) +mpo = fermion.build_mpo( + edges, L=3, t=1.0, U=2.0, mu=0.1, fermionic=False +) opt = pepsy.MpoOptimizer( mpo, @@ -63,9 +66,9 @@ compression. `mode="svd"`, `mode="mpo"`, and `mode="dmrg"` use the same block-aware path for native Symmray MPOs; the optimizer does not require a dense conversion of the input MPO. `ham_tn.build_mpo(..., fermionic=True)` is also routed to the native -`Fermion.to_mpo(...)` entry point. Use `to_backend=...` on the model-facing -builder when the stored blocks must be moved to Torch or another supported -backend. +`Fermion.build_mpo(...)` entry point. `Fermion.to_mpo(...)` remains a +compatibility alias. Use `to_backend=...` on the model-facing builder when +the stored blocks must be moved to Torch or another supported backend. Native MPO assembly/replay is also measurable with a native fermionic MPS. `MpsEnergyOptimizer` applies the native MPO sitewise as a factorized graded diff --git a/docs/api/optimizers/tree.md b/docs/api/optimizers/tree.md index fe7540a..0152ab3 100644 --- a/docs/api/optimizers/tree.md +++ b/docs/api/optimizers/tree.md @@ -48,34 +48,35 @@ Nevergrad objectives permute only the remaining leaf sites. ### Layout-aware native MPOs -After selecting a plan, build a Hamiltonian MPO in the plan's logical tree -order with `plan.to_mpo(...)` or `tree_mpo(plan, hamiltonian)`. The default is -the native graded Symmray path, including `U1FermionicArray` and -`U1U1FermionicArray` tensors: +After selecting a plan, the canonical tree-native operator is built with +`plan.build_tree_operator(...)` or `Fermion.build_tree_operator(...)`. The +legacy `tree_mpo(plan, hamiltonian)` / `plan.to_mpo(...)` builders return the +ordinary compatibility chain MPO with the tree operator attached. The native +path includes `U1FermionicArray` and `U1U1FermionicArray` tensors: ```python from pepsy import Fermion -from pepsy.optimizers.tree import TreeLayoutFinder, tree_mpo +from pepsy.optimizers.tree import TreeLayoutFinder finder = TreeLayoutFinder(gates, n=8, max_arity=2) plan = finder.run(order="quality") fermion = Fermion(spinful=True, symmetry="U1U1") hamiltonian = fermion.hamiltonian(edges, t=1.0, U=2.0, mu=0.1) -mpo = tree_mpo(plan, hamiltonian) # fermionic=True by default -# equivalent: -# mpo = plan.to_mpo(hamiltonian) -energy = opt.tn.expectation_mpo_exact(mpo, range(plan.n)) +tree_operator = plan.build_tree_operator(hamiltonian) +chain_mpo = tree_operator.chain_mpo +energy = opt.tn.expectation_mpo_exact(chain_mpo, range(plan.n)) ``` -For the operator-level API, use `TreeMPO` (or -`fermion.to_tree_mpo(...)`). It exposes both representations without mixing -their tensor-network geometries: +For a model-facing operator build, use `Fermion.build_tree_operator(...)`. +`Fermion.to_tree_mpo(...)` and `TreePlan.to_tree_mpo(...)` remain compatibility +aliases. `TreeMPO` exposes the optional chain representation without mixing +the two tensor-network geometries: ```python from pepsy import Fermion -tree_operator = fermion.to_tree_mpo( +tree_operator = fermion.build_tree_operator( hamiltonian=hamiltonian, tree=plan, compress=True, @@ -153,6 +154,18 @@ Pass `fermionic=False` only for dense ordinary/Jordan--Wigner-compatible terms. Existing `OneDMap` lattice maps remain unchanged and should continue to be used for regular 2D/3D coordinate layouts. +`TreeMPO` subclasses Quimb's `TensorNetworkGenOperator`, in the same way that +`TreeTensorNetwork` subclasses `TensorNetworkGenVector`. It is the tree twin +of Quimb's `MatrixProductOperator`: the common operator surface includes +`sites`, `nsites`, `site_tag`, `upper_ind`, `lower_ind`, `to_dense`, `H`, +`copy`, `identity`, `from_dense`, `add_MPO`, `singular_values`, and +`amplitude`. Tree-specific geometry is exposed through `plan`, `node_tensor`, +`neighbors`, and `bond`; `canonicalize`/`compress` perform the corresponding +tree-wide QR/SVD sweeps. It cannot inherit Quimb's chain-only +`MatrixProductOperator` implementation because a branched tree has no single +left/right ordering. The optional `chain_mpo` remains the separate object for +chain workflows. + The conventional binary TTN with a three-leg top tensor is the default when there are at least three leaves and no `root_qubit`. Pass `max_arity=2, top_arity=3` explicitly to `TreePlan.from_order`, diff --git a/docs/api/tensors/symmetric.md b/docs/api/tensors/symmetric.md index c333b9f..4ab511a 100644 --- a/docs/api/tensors/symmetric.md +++ b/docs/api/tensors/symmetric.md @@ -594,14 +594,14 @@ mpo = fermion.build_mpo( ) ``` -This returns the existing symmetry-preserving Jordan-Wigner-compatible MPO -convention. For the native graded path, use ``to_mpo``: +This returns the native graded MPO by default. For the explicit +Jordan-Wigner-compatible convention, pass ``fermionic=False``: ```python -mpo = fermion.to_mpo( - [(0, 1), (1, 2)], L=3, t=1.0, U=8.0, mu=0.0, max_bond=16 +mpo = fermion.build_mpo( + [(0, 1), (1, 2)], L=3, t=1.0, U=8.0, mu=0.0, max_bond=16, + fermionic=False, ) -assert all(type(tensor.data).__name__ == "U1U1FermionicArray" for tensor in mpo) ``` Native fermionic gate streams from the same ``Fermion`` model can be passed to @@ -673,9 +673,10 @@ assembly, replay, and exact energy measurement are supported. Native MPO energy applies the operator sitewise as a factorized graded MPO-MPS contraction, so its cost is controlled by the MPS and MPO bond dimensions. -``Fermion.build_mpo(..., fermionic=True)`` selects the same native construction -as ``Fermion.to_mpo(...)``. Omitting ``fermionic=True`` retains the explicit -Jordan--Wigner compatibility MPO path. +``Fermion.build_mpo(...)`` is the canonical native construction and defaults +to graded Symmray tensors. Passing ``fermionic=False`` selects the explicit +Jordan--Wigner compatibility MPO path. ``Fermion.to_mpo(...)`` remains a +compatibility alias. For periodic square lattices encoded as long-range edges in an OBC MPS/MPO, ``mode="folded-snake"`` alternates opposite columns before snaking. On a 6 by diff --git a/src/pepsy/__init__.py b/src/pepsy/__init__.py index 4acfc2a..ca71d60 100644 --- a/src/pepsy/__init__.py +++ b/src/pepsy/__init__.py @@ -156,6 +156,7 @@ "TreeMPO": ".optimizers", "TreeOptimizer": ".optimizers", "TreePlan": ".optimizers", + "build_tree_operator": ".optimizers", "tree_mpo": ".optimizers", "square_lattice_zigzag": ".optimizers", "TreeStabOptimizer": ".optimizers", @@ -362,7 +363,7 @@ def __getattr__(name): y, z, ) - from .optimizers import DeferredInjectionRecord, DeferredInjectionReport, DeferredProjectionRecord, GlobalOptimizer, ImmediateInjectionReport, ImmediateProjectionRecord, MeasurementRecord, MpoOptimizer, MpsEnergyOptimizer, MpsOptimizer, MpsStabOptimizer, NormEventRecord, PepsEnergyOptimizer, PepsOptimizer, STNState, StabilizerMpsSettingsAdvice, StabilizerMpsRunResult, StabilizerMpsSimulator, StabilizerTreeRunResult, StreamAnalysisRecord, SimpleUpdateGen, SymDMRG2, SweepOptimizer, TreeEnergyOptimizer, TreeLayoutFinder, TreeMPO, TreeOptimizer, TreePlan, TreeStabOptimizer, run_stabilizer_mps_stream, run_stabilizer_tree_stream # noqa: F401 + from .optimizers import DeferredInjectionRecord, DeferredInjectionReport, DeferredProjectionRecord, GlobalOptimizer, ImmediateInjectionReport, ImmediateProjectionRecord, MeasurementRecord, MpoOptimizer, MpsEnergyOptimizer, MpsOptimizer, MpsStabOptimizer, NormEventRecord, PepsEnergyOptimizer, PepsOptimizer, STNState, StabilizerMpsSettingsAdvice, StabilizerMpsRunResult, StabilizerMpsSimulator, StabilizerTreeRunResult, StreamAnalysisRecord, SimpleUpdateGen, SymDMRG2, SweepOptimizer, TreeEnergyOptimizer, TreeLayoutFinder, TreeMPO, TreeOptimizer, TreePlan, TreeStabOptimizer, build_tree_operator, run_stabilizer_mps_stream, run_stabilizer_tree_stream # noqa: F401 from .sampling import FermionConfigurationEncoding, MpsDiagonalEstimate, MpsBatchSampleResult, MpsSampleResult, MpsSampler, PEPSSampleResult, PepsBpSampler, TreeBatchSampleResult, TreeSampleResult, TreeSampler, VecSampler # noqa: F401 from .solvers import FDSolver # noqa: F401 from .tensors import ( # noqa: F401 diff --git a/src/pepsy/operators/hamiltonians.py b/src/pepsy/operators/hamiltonians.py index 1b5e287..38739b2 100644 --- a/src/pepsy/operators/hamiltonians.py +++ b/src/pepsy/operators/hamiltonians.py @@ -517,7 +517,7 @@ def build_mpo( fermionic : bool | None, default=None Native graded encoding flag for the fermion-model form. ``None`` and ``False`` select the Jordan-Wigner-compatible MPO builder; - ``True`` selects ``Fermion.to_mpo(...)``. + ``True`` selects the native graded ``Fermion.build_mpo(...)``. charge_sectors : bool, default=False When native construction is enabled, return one MPO per operator charge as ``{charge: mpo}`` instead of requiring one homogeneous @@ -562,8 +562,7 @@ def build_mpo( fermionic_use = False if fermionic is None else bool(fermionic) if charge_sectors and not fermionic_use: raise ValueError("charge_sectors=True requires fermionic=True.") - mpo_builder = fermion.to_mpo if fermionic_use else fermion.build_mpo - return mpo_builder( + return fermion.build_mpo( ints, L=self.L, mapper=mapper_use, diff --git a/src/pepsy/optimizers/__init__.py b/src/pepsy/optimizers/__init__.py index 8d92680..f267946 100644 --- a/src/pepsy/optimizers/__init__.py +++ b/src/pepsy/optimizers/__init__.py @@ -42,6 +42,7 @@ "TreeMPO": ".tree", "TreeOptimizer": ".tree", "TreePlan": ".tree", + "build_tree_operator": ".tree", "tree_mpo": ".tree", "square_lattice_zigzag": "._layout_orders", "TreeStabOptimizer": ".tree_stabilizer", diff --git a/src/pepsy/optimizers/tree/__init__.py b/src/pepsy/optimizers/tree/__init__.py index 33c8a66..ca3b2da 100644 --- a/src/pepsy/optimizers/tree/__init__.py +++ b/src/pepsy/optimizers/tree/__init__.py @@ -8,7 +8,7 @@ from .layout import TreeLayoutFinder, TreePlan from .optimizer import TreeOptimizer -from .operators import TreeMPO, tree_mpo +from .operators import TreeMPO, build_tree_operator, tree_mpo from .ttn import TreeTensorNetwork __all__ = [ @@ -17,5 +17,6 @@ "TreeLayoutFinder", "TreePlan", "TreeMPO", + "build_tree_operator", "tree_mpo", ] diff --git a/src/pepsy/optimizers/tree/layout.py b/src/pepsy/optimizers/tree/layout.py index f765915..bb877a5 100644 --- a/src/pepsy/optimizers/tree/layout.py +++ b/src/pepsy/optimizers/tree/layout.py @@ -1080,22 +1080,21 @@ def to_mpo(self, hamiltonian, **kwargs): return tree_mpo(self, hamiltonian, **kwargs) - def to_tree_mpo(self, hamiltonian, **kwargs): - """Build the public :class:`TreeMPO` operator for this plan. + def build_tree_operator(self, hamiltonian, **kwargs): + """Build the canonical :class:`TreeMPO` operator for this plan. - The returned class keeps the optional chain MPO available as + The returned object keeps the optional chain MPO available as ``.chain_mpo`` and exposes the TreePlan-routed representation through - ``.tree_networks`` and ``.expectation``. + ``.tree_networks`` and ``.expectation``. With + ``charge_sectors=True`` the result is ``{charge: TreeMPO}``. """ - from .operators import tree_mpo + from .operators import build_tree_operator - chain_mpo = tree_mpo(self, hamiltonian, **kwargs) - if isinstance(chain_mpo, dict): - return { - charge: mpo.pepsy_tree_operator - for charge, mpo in chain_mpo.items() - } - return chain_mpo.pepsy_tree_operator + return build_tree_operator(self, hamiltonian, **kwargs) + + # Compatibility spelling retained while ``build_tree_operator`` becomes + # the single plan-facing tree operator builder. + to_tree_mpo = build_tree_operator def is_leaf(self, nid): return len(self.children.get(nid, ())) == 0 diff --git a/src/pepsy/optimizers/tree/operators.py b/src/pepsy/optimizers/tree/operators.py index efe53cc..dd93776 100644 --- a/src/pepsy/optimizers/tree/operators.py +++ b/src/pepsy/optimizers/tree/operators.py @@ -15,10 +15,11 @@ staggered eta-pair observable and keeps its tree bond independent of the lattice size. -The returned public object remains a regular Quimb MPO for compatibility with -MPS/MPO APIs. It carries the tree operator as private metadata consumed by -``TreeTensorNetwork.expectation_mpo_exact``. The two networks remain -separate throughout the contraction. +``TreeMPO`` is a Quimb ``TensorNetworkGenOperator`` over the TreePlan geometry, +analogous to ``TreeTensorNetwork`` being a ``TensorNetworkGenVector``. The +``tree_mpo`` compatibility builder additionally returns a regular Quimb MPO +for MPS/MPO APIs and attaches the ``TreeMPO`` to it. The two representations +remain separate throughout the tree contraction. """ from __future__ import annotations @@ -27,11 +28,13 @@ from numbers import Integral import warnings +import autoray as ar import numpy as np +import quimb.tensor as qtn from .layout import TreePlan -__all__ = ["TreeMPO", "tree_mpo"] +__all__ = ["TreeMPO", "build_tree_operator", "tree_mpo"] def _tree_plan_signature(plan): @@ -50,12 +53,58 @@ def _tree_plan_signature(plan): ) -class TreeMPO: +def _tree_node_selector(plan, selector): + """Resolve one public tree-node selector to a structural node id.""" + if isinstance(selector, str): + if selector.startswith("N"): + selector = selector[1:] + else: + raise ValueError( + "TreeMPO geometry selectors must be node ids or N tags." + ) + try: + node = int(selector) + except (TypeError, ValueError) as exc: + raise TypeError(f"invalid TreeMPO node selector {selector!r}.") from exc + if node not in plan.children: + raise ValueError(f"{node!r} is not a TreePlan node.") + return node + + +def _tree_region_selector(plan, selector): + """Resolve one or more node selectors to a connected TreePlan region.""" + if isinstance(selector, (tuple, list, set, frozenset)): + nodes = tuple(_tree_node_selector(plan, node) for node in selector) + else: + nodes = (_tree_node_selector(plan, selector),) + if not nodes: + raise ValueError("TreeMPO canonical regions cannot be empty.") + region = set(nodes) + for node in nodes[1:]: + region.update(plan.node_path(nodes[0], node)) + return frozenset(region) + + +def _tree_subtree_span(plan, nodes): + """Return the minimal connected node set spanning ``nodes``.""" + nodes = tuple(nodes) + if not nodes: + raise ValueError("need at least one tree node to span a subtree.") + region = {nodes[0]} + for node in nodes[1:]: + region.update(plan.node_path(nodes[0], node)) + return frozenset(region) + + +class TreeMPO(qtn.TensorNetworkGenOperator): """TreePlan-aware operator with dense and native Symmray backends. ``TreeMPO`` is the operator-level API for measurements on a - :class:`TreeTensorNetwork`. It deliberately keeps the optional linear - chain MPO separate from the tree representation: + :class:`TreeTensorNetwork`. It subclasses Quimb's generalized operator + network, so common methods such as ``sites``, ``site_tag``, ``upper_ind``, + ``lower_ind``, ``to_dense``, ``H``, and ``copy`` operate on its primary + TreePlan network. It deliberately keeps the optional linear chain MPO + separate from the tree representation: ``chain_mpo`` The ordinary Quimb ``MatrixProductOperator`` produced by @@ -75,10 +124,26 @@ class TreeMPO: compact network instead. """ + # Match Quimb's generalized operator API while retaining the additional + # TreePlan/chain representation metadata owned by this class. + _EXTRA_PROPS = qtn.TensorNetworkGenOperator._EXTRA_PROPS + ( + "_plan", + "_node_tag_id", + "_pepsy_backend", + "tree_networks", + "_canonical_region", + "chain_mpo", + "terms", + "fermionic", + "symmetry", + "cutoff", + "compressed", + ) + def __init__( self, - plan, - tree_networks, + plan=None, + tree_networks=None, *, chain_mpo=None, terms=None, @@ -87,24 +152,72 @@ def __init__( symmetry=None, cutoff=1e-12, compressed=False, + sites=None, + site_tag_id="I{}", + upper_ind_id="k{}", + lower_ind_id="b{}", + node_tag_id="N{}", + virtual=True, + deep=False, ): + if isinstance(plan, TreeMPO) and tree_networks is None: + source = plan + plan = source.plan + networks = tuple( + network.copy(virtual=virtual, deep=deep) + for network in source.tree_networks + ) + chain_mpo = ( + None + if source.chain_mpo is None + else source.chain_mpo.copy(virtual=virtual, deep=deep) + ) + terms = source.terms + backend = source.backend + fermionic = source.fermionic + symmetry = source.symmetry + cutoff = source.cutoff + compressed = source.compressed + sites = source.sites + site_tag_id = source.site_tag_id + upper_ind_id = source.upper_ind_id + lower_ind_id = source.lower_ind_id + node_tag_id = source.node_tag_id + elif tree_networks is None: + raise TypeError("TreeMPO requires a tree operator network.") + else: + networks = ( + tuple(tree_networks) + if isinstance(tree_networks, (tuple, list)) + else (tree_networks,) + ) + if not isinstance(plan, TreePlan): raise TypeError("plan must be a TreePlan.") - if isinstance(tree_networks, (tuple, list)): - networks = tuple(tree_networks) - else: - networks = (tree_networks,) if not networks or any(network is None for network in networks): raise ValueError("TreeMPO requires at least one tree operator network.") - self.plan = plan + # The first network is the primary generalized tree operator. Use a + # virtual Quimb view so inherited operator methods such as + # ``sites``, ``upper_ind``, ``lower_ind``, ``to_dense``, ``bond``, and + # ``H`` operate on the same tensors as ``tree_networks[0]``. + super().__init__(networks[0], virtual=True) + self._plan = plan self.tree_networks = networks self.chain_mpo = chain_mpo self.terms = None if terms is None else dict(terms) - self.backend = str(backend) + self._pepsy_backend = str(backend) self.fermionic = bool(fermionic) self.symmetry = symmetry self.cutoff = float(cutoff) self.compressed = bool(compressed) + self._sites = ( + tuple(sorted(plan.node_of_qubit)) if sites is None else tuple(sites) + ) + self._site_tag_id = site_tag_id + self._upper_ind_id = upper_ind_id + self._lower_ind_id = lower_ind_id + self._node_tag_id = node_tag_id + self._canonical_region = None self.pepsy_tree_plan_signature = _tree_plan_signature(plan) if chain_mpo is not None: chain_mpo.pepsy_tree_plan_signature = self.pepsy_tree_plan_signature @@ -114,6 +227,73 @@ def __init__( chain_mpo.pepsy_tree_operator = self chain_mpo.pepsy_tree_operator_networks = self.tree_networks + @property + def backend(self): + """Return the logical Pepsy backend label for this operator.""" + return self._pepsy_backend + + @property + def pepsy_backend(self): + """Compatibility view used by Quimb's structured-network copier.""" + return self._pepsy_backend + + @property + def plan(self): + """The :class:`TreePlan` describing the operator geometry.""" + return self._plan + + @property + def node_tag_id(self): + """Format string for structural tree-node tags.""" + return self._node_tag_id + + @property + def site_ind_id(self): + """Alias for the operator's upper physical-index format.""" + return self.upper_ind_id + + @site_ind_id.setter + def site_ind_id(self, value): + self.upper_ind_id = value + + def site_ind(self, site): + """Return the ket-like physical index for ``site``.""" + return self.upper_ind(site) + + @property + def root(self): + """The structural root node id.""" + return self.plan.root + + @property + def canonical_region(self): + """The currently canonicalized connected operator region.""" + return self._canonical_region + + @property + def orthogonality_center(self): + """The single canonical node, or ``None`` for a larger region.""" + region = self.canonical_region + return next(iter(region)) if region is not None and len(region) == 1 else None + + @property + def fermionic(self): + """Whether the operator stores native fermionic arrays.""" + return self._fermionic + + @fermionic.setter + def fermionic(self, value): + self._fermionic = bool(value) + + @property + def symmetry(self): + """Native Symmray symmetry label, if present.""" + return self._symmetry + + @symmetry.setter + def symmetry(self, value): + self._symmetry = value + @classmethod def from_hamiltonian( cls, @@ -203,6 +383,168 @@ def from_terms( operator.compress(max_bond=max_bond, cutoff=cutoff) return operator + @classmethod + def from_dense( + cls, + plan, + array=None, + dims=2, + *, + tree=None, + sites=None, + tags=None, + site_tag_id="I{}", + upper_ind_id="k{}", + lower_ind_id="b{}", + node_tag_id="N{}", + **split_opts, + ): + """Build an exact tree operator from a dense matrix. + + This is the tree analogue of ``MatrixProductOperator.from_dense``. + The matrix is decomposed over the supplied ``TreePlan`` with lossless + leaf-to-root SVDs. Only the physical site ordering differs from the + chain constructor: ``sites`` labels the plan's logical qubits. + """ + if not isinstance(plan, TreePlan): + if tree is None: + raise TypeError("pass a TreePlan with `tree=` or as the first argument.") + if array is not None: + raise TypeError("dense array was supplied more than once.") + array = plan + plan = tree + elif tree is not None and tree is not plan: + raise ValueError("plan and tree specify different TreePlans.") + if array is None: + raise TypeError("TreeMPO.from_dense requires a dense matrix.") + if not isinstance(plan, TreePlan): + raise TypeError("plan must be a TreePlan.") + if sites is None: + sites = tuple(sorted(plan.node_of_qubit)) + else: + sites = tuple(int(site) for site in sites) + if sites != tuple(sorted(sites)): + raise ValueError("TreeMPO.from_dense requires sorted site labels.") + if set(sites) != set(plan.node_of_qubit): + raise ValueError( + "TreeMPO.from_dense currently requires one matrix site per tree site." + ) + if isinstance(dims, Integral): + dims = (int(dims),) * len(sites) + else: + dims = tuple(int(dim) for dim in dims) + if len(dims) != len(sites): + raise ValueError("dims must have one entry per TreePlan site.") + if np.prod(dims, dtype=int) ** 2 != np.size(array): + raise ValueError("array size does not match the supplied physical dims.") + network = _tree_operator_from_dense( + plan, + array, + sites=sites, + dims=dims, + split_opts=split_opts, + site_tag_id=site_tag_id, + upper_ind_id=upper_ind_id, + lower_ind_id=lower_ind_id, + node_tag_id=node_tag_id, + ) + if tags is not None: + network.add_tag(tags) + return cls( + plan, + network, + backend="dense", + fermionic=False, + sites=sites, + site_tag_id=site_tag_id, + upper_ind_id=upper_ind_id, + lower_ind_id=lower_ind_id, + node_tag_id=node_tag_id, + ) + + @classmethod + def from_fill_fn( + cls, + fill_fn, + plan, + bond_dim, + *, + phys_dim=2, + dtype=float, + sites=None, + tags=None, + site_tag_id="I{}", + upper_ind_id="k{}", + lower_ind_id="b{}", + node_tag_id="N{}", + ): + """Build a tree operator from a tensor filling function. + + ``fill_fn`` is called as ``fill_fn(shape)`` for each plan node, where + ``shape`` is ordered as physical upper/lower legs followed by the + node's tree bonds. A scalar ``bond_dim`` or one value per edge is + accepted through the same uniform tree convention. + """ + if sites is None: + sites = tuple(sorted(plan.node_of_qubit)) + else: + sites = tuple(sites) + network = _tree_operator_from_fill_fn( + plan, + fill_fn, + bond_dim=bond_dim, + phys_dim=phys_dim, + dtype=dtype, + site_tag_id=site_tag_id, + upper_ind_id=upper_ind_id, + lower_ind_id=lower_ind_id, + node_tag_id=node_tag_id, + ) + if tags is not None: + network.add_tag(tags) + return cls( + plan, + network, + backend="dense", + fermionic=False, + sites=sites, + site_tag_id=site_tag_id, + upper_ind_id=upper_ind_id, + lower_ind_id=lower_ind_id, + node_tag_id=node_tag_id, + ) + + @classmethod + def rand( + cls, + plan, + bond_dim, + *, + phys_dim=2, + dtype=complex, + seed=None, + **operator_opts, + ): + """Build a random dense TreeMPO with uniform virtual bond size.""" + rng = np.random.default_rng(seed) + + def fill(shape): + if np.issubdtype(np.dtype(dtype), np.complexfloating): + return ( + rng.standard_normal(shape) + + 1j * rng.standard_normal(shape) + ).astype(dtype) + return rng.standard_normal(shape).astype(dtype) + + return cls.from_fill_fn( + fill, + plan, + bond_dim=bond_dim, + phys_dim=phys_dim, + dtype=dtype, + **operator_opts, + ) + @property def tree_network(self): """Return the sole tree network, or raise for a term sum.""" @@ -213,6 +555,166 @@ def tree_network(self): ) return self.tree_networks[0] + @property + def nqubits(self): + """Number of logical physical sites in the TreePlan.""" + return self.plan.n + + @property + def top_arity(self): + """Number of virtual child bonds at the structural root.""" + return self.plan.top_arity + + @property + def max_virtual_degree(self): + """Largest number of virtual tree bonds on one operator tensor.""" + return self.plan.max_virtual_degree() + + @property + def max_tensor_rank(self): + """Largest virtual/physical leg count on one operator tensor.""" + return self.plan.max_tensor_rank() + + def node_tag(self, node): + """Return the structural tag for a TreePlan node.""" + return self._node_tag_id.format(int(node)) + + def node_tensor(self, node): + """Return a primary TTNO tensor by TreePlan node id.""" + return self.tree_networks[0][self.node_tag(node)] + + def _select_tids(self, tids, virtual=True, with_exponent=False): + """Select a structured view while keeping its primary network live.""" + selected = super()._select_tids( + tids, + virtual=virtual, + with_exponent=with_exponent, + ) + # Quimb's generic ``new(like=...)`` copies the extra properties from + # the source, including ``tree_networks``. Replace that source tuple + # with the selected view so inherited selection methods never mutate + # or inspect the original operator by accident. + selected.tree_networks = (qtn.TensorNetwork(selected, virtual=True),) + selected.chain_mpo = None + return selected + + def neighbors(self, node): + """Return the TreePlan neighbors of a structural node.""" + node = int(node) + if node not in self.plan.children: + raise ValueError(f"{node!r} is not a TreePlan node.") + return tuple(self.plan.children[node]) + ( + (self.plan.parent[node],) + if self.plan.parent.get(node) is not None + else () + ) + + def is_leaf(self, node): + """Whether ``node`` is a structural leaf.""" + return self.plan.is_leaf(int(node)) + + def parent(self, node): + """Return the parent structural node, or ``None`` at the root.""" + node = int(node) + if node not in self.plan.children: + raise ValueError(f"{node!r} is not a TreePlan node.") + return self.plan.parent.get(node) + + def children(self, node): + """Return the structural children of ``node``.""" + node = int(node) + if node not in self.plan.children: + raise ValueError(f"{node!r} is not a TreePlan node.") + return self.plan.children[node] + + def node_path(self, node1, node2): + """Return the inclusive structural path between two nodes.""" + return self.plan.node_path(int(node1), int(node2)) + + def leaf_of_qubit(self, qubit): + """Return the structural leaf carrying ``qubit``.""" + return self.plan.leaf_of_qubit[int(qubit)] + + def qubit_of_leaf(self, node): + """Return the qubit carried by a structural leaf.""" + return self.plan.qubit_of_leaf[int(node)] + + def qubit_of_node(self, node): + """Return the qubit carried by a node, or ``None`` if virtual.""" + return self.plan.qubit_of_node.get(int(node)) + + def node_of_qubit(self, qubit): + """Return the structural node carrying ``qubit``.""" + return self.plan.node_of_qubit[int(qubit)] + + def tree_distance(self, qubit1, qubit2): + """Return the structural distance between two physical sites.""" + return self.plan.tree_distance(int(qubit1), int(qubit2)) + + def steiner_nodes(self, nodes): + """Return the minimal connected subtree spanning ``nodes``.""" + return self.plan.steiner_nodes(tuple(int(node) for node in nodes)) + + def subtree_span(self, nodes): + """Return the minimal connected subtree spanning arbitrary nodes.""" + return _tree_subtree_span( + self.plan, tuple(int(node) for node in nodes), + ) + + def is_binary(self, *, allow_ternary_root=True): + """Whether this operator's tree is binary below its root.""" + return self.plan.is_binary(allow_ternary_root=allow_ternary_root) + + def bond(self, node, neighbor): + """Return the live operator bond between adjacent TreePlan nodes.""" + node = int(node) + neighbor = int(neighbor) + if neighbor not in self.neighbors(node): + raise ValueError( + f"nodes {node} and {neighbor} are not adjacent in the tree." + ) + shared = qtn.bonds( + self.node_tensor(node), self.node_tensor(neighbor), + ) + if len(shared) != 1: + raise ValueError( + f"nodes {node} and {neighbor} must share exactly one bond; " + f"found {sorted(shared)}." + ) + return next(iter(shared)) + + def validate(self): + """Validate the primary TTNO against its TreePlan geometry.""" + network = self.tree_networks[0] + for node in self.plan.nodes(): + tensor = self.node_tensor(node) + expected = set(self.neighbors(node)) + physical = self.plan.qubit_of_node.get(node) + expected_inds = { + f"_pepsy_tnno_{min(node, other)}_{max(node, other)}" + for other in expected + } + if physical is not None: + expected_inds.update(( + self.upper_ind(physical), + self.lower_ind(physical), + )) + if set(tensor.inds) != expected_inds: + raise ValueError( + f"TreeMPO node {node} has unexpected indices: " + f"{tensor.inds!r}." + ) + for node in self.plan.nodes(): + for neighbor in self.plan.children[node]: + self.bond(node, neighbor) + if set(network.outer_inds()) != { + self.upper_ind(site) for site in self.sites + } | { + self.lower_ind(site) for site in self.sites + }: + raise ValueError("TreeMPO has unexpected outer physical indices.") + return self + def max_bond(self): """Return the largest virtual bond among the tree networks.""" bonds = [] @@ -221,14 +723,458 @@ def max_bond(self): bonds.append(network.ind_size(index)) return max(bonds, default=1) - def canonicalize(self, center=None): - """Canonicalize every stored TTNO around one TreePlan node.""" + def bond_size(self, node, neighbor): + """Return the dimension of one live operator tree bond.""" + return self.node_tensor(node).ind_size(self.bond(node, neighbor)) + + def bond_sizes(self): + """Return operator bond dimensions in deterministic tree-edge order.""" + return tuple( + self.bond_size(node, child) + for node in self.plan.nodes() + for child in self.plan.children[node] + ) + + def edge_nodes(self): + """Return all directed parent-child tree edges.""" + return tuple( + (node, child) + for node in self.plan.nodes() + for child in self.plan.children[node] + ) + + @property + def L(self): + """Number of logical physical sites, as in a chain MPO.""" + return self.nsites + + @property + def cyclic(self): + """TreeMPOs are open tree networks, never cyclic chains.""" + return False + + def to_dense(self, *inds_seq, to_qarray=False, **contract_opts): + """Contract the complete operator, summing internal term networks.""" + if len(self.tree_networks) == 1: + return qtn.TensorNetworkGenOperator.to_dense( + self, + *inds_seq, + to_qarray=to_qarray, + **contract_opts, + ) + if not inds_seq: + inds_seq = (self.upper_inds_present, self.lower_inds_present) + values = [] + for network in self.tree_networks: + view = qtn.TensorNetworkGenOperator( + network, + virtual=True, + ) + view._sites = self.sites + view._site_tag_id = self.site_tag_id + view._upper_ind_id = self.upper_ind_id + view._lower_ind_id = self.lower_ind_id + values.append(view.to_dense(*inds_seq, **contract_opts)) + result = values[0] + for value in values[1:]: + result = result + value + if to_qarray: + import quimb as qu + + return qu.qarray(result) + return result + + def identity(self, *, phys_dim=None, dtype=None): + """Return the exact bond-one identity TreeMPO on this plan.""" + if phys_dim is None: + phys_dim = tuple(self.phys_dim(site) for site in self.sites) + if dtype is None: + dtype = self.dtype + network = _identity_tree_operator( + self.plan, + phys_dim=phys_dim, + dtype=dtype, + site_tag_id=self.site_tag_id, + upper_ind_id=self.upper_ind_id, + lower_ind_id=self.lower_ind_id, + node_tag_id=self.node_tag_id, + ) + return type(self)( + self.plan, + network, + backend="dense", + fermionic=False, + sites=self.sites, + site_tag_id=self.site_tag_id, + upper_ind_id=self.upper_ind_id, + lower_ind_id=self.lower_ind_id, + node_tag_id=self.node_tag_id, + ) + + def add_MPO( + self, + other, + inplace=False, + negate=False, + compress=False, + **compress_opts, + ): + """Add another matching tree operator by arbitrary-geometry direct sum.""" + if not isinstance(other, TreeMPO): + other = getattr(other, "pepsy_tree_operator", None) + if not isinstance(other, TreeMPO): + raise TypeError("other must be a TreeMPO or an annotated chain MPO.") + if self.pepsy_tree_plan_signature != other.pepsy_tree_plan_signature: + raise ValueError("TreeMPOs must use the same TreePlan.") + if self.fermionic != other.fermionic: + raise TypeError("cannot add dense and native TreeMPOs.") + if len(self.tree_networks) != len(other.tree_networks): + raise ValueError("TreeMPO term-network counts must match.") + + networks = [] + for left, right in zip(self.tree_networks, other.tree_networks): + networks.append(qtn.tensor_network_ag_sum( + left, + right, + site_tags=tuple(self.node_tag(node) for node in self.plan.nodes()), + negate=negate, + compress=compress, + **compress_opts, + )) + chain = None + if self.chain_mpo is not None and other.chain_mpo is not None: + chain = self.chain_mpo.add_MPO( + other.chain_mpo, + inplace=False, + negate=negate, + compress=compress, + **compress_opts, + ) + terms = None + if self.terms is not None and other.terms is not None: + terms = dict(self.terms) + for support, value in other.terms.items(): + if support in terms: + terms[support] = terms[support] + ((-1) if negate else 1) * value + else: + terms[support] = ((-1) if negate else 1) * value + result = type(self)( + self.plan, + tuple(networks), + chain_mpo=chain, + terms=terms, + backend=self.backend, + fermionic=self.fermionic, + symmetry=self.symmetry, + cutoff=self.cutoff, + compressed=compress, + sites=self.sites, + site_tag_id=self.site_tag_id, + upper_ind_id=self.upper_ind_id, + lower_ind_id=self.lower_ind_id, + node_tag_id=self.node_tag_id, + ) + if inplace: + self.__dict__.clear() + self.__dict__.update(result.__dict__) + return self + return result + + add_MPO_ = lambda self, other, **kwargs: self.add_MPO( # noqa: E731 + other, inplace=True, **kwargs, + ) + + def matrix_element(self, bra, ket=None): + """Return ```` for computational-basis strings.""" + if ket is None: + ket = bra + bra = tuple(int(value) for value in bra) + ket = tuple(int(value) for value in ket) + if len(bra) != self.nsites or len(ket) != self.nsites: + raise ValueError("basis configurations must match TreeMPO.nsites.") + selector = {} + for site, bra_value, ket_value in zip(self.sites, bra, ket): + selector[self.upper_ind(site)] = bra_value + selector[self.lower_ind(site)] = ket_value + value = 0.0 + for network in self.tree_networks: + value = value + network.isel(selector).contract(all) + return value + + def amplitude(self, configuration): + """Return the diagonal computational-basis matrix element.""" + return self.matrix_element(configuration) + + def singular_values(self, node, neighbor=None, *, method="svd"): + """Return singular values across one tree operator edge.""" + if neighbor is None: + try: + node, neighbor = node + except (TypeError, ValueError) as exc: + raise TypeError("singular_values needs an operator edge.") from exc + node = _tree_node_selector(self.plan, node) + neighbor = _tree_node_selector(self.plan, neighbor) + if neighbor not in self.neighbors(node): + raise ValueError("singular_values requires adjacent tree nodes.") + work = self.copy() + work.canonicalize(center=neighbor) + tensor = work.node_tensor(node) + bond = work.bond(node, neighbor) + return tensor.singular_values( + tuple(ind for ind in tensor.inds if ind != bond), + method=method, + ) + + def rand_state(self, bond_dim, **state_opts): + """Return a random :class:`TreeTensorNetwork` on the same plan.""" + from .ttn import TreeTensorNetwork + + return TreeTensorNetwork.rand(self.plan, D=bond_dim, **state_opts) + + def show(self, *, bond_dims=True, node_ids=False, color=False): + """Print a compact top-down tree drawing for this operator.""" + del color + + def render(node, prefix="", is_last=True): + qubit = self.plan.qubit_of_node.get(node) + label = f"N{node}" if node_ids else "●" + if qubit is not None: + label += f" q{qubit}" + lines = [prefix + ("└─ " if is_last else "├─ ") + label] + children = tuple(self.plan.children[node]) + for index, child in enumerate(children): + edge = self.bond_size(node, child) if bond_dims else None + edge_label = f" [{edge}]" if edge is not None else "" + child_lines = render( + child, + prefix + (" " if is_last else "│ "), + index == len(children) - 1, + ) + child_lines[0] = child_lines[0] + edge_label + lines.extend(child_lines) + return lines + + lines = render(self.plan.root, "", True) + print("\n".join(lines)) + + def canonicalize(self, center=None, *, inplace=True): + """Canonicalize every stored TTNO around one TreePlan node. + + This is the tree equivalent of an MPO mixed-canonical gauge. The + default is inplace, matching Quimb's MPO canonicalization methods; + pass ``inplace=False`` to obtain an independent operator. + """ if center is None: center = self.plan.root - for network in self.tree_networks: - _canonicalize_tree_operator(network, self.plan, center) + center = _tree_node_selector(self.plan, center) + target = self if inplace else self.copy() + for network in target.tree_networks: + _canonicalize_tree_operator(network, target.plan, center) + target._canonical_region = frozenset({center}) + return target + + def canonicalize_(self, center=None): + """Inplace alias for :meth:`canonicalize`.""" + return self.canonicalize(center=center, inplace=True) + + canonize = canonicalize_ + + def invalidate_canonical_form(self): + """Forget operator gauge metadata after an unmanaged tensor edit.""" + self._canonical_region = None return self + def isometry_direction(self, node): + """Return the neighbour receiving a node's canonical QR factor.""" + node = _tree_node_selector(self.plan, node) + tensor = self.node_tensor(node) + if tensor.left_inds is None: + return None + right_inds = [ind for ind in tensor.inds if ind not in tensor.left_inds] + if len(right_inds) != 1: + return None + for neighbor in self.neighbors(node): + if right_inds[0] == self.bond(node, neighbor): + return neighbor + return None + + def isometry_map(self): + """Return the live QR orientation map for all TreePlan nodes.""" + return { + node: self.isometry_direction(node) + for node in self.plan.nodes() + } + + def is_subtree_canonical_form(self, nodes=None, *, span=False): + """Check the lossless QR metadata around a connected operator region.""" + if nodes is None: + region = self.canonical_region + if region is None: + return False + else: + region = _tree_region_selector(self.plan, nodes) if span else frozenset( + _tree_node_selector(self.plan, node) for node in nodes + ) + if _tree_subtree_span(self.plan, region) != region: + return False + for node in self.plan.nodes(): + if node in region: + continue + path = min( + ( + self.plan.node_path(node, target) + for target in region + ), + key=len, + ) + if self.isometry_direction(node) != path[1]: + return False + return True + + def is_canonical_form(self, center=None): + """Check whether the operator has a one-node canonical region.""" + if center is None: + center = self.orthogonality_center + if center is None: + return False + return self.is_subtree_canonical_form((center,)) + + def shift_orthogonality_center(self, current, new): + """Move the operator QR centre to another TreePlan node.""" + del current + return self.canonicalize(center=new, inplace=True) + + def calc_current_orthog_center(self): + """Return the current operator canonical region bounds.""" + region = self.canonical_region + if not region: + return None + ordered = sorted(region) + return ordered[0], ordered[-1] + + def left_canonicalize(self, *, center=None, inplace=False, **kwargs): + """MPO-compatible alias for a root-oriented tree QR sweep.""" + del kwargs + return self.canonicalize( + center=self.plan.root if center is None else center, + inplace=inplace, + ) + + left_canonicalize_ = lambda self, **kwargs: self.left_canonicalize( # noqa: E731 + inplace=True, **kwargs, + ) + left_canonize = left_canonicalize_ + + def right_canonicalize(self, *, center=None, inplace=False, **kwargs): + """MPO-compatible alias for a root-oriented tree QR sweep.""" + del kwargs + return self.canonicalize( + center=self.plan.root if center is None else center, + inplace=inplace, + ) + + right_canonicalize_ = lambda self, **kwargs: self.right_canonicalize( # noqa: E731 + inplace=True, **kwargs, + ) + right_canonize = right_canonicalize_ + + def compress_site(self, node, *, max_bond=None, cutoff=None, **kwargs): + """Compress all tree bonds consistently around ``node``.""" + del kwargs + self.canonicalize(center=node) + return self.compress(max_bond=max_bond, cutoff=cutoff) + + def left_compress(self, *, max_bond=None, cutoff=None, **kwargs): + """Tree analogue of a left-to-right MPO compression sweep.""" + del kwargs + return self.compress(max_bond=max_bond, cutoff=cutoff) + + def right_compress(self, *, max_bond=None, cutoff=None, **kwargs): + """Tree analogue of a right-to-left MPO compression sweep.""" + del kwargs + return self.compress(max_bond=max_bond, cutoff=cutoff) + + def canonize_around( + self, tags, which="all", *, inplace=False, **canonize_opts, + ): + """Quimb-style alias for TreePlan-centered TTNO canonicalization. + + Tree operators have a rooted geometry rather than a one-dimensional + tag interval, so the supported target is one TreePlan node. The + additional Quimb options are accepted for API familiarity and are + intentionally ignored after validating that the target is a node. + """ + del which, canonize_opts + if isinstance(tags, (tuple, list, set, frozenset)): + if len(tags) == 0: + raise ValueError("TreeMPO.canonize_around needs one node.") + if len(tags) != 1: + target = self if inplace else self.copy() + region = _tree_region_selector(target.plan, tags) + for network in target.tree_networks: + _canonicalize_tree_operator_region( + network, target.plan, region, + ) + target._canonical_region = region + return target + target = self if inplace else self.copy() + return target.canonicalize(center=tags, inplace=True) + + def canonize_around_(self, tags, **kwargs): + """In-place Quimb-style alias for :meth:`canonize_around`.""" + kwargs["inplace"] = True + return self.canonize_around(tags, **kwargs) + + def canonize_between( + self, tags1, tags2, *, inplace=False, absorb="right", **canonize_opts, + ): + """Canonicalize the operator exterior to a TreePlan path. + + A path is the tree analogue of the mixed-canonical interval used by + an MPS. ``absorb`` and other Quimb gauge options are accepted for API + compatibility; the lossless native QR policy controls the operation. + """ + del absorb, canonize_opts + node1 = _tree_node_selector(self.plan, tags1) + node2 = _tree_node_selector(self.plan, tags2) + region = frozenset(self.plan.node_path(node1, node2)) + target = self if inplace else self.copy() + for network in target.tree_networks: + _canonicalize_tree_operator_region(network, target.plan, region) + target._canonical_region = region + return target + + def canonize_between_(self, tags1, tags2, **kwargs): + """In-place alias for :meth:`canonize_between`.""" + kwargs["inplace"] = True + return self.canonize_between(tags1, tags2, **kwargs) + + def compress_between( + self, tags1, tags2, max_bond=None, cutoff=1e-10, **compress_opts, + ): + """Quimb-style compression entry point for a tree operator. + + A TreeMPO compression is a global leaf-to-root sweep so every edge + sees the complete operator sum. ``tags1`` and ``tags2`` identify an + adjacent TreePlan edge for validation; the configured sweep then + compresses all TreePlan bonds consistently. + """ + inplace = compress_opts.pop("inplace", True) + del compress_opts + node1 = _tree_node_selector(self.plan, tags1) + node2 = _tree_node_selector(self.plan, tags2) + if node2 not in self.neighbors(node1): + raise ValueError( + "TreeMPO.compress_between requires adjacent TreePlan nodes." + ) + target = self if inplace else self.copy() + return target.compress(max_bond=max_bond, cutoff=cutoff) + + def compress_between_(self, tags1, tags2, **kwargs): + """In-place alias for :meth:`compress_between`.""" + kwargs["inplace"] = True + return self.compress_between(tags1, tags2, **kwargs) + def compress(self, *, max_bond=None, cutoff=None): """Compress the TTNO on every TreePlan edge with native SVD.""" if cutoff is None: @@ -245,14 +1191,29 @@ def compress(self, *, max_bond=None, cutoff=None): self.cutoff = cutoff self.compressed = True self.pepsy_compression_report = reports[0] if len(reports) == 1 else reports + self._canonical_region = None return self - def copy(self): - """Copy the operator and both of its optional representations.""" - chain_mpo = None if self.chain_mpo is None else self.chain_mpo.copy() + def copy(self, virtual=False, deep=False, *, conj=False, transpose=False): + """Copy the operator and both of its optional representations. + + The signature follows Quimb's tensor-network ``copy`` API. ``virtual`` + keeps tensor data shared while copying the network structure; ``deep`` + requests independent numeric data as in the underlying Quimb views. + ``conj`` and ``transpose`` are accepted as convenient MPO-compatible + view operations. + """ + chain_mpo = ( + None + if self.chain_mpo is None + else self.chain_mpo.copy(virtual=virtual, deep=deep) + ) copied = type(self)( self.plan, - tuple(network.copy() for network in self.tree_networks), + tuple( + network.copy(virtual=virtual, deep=deep) + for network in self.tree_networks + ), chain_mpo=chain_mpo, terms=self.terms, backend=self.backend, @@ -260,6 +1221,11 @@ def copy(self): symmetry=self.symmetry, cutoff=self.cutoff, compressed=self.compressed, + sites=self.sites, + site_tag_id=self.site_tag_id, + upper_ind_id=self.upper_ind_id, + lower_ind_id=self.lower_ind_id, + node_tag_id=self.node_tag_id, ) if chain_mpo is not None: chain_mpo.pepsy_tree_plan_signature = copied.pepsy_tree_plan_signature @@ -268,8 +1234,115 @@ def copy(self): ) chain_mpo.pepsy_tree_operator = copied chain_mpo.pepsy_tree_operator_networks = copied.tree_networks + if hasattr(self, "pepsy_compression_report"): + copied.pepsy_compression_report = self.pepsy_compression_report + copied._canonical_region = self.canonical_region + if transpose: + copied._transpose_operator_inplace() + if conj: + copied.conj(inplace=True) return copied + def _transpose_operator_inplace(self): + """Transpose every local upper/lower physical pair in place.""" + for network in self.tree_networks: + for tensor in network: + physical_axes = [] + for site in self.sites: + upper = self.upper_ind(site) + lower = self.lower_ind(site) + if upper in tensor.inds and lower in tensor.inds: + physical_axes.append(( + tensor.inds.index(upper), + tensor.inds.index(lower), + )) + if not physical_axes: + continue + permutation = list(range(tensor.ndim)) + for upper_axis, lower_axis in physical_axes: + permutation[upper_axis], permutation[lower_axis] = ( + permutation[lower_axis], permutation[upper_axis] + ) + tensor.modify(data=ar.do("transpose", tensor.data, permutation)) + if self.chain_mpo is not None: + for tensor in self.chain_mpo: + axes = [] + for site in self.sites: + upper = self.upper_ind(site) + lower = self.lower_ind(site) + if upper in tensor.inds and lower in tensor.inds: + axes.append(( + tensor.inds.index(upper), + tensor.inds.index(lower), + )) + permutation = list(range(tensor.ndim)) + for upper_axis, lower_axis in axes: + permutation[upper_axis], permutation[lower_axis] = ( + permutation[lower_axis], permutation[upper_axis] + ) + if axes: + tensor.modify(data=ar.do("transpose", tensor.data, permutation)) + return self + + def conj( + self, + mangle_inner=False, + output_inds=None, + phase_dual=True, + inplace=False, + ): + """Conjugate every stored tree operator like a Quimb operator view.""" + if inplace: + for network in self.tree_networks: + network.conj( + mangle_inner=mangle_inner, + output_inds=output_inds, + phase_dual=phase_dual, + inplace=True, + ) + if self.chain_mpo is not None: + self.chain_mpo.conj( + mangle_inner=mangle_inner, + output_inds=output_inds, + phase_dual=phase_dual, + inplace=True, + ) + return self + + networks = tuple( + network.conj( + mangle_inner=mangle_inner, + output_inds=output_inds, + phase_dual=phase_dual, + ) + for network in self.tree_networks + ) + chain_mpo = ( + None + if self.chain_mpo is None + else self.chain_mpo.conj( + mangle_inner=mangle_inner, + output_inds=output_inds, + phase_dual=phase_dual, + ) + ) + return type(self)( + self.plan, + networks, + chain_mpo=chain_mpo, + terms=self.terms, + backend=self.backend, + fermionic=self.fermionic, + symmetry=self.symmetry, + cutoff=self.cutoff, + compressed=self.compressed, + sites=self.sites, + site_tag_id=self.site_tag_id, + upper_ind_id=self.upper_ind_id, + lower_ind_id=self.lower_ind_id, + node_tag_id=self.node_tag_id, + ) + def expectation(self, state, *, normalized=True, optimize="auto"): """Evaluate ```` in one public operation.""" import quimb.tensor as qtn # pylint: disable=import-outside-toplevel @@ -762,6 +1835,224 @@ def _tree_operator_peel_order(plan, nodes): return tuple(order), next(iter(remaining)) +def _tree_operator_from_dense( + plan, + array, + *, + sites, + dims, + split_opts, + site_tag_id, + upper_ind_id, + lower_ind_id, + node_tag_id, +): + """Decompose a dense matrix exactly across a TreePlan.""" + data = ar.do("reshape", array, tuple(dims) + tuple(dims)) + upper = tuple(upper_ind_id.format(site) for site in sites) + lower = tuple(lower_ind_id.format(site) for site in sites) + blob = qtn.Tensor(data, inds=upper + lower) + + owned = {node: [] for node in plan.nodes()} + for site in sites: + owned[plan.node_of_qubit[site]].extend(( + upper_ind_id.format(site), + lower_ind_id.format(site), + )) + factors = {} + peel_order, hub = _tree_operator_peel_order(plan, set(plan.nodes())) + opts = dict(split_opts) + opts.setdefault("method", "svd") + opts.setdefault("absorb", "right") + opts.setdefault("cutoff", 0.0) + opts.setdefault("get", "tensors") + + for node, neighbor in peel_order: + bond_ind = f"_pepsy_tnno_{min(node, neighbor)}_{max(node, neighbor)}" + left, blob = blob.split( + left_inds=tuple(owned[node]), + right_inds=tuple(ind for ind in blob.inds if ind not in owned[node]), + bond_ind=bond_ind, + **opts, + ) + factors[node] = left + owned[neighbor].append(bond_ind) + factors[hub] = blob + + tensors = [] + for node in plan.nodes(): + tensor = factors[node] + qubit = plan.qubit_of_node.get(node) + neighbors = tuple(plan.children[node]) + ( + (plan.parent[node],) if plan.parent.get(node) is not None else () + ) + desired = [ + *( ( + upper_ind_id.format(qubit), + lower_ind_id.format(qubit), + ) if qubit is not None else () ), + *( + f"_pepsy_tnno_{min(node, neighbor)}_{max(node, neighbor)}" + for neighbor in neighbors + ), + ] + tensor = tensor.transpose(*desired) + tensor.add_tag(node_tag_id.format(node)) + if qubit is not None: + tensor.add_tag(site_tag_id.format(qubit)) + tensors.append(tensor) + + network = qtn.TensorNetwork(tensors) + network.pepsy_tree_operator_kind = "dense_tree_tnno" + network.pepsy_tree_operator_is_ttno = True + network.pepsy_tree_operator_bond = max( + (network.ind_size(index) for index in network.inner_inds()), + default=1, + ) + network.pepsy_tree_operator_raw_bond = network.pepsy_tree_operator_bond + return network + + +def _tree_operator_from_fill_fn( + plan, + fill_fn, + *, + bond_dim, + phys_dim, + dtype, + site_tag_id, + upper_ind_id, + lower_ind_id, + node_tag_id, +): + """Build a regular dense TTNO from local filled tensors.""" + if isinstance(bond_dim, Integral): + edge_dims = { + (min(node, child), max(node, child)): int(bond_dim) + for node in plan.nodes() + for child in plan.children[node] + } + else: + edge_values = tuple(int(value) for value in bond_dim) + edges = tuple( + (node, child) + for node in plan.nodes() + for child in plan.children[node] + ) + if len(edge_values) != len(edges): + raise ValueError("bond_dim must be one value per TreePlan edge.") + edge_dims = { + (min(node, child), max(node, child)): value + for (node, child), value in zip(edges, edge_values) + } + + if isinstance(phys_dim, Integral): + physical_dims = {site: int(phys_dim) for site in plan.node_of_qubit} + else: + values = tuple(int(value) for value in phys_dim) + sites = tuple(sorted(plan.node_of_qubit)) + if len(values) != len(sites): + raise ValueError("phys_dim must have one value per tree site.") + physical_dims = dict(zip(sites, values)) + + tensors = [] + for node in plan.nodes(): + qubit = plan.qubit_of_node.get(node) + neighbors = tuple(plan.children[node]) + ( + (plan.parent[node],) if plan.parent.get(node) is not None else () + ) + inds = [ + *(( + upper_ind_id.format(qubit), + lower_ind_id.format(qubit), + ) if qubit is not None else ()), + *( + f"_pepsy_tnno_{min(node, neighbor)}_{max(node, neighbor)}" + for neighbor in neighbors + ), + ] + shape = [ + *( (physical_dims[qubit], physical_dims[qubit]) + if qubit is not None else () ), + *( + edge_dims[(min(node, neighbor), max(node, neighbor))] + for neighbor in neighbors + ), + ] + try: + data = fill_fn(tuple(shape)) + except TypeError: + data = fill_fn(node, tuple(shape)) + tensor = qtn.Tensor(np.asarray(data, dtype=dtype), inds=inds) + tensor.add_tag(node_tag_id.format(node)) + if qubit is not None: + tensor.add_tag(site_tag_id.format(qubit)) + tensors.append(tensor) + + network = qtn.TensorNetwork(tensors) + network.pepsy_tree_operator_kind = "dense_tree_tnno" + network.pepsy_tree_operator_is_ttno = True + return network + + +def _identity_tree_operator( + plan, + *, + phys_dim=2, + dtype=complex, + site_tag_id="I{}", + upper_ind_id="k{}", + lower_ind_id="b{}", + node_tag_id="N{}", +): + """Build an exact bond-one identity TTNO.""" + if isinstance(phys_dim, Integral): + physical_dims = {site: int(phys_dim) for site in plan.node_of_qubit} + else: + sites = tuple(sorted(plan.node_of_qubit)) + values = tuple(int(value) for value in phys_dim) + if len(values) != len(sites): + raise ValueError("phys_dim must have one value per tree site.") + physical_dims = dict(zip(sites, values)) + + tensors = [] + for node in plan.nodes(): + qubit = plan.qubit_of_node.get(node) + neighbors = tuple(plan.children[node]) + ( + (plan.parent[node],) if plan.parent.get(node) is not None else () + ) + inds = [ + *(( + upper_ind_id.format(qubit), + lower_ind_id.format(qubit), + ) if qubit is not None else ()), + *( + f"_pepsy_tnno_{min(node, neighbor)}_{max(node, neighbor)}" + for neighbor in neighbors + ), + ] + shape = [ + *((physical_dims[qubit], physical_dims[qubit]) + if qubit is not None else ()), + *(1 for _ in neighbors), + ] + data = np.zeros(shape, dtype=dtype) + if qubit is None: + data[...] = 1 + else: + data[(slice(None), slice(None)) + (0,) * len(neighbors)] = np.eye( + physical_dims[qubit], dtype=dtype, + ) + tensor = qtn.Tensor(data, inds=inds, tags=[node_tag_id.format(node)]) + if qubit is not None: + tensor.add_tag(site_tag_id.format(qubit)) + tensors.append(tensor) + network = qtn.TensorNetwork(tensors) + network.pepsy_tree_operator_kind = "dense_tree_identity" + network.pepsy_tree_operator_is_ttno = True + return network + + def _native_tree_term_network( plan, term, support, *, symmetry, cutoff=1e-12, dtype=None, ): @@ -1239,17 +2530,37 @@ def _tree_operator_qr(tensor, *, left_inds, bond_ind): def _canonicalize_tree_operator(network, plan, center): """Canonicalize a tree operator by lossless QR from leaves to center.""" + return _canonicalize_tree_operator_region(network, plan, {center}) + + +def _canonicalize_tree_operator_region(network, plan, region): + """Canonicalize the complement of a connected tree region inwards.""" import quimb.tensor as qtn # pylint: disable=import-outside-toplevel - if center not in plan.children: - raise ValueError(f"operator canonicalization center {center!r} is invalid.") + region = frozenset(region) + if not region or not region.issubset(plan.children): + raise ValueError("operator canonicalization region is invalid.") + if _tree_subtree_span(plan, region) != region: + raise ValueError("operator canonicalization region must be connected.") + + def distance_to_region(node): + return min(len(plan.node_path(node, target)) for target in region) + order = sorted( - (node for node in plan.nodes() if node != center), - key=lambda node: len(plan.node_path(node, center)), + (node for node in plan.nodes() if node not in region), + key=distance_to_region, reverse=True, ) for node in order: - neighbor = plan.node_path(node, center)[1] + path = min( + ( + plan.node_path(node, target) + for target in region + if target != node + ), + key=len, + ) + neighbor = path[1] tensor = _tree_operator_tensor(network, node) target = _tree_operator_tensor(network, neighbor) bond = _tree_operator_bond(network, plan, node, neighbor) @@ -1269,7 +2580,9 @@ def _canonicalize_tree_operator(network, plan, center): inds=merged.inds, left_inds=None, ) - network.pepsy_tree_operator_center = int(center) + network.pepsy_tree_operator_center = ( + next(iter(region)) if len(region) == 1 else None + ) network.pepsy_tree_operator_canonical = True return network @@ -1829,6 +3142,7 @@ def _annotate_tree_mpo( compressed=False, cutoff=1e-12, max_bond=None, + to_backend=None, ): """Attach a public :class:`TreeMPO` to a compatibility chain MPO.""" mpo.pepsy_tree_plan_signature = _tree_plan_signature(plan) @@ -1860,6 +3174,16 @@ def _annotate_tree_mpo( mpo.pepsy_tree_operator_networks = operator.tree_networks if compressed: operator.compress(max_bond=max_bond, cutoff=cutoff) + if to_backend is not None: + # Move the primary TreeMPO representation onto the requested backend + # so ``TreeMPO.expectation`` contracts against a matching tree state. + # The compatibility chain MPO is converted separately by the caller; + # here the structured operator networks are moved in place after any + # numpy-side compression completes. + from ...tensors.symmetric import _apply_to_tensor_network_arrays + + for network in operator.tree_networks: + _apply_to_tensor_network_arrays(network, to_backend) return mpo @@ -2002,6 +3326,7 @@ def tree_mpo( compressed=compress, cutoff=cutoff, max_bond=max_bond, + to_backend=to_backend, ) return result tree_operator = _build_tree_operator( @@ -2027,4 +3352,49 @@ def tree_mpo( compressed=compress, cutoff=cutoff, max_bond=max_bond, + to_backend=to_backend, + ) + + +def build_tree_operator( + plan, + hamiltonian, + *, + max_bond=None, + cutoff=1e-12, + upper_ind_id="k{}", + lower_ind_id="b{}", + site_tag_id="I{}", + compress=True, + dtype=None, + fermionic=True, + charge_sectors=False, + to_backend=None, +): + """Build the canonical native :class:`TreeMPO` for a ``TreePlan``. + + The compatibility :func:`tree_mpo` builder returns the ordinary linear + chain MPO and attaches this tree operator to it. This function returns + only the tree-native operator, keeping the two tensor-network geometries + explicit. With ``charge_sectors=True`` it returns ``{charge: TreeMPO}``. + """ + built = tree_mpo( + plan, + hamiltonian, + max_bond=max_bond, + cutoff=cutoff, + upper_ind_id=upper_ind_id, + lower_ind_id=lower_ind_id, + site_tag_id=site_tag_id, + compress=compress, + dtype=dtype, + fermionic=fermionic, + charge_sectors=charge_sectors, + to_backend=to_backend, ) + if isinstance(built, dict): + return { + charge: mpo.pepsy_tree_operator + for charge, mpo in built.items() + } + return built.pepsy_tree_operator diff --git a/src/pepsy/optimizers/tree/optimizer.py b/src/pepsy/optimizers/tree/optimizer.py index 59c2a18..4cb0580 100644 --- a/src/pepsy/optimizers/tree/optimizer.py +++ b/src/pepsy/optimizers/tree/optimizer.py @@ -4145,7 +4145,8 @@ def apply_submpo(self, submpo, where, *, max_bond=None, cutoff=None): if state_native: raise TypeError( "native fermionic TreeTensorNetwork requires a native " - "Symmray MPO. Build it with tree_mpo(..., fermionic=True) " + "Symmray MPO. Build it with build_tree_operator(...) or " + "tree_mpo(..., fermionic=True) " "or supply a model-native MPO." ) raise TypeError( diff --git a/src/pepsy/tensors/symmetric.py b/src/pepsy/tensors/symmetric.py index ea49e81..9d7e187 100644 --- a/src/pepsy/tensors/symmetric.py +++ b/src/pepsy/tensors/symmetric.py @@ -10223,8 +10223,9 @@ def hamiltonian( def build_mpo( self, - terms_or_edges, + terms_or_edges=None, *, + hamiltonian=None, L=None, mapper=None, idx2coo=None, @@ -10236,31 +10237,52 @@ def build_mpo( lower_ind_id="b{}", site_tag_id="I{}", dtype=None, - fermionic=False, + fermionic=True, charge_sectors=False, to_backend=None, **params, ): - """Build a Symmray MPO directly from this fermion model. - - This is the model-facing shorthand for - ``fermion.hamiltonian(...).to_mpo(...)``. The resulting MPO uses the - same symmetry as :meth:`SymHamiltonian.to_mpo`. By default this - model-facing helper builds the current bosonic/Jordan-Wigner - compatibility MPO and can be evolved with ``jw_trotter_gates``. - Pass ``fermionic=True`` to select native graded ``FermionicArray`` - construction; this is equivalent to calling :meth:`to_mpo` with the - same model parameters. - ``t``, ``U``/``V``, and ``mu`` remain explicit build parameters and are - forwarded to :meth:`hamiltonian`. + """Build the model-facing one-dimensional MPO. + + This is the canonical ``Fermion`` entry point for a chain MPO. Pass + either model terms or edges in ``terms_or_edges`` or an existing + :class:`SymHamiltonian` with ``hamiltonian=``. Native graded + ``FermionicArray`` tensors are built by default; pass + ``fermionic=False`` for the explicit Jordan--Wigner compatibility + MPO. + + ``to_mpo`` is retained as a compatibility alias of this method. + ``t``, ``U``/``V``, and ``mu`` remain explicit build parameters and + are forwarded to :meth:`hamiltonian`. """ to_backend = self.to_backend if to_backend is None else to_backend - hamiltonian = self.hamiltonian( - terms_or_edges, - to_backend=to_backend, - **params, - ) - return hamiltonian.to_mpo( + if hamiltonian is not None: + if terms_or_edges is not None: + raise TypeError( + "Pass either terms_or_edges or hamiltonian, not both." + ) + if not isinstance(hamiltonian, SymHamiltonian): + raise TypeError("hamiltonian must be a SymHamiltonian instance.") + target = hamiltonian + elif isinstance(terms_or_edges, SymHamiltonian): + target = terms_or_edges + else: + if terms_or_edges is None: + raise TypeError("build_mpo requires terms_or_edges or hamiltonian.") + target = self.hamiltonian( + terms_or_edges, + to_backend=to_backend, + **params, + ) + params = {} + + if params: + names = ", ".join(sorted(params)) + raise TypeError( + "Model parameters cannot be supplied with an existing " + f"SymHamiltonian: {names}." + ) + return target.to_mpo( L=L, mapper=mapper, idx2coo=idx2coo, @@ -10277,6 +10299,10 @@ def build_mpo( to_backend=to_backend, ) + # ``to_mpo`` was the original model-facing name. Keep one implementation + # so the two spellings cannot drift in defaults or supported arguments. + to_mpo = build_mpo + def build_pepo( self, terms_or_edges=None, @@ -10322,89 +10348,7 @@ def build_pepo( **params, ) - def to_mpo( - self, - terms_or_edges=None, - *, - hamiltonian=None, - L=None, - mapper=None, - idx2coo=None, - coo2idx=None, - max_bond=None, - cutoff=1e-12, - compress=True, - upper_ind_id="k{}", - lower_ind_id="b{}", - site_tag_id="I{}", - dtype=None, - fermionic=True, - charge_sectors=False, - to_backend=None, - **params, - ): - """Build a native graded MPO from a fermionic term collection. - - ``terms_or_edges`` may be lattice edges for the built-in model or a - mapping such as ``{(0, 2, 4): fermion.operator_term(...)}``. The - latter supports arbitrary homogeneous-charge term support, including - non-contiguous sites. Pass an existing :class:`SymHamiltonian` with - ``hamiltonian=`` when the terms have already been assembled. - Set ``charge_sectors=True`` to return one native MPO per charge for a - mixed-charge term collection. - - The native path is selected by default and returns MPO tensors backed - by Symmray ``FermionicArray`` objects. Set ``fermionic=False`` to - request the explicit Jordan--Wigner compatibility MPO instead. - - ``to_backend`` overrides the backend configured on this ``Fermion`` - instance for both the Hamiltonian terms and the returned MPO blocks. - """ - to_backend = self.to_backend if to_backend is None else to_backend - if hamiltonian is not None: - if terms_or_edges is not None: - raise TypeError( - "Pass either terms_or_edges or hamiltonian, not both." - ) - if not isinstance(hamiltonian, SymHamiltonian): - raise TypeError("hamiltonian must be a SymHamiltonian instance.") - target = hamiltonian - elif isinstance(terms_or_edges, SymHamiltonian): - target = terms_or_edges - else: - if terms_or_edges is None: - raise TypeError("to_mpo requires terms_or_edges or hamiltonian.") - target = self.hamiltonian( - terms_or_edges, - to_backend=to_backend, - **params, - ) - params = {} - - if params: - names = ", ".join(sorted(params)) - raise TypeError( - "Model parameters cannot be supplied with an existing " - f"SymHamiltonian: {names}." - ) - return target.to_mpo( - L=L, - mapper=mapper, - idx2coo=idx2coo, - coo2idx=coo2idx, - max_bond=max_bond, - cutoff=cutoff, - compress=compress, - upper_ind_id=upper_ind_id, - lower_ind_id=lower_ind_id, - site_tag_id=site_tag_id, - dtype=dtype, - fermionic=fermionic, - charge_sectors=charge_sectors, - to_backend=to_backend, - ) - - def to_tree_mpo( + def build_tree_operator( self, terms_or_edges=None, *, @@ -10420,19 +10364,22 @@ def to_tree_mpo( to_backend=None, **params, ): - """Build a :class:`pepsy.TreeMPO` for a selected ``TreePlan``. + """Build the native :class:`pepsy.TreeMPO` for a selected plan. ``tree`` and ``plan`` are aliases. The returned object exposes the optional linear representation as ``.chain_mpo`` and the TreePlan representation through ``.tree_networks`` and ``.expectation``. Native ``fermionic=True`` keeps Symmray's graded tensors intact for U1, U1U1, and other supported symmetries. + + This is the canonical ``Fermion`` tree-operator entry point. + ``to_tree_mpo`` and ``build_tree_mpo`` remain compatibility aliases. """ if tree is not None and plan is not None: raise TypeError("pass only one of tree= or plan=") plan = tree if tree is not None else plan if plan is None: - raise TypeError("to_tree_mpo requires tree= or plan=.") + raise TypeError("build_tree_operator requires tree= or plan=.") if hamiltonian is not None: if terms_or_edges is not None: raise TypeError( @@ -10446,7 +10393,7 @@ def to_tree_mpo( else: if terms_or_edges is None: raise TypeError( - "to_tree_mpo requires terms_or_edges or hamiltonian." + "build_tree_operator requires terms_or_edges or hamiltonian." ) target = self.hamiltonian( terms_or_edges, @@ -10460,9 +10407,9 @@ def to_tree_mpo( "Model parameters cannot be supplied with an existing " f"SymHamiltonian: {names}." ) - from ..optimizers.tree import tree_mpo + from ..optimizers.tree import build_tree_operator - built = tree_mpo( + return build_tree_operator( plan, target, max_bond=max_bond, @@ -10473,14 +10420,10 @@ def to_tree_mpo( charge_sectors=charge_sectors, to_backend=to_backend, ) - if isinstance(built, dict): - return { - charge: mpo.pepsy_tree_operator - for charge, mpo in built.items() - } - return built.pepsy_tree_operator - build_tree_mpo = to_tree_mpo + # Keep the historical spellings as aliases of the one canonical builder. + to_tree_mpo = build_tree_operator + build_tree_mpo = build_tree_operator def to_pepo( self, diff --git a/tests/test_optimize_mpo.py b/tests/test_optimize_mpo.py index 362f5c1..5745568 100644 --- a/tests/test_optimize_mpo.py +++ b/tests/test_optimize_mpo.py @@ -551,6 +551,7 @@ def test_mpo_optimizer_adapts_long_range_native_gate_to_jw_symmray_mpo(): t=1.0, U=2.0, mu=0.1, + fermionic=False, ) gates = fermion.strang_gate_stream( [(0, 3)], @@ -583,6 +584,7 @@ def test_mpo_optimizer_handles_fermion_symmray_mpo_and_native_gate_stream(mode): t=1.0, U=2.0, mu=0.1, + fermionic=False, max_bond=16, cutoff=1e-12, ) @@ -610,7 +612,9 @@ def test_fermion_build_mpo_and_ham_tn_adapter_preserve_symmetry(): edges = [(0, 1), (1, 2)] builder = py.ham_tn(Lx=3, Ly=1, data_type="complex128") - direct = fermion.build_mpo(edges, L=3, t=1.0, U=0.0, mu=0.0) + direct = fermion.build_mpo( + edges, L=3, t=1.0, U=0.0, mu=0.0, fermionic=True, + ) adapted = builder.build_mpo( fermion=fermion, edges=edges, @@ -618,6 +622,7 @@ def test_fermion_build_mpo_and_ham_tn_adapter_preserve_symmetry(): t=1.0, U=0.0, mu=0.0, + fermionic=True, ) positional = builder.build_mpo( fermion, @@ -625,16 +630,24 @@ def test_fermion_build_mpo_and_ham_tn_adapter_preserve_symmetry(): t=1.0, U=0.0, mu=0.0, + fermionic=True, ) assert direct.L == adapted.L == positional.L == 3 - assert all(type(tensor.data).__name__ == "U1U1Array" for tensor in direct) - assert all(type(tensor.data).__name__ == "U1U1Array" for tensor in adapted) - assert all(type(tensor.data).__name__ == "U1U1Array" for tensor in positional) + assert all( + type(tensor.data).__name__ == "U1U1FermionicArray" for tensor in direct + ) + assert all( + type(tensor.data).__name__ == "U1U1FermionicArray" for tensor in adapted + ) + assert all( + type(tensor.data).__name__ == "U1U1FermionicArray" + for tensor in positional + ) -def test_build_mpo_can_select_native_fermionic_construction(): - """The model-facing builder exposes the native MPO path explicitly.""" +def test_build_mpo_defaults_to_native_and_to_mpo_is_its_alias(): + """The model-facing builder has one native default and one alias.""" pytest.importorskip("symmray") fermion = py.Fermion(spinful=True, symmetry="U1U1") @@ -644,7 +657,6 @@ def test_build_mpo_can_select_native_fermionic_construction(): t=1.0, U=0.0, mu=0.0, - fermionic=True, compress=False, ) direct = fermion.to_mpo( @@ -660,6 +672,7 @@ def test_build_mpo_can_select_native_fermionic_construction(): type(tensor.data).__name__ == "U1U1FermionicArray" for tensor in native ) + assert type(fermion).to_mpo is type(fermion).build_mpo assert native.to_dense().allclose(direct.to_dense()) @@ -674,6 +687,7 @@ def test_mpo_optimizer_explicit_compress_handles_empty_symmray_stream(): t=1.0, U=0.0, mu=0.0, + fermionic=False, compress=False, ) raw_bond = mpo.max_bond() diff --git a/tests/test_public_api.py b/tests/test_public_api.py index a6f9bfc..363f340 100644 --- a/tests/test_public_api.py +++ b/tests/test_public_api.py @@ -60,7 +60,7 @@ def test_tree_optimizers_are_available_from_high_level_api(): "SimulatorCandidate", "SimulatorPlan", "SimulatorPlanner", "recommend_simulator", "TreeEnergyOptimizer", "TreeLayoutFinder", - "TreeMPO", "TreeOptimizer", "tree_mpo", + "TreeMPO", "TreeOptimizer", "build_tree_operator", "tree_mpo", "TreePlan", "TreeStabOptimizer", "TreeTensorNetwork", @@ -137,7 +137,7 @@ def test_internal_symbols_not_exported(): "PepsEnergyOptimizer", "PepsOptimizer", "SimpleUpdateGen", "SymDMRG2", "PEPSSampleResult", "PepsBpSampler", "CoherentCrosstalkModel", "compile_stim_circuit", "run_coalesced_noisy_shots", "run_coalesced_stim_shots", "run_coalesced_trajectory_shots", "run_noisy_shots", "run_stabilizer_mps_stream", "run_stabilizer_tree_stream", "run_stim_shots", "run_trajectory_shots", "sample_coalesced_bits", "sample_noisy_gate_stream", "sample_noisy_gate_streams", "sample_stim_circuit", "sample_stim_circuits", "sample_trajectory_stream", "TreeEnergyOptimizer", "TreeLayoutFinder", - "TreeMPO", "TreeOptimizer", "tree_mpo", + "TreeMPO", "TreeOptimizer", "build_tree_operator", "tree_mpo", "TreePlan", "TreeStabOptimizer", "TreeTensorNetwork", diff --git a/tests/test_tree_mpo.py b/tests/test_tree_mpo.py index b84135f..271fae6 100644 --- a/tests/test_tree_mpo.py +++ b/tests/test_tree_mpo.py @@ -2,9 +2,10 @@ import numpy as np import pytest +import quimb.tensor as qtn import pepsy -from pepsy.optimizers.tree import TreeMPO, TreePlan, tree_mpo +from pepsy.optimizers.tree import TreeMPO, TreePlan, build_tree_operator, tree_mpo pytest.importorskip("symmray") @@ -36,6 +37,7 @@ def test_tree_mpo_preserves_native_fermionic_symmetry_and_expectation(symmetry): dtype="complex64", ) + assert pepsy.build_tree_operator is build_tree_operator assert pepsy.tree_mpo is tree_mpo assert mpo.pepsy_tree_order == (2, 0, 3, 1) assert mpo.pepsy_tree_native is True @@ -260,7 +262,7 @@ def test_fermion_tree_mpo_class_is_native_and_keeps_chain_representation(symmetr [(0, 1), (1, 2)], t=1.0, U=2.0, mu=0.1, ) plan = TreePlan.from_order(range(4), structure="balanced", top_arity=2) - operator = fermion.to_tree_mpo( + operator = fermion.build_tree_operator( hamiltonian=hamiltonian, tree=plan, compress=True, @@ -281,6 +283,8 @@ def test_fermion_tree_mpo_class_is_native_and_keeps_chain_representation(symmetr ) assert isinstance(operator, TreeMPO) + assert type(fermion).to_tree_mpo is type(fermion).build_tree_operator + assert type(fermion).build_tree_mpo is type(fermion).build_tree_operator assert operator.backend == "symmray" assert operator.chain_mpo is not None assert len(operator.tree_networks) == 1 @@ -331,3 +335,122 @@ def test_sparse_pair_terms_do_not_invent_missing_pairs(): operator.compress(max_bond=16) assert operator.max_bond() <= 16 assert operator.max_bond() <= raw_bond + + +def test_to_tree_mpo_applies_to_backend_to_operator_networks(): + """``to_tree_mpo(to_backend=...)`` moves the primary TreeMPO operator. + + The compatibility chain MPO was always backend-converted, but the primary + ``TreeMPO.tree_networks`` are what ``TreeMPO.expectation`` contracts. They + must land on the requested backend so the operator matches a tree state + built on that backend. + """ + fermion = pepsy.Fermion(spinful=True, symmetry="U1", dtype="complex128") + hamiltonian = fermion.hamiltonian([(0, 1), (1, 2)], t=1.0, U=2.0, mu=0.1) + plan = TreePlan.from_order(range(4), structure="balanced", top_arity=2) + + seen = [] + + def to_backend(array): + converted = np.asarray(array, dtype=np.complex64) + seen.append(np.dtype(converted.dtype)) + return converted + + operator = plan.to_tree_mpo( + hamiltonian, fermionic=True, compress=True, to_backend=to_backend, + ) + + assert isinstance(operator, TreeMPO) + assert seen, "to_backend was never applied to the operator networks" + for network in operator.tree_networks: + for tensor in network: + assert np.dtype(tensor.data.dtype) == np.dtype("complex64") + + +def test_tree_mpo_matches_quimb_generalized_operator_surface(): + """TreeMPO shares Quimb's operator API without pretending to be a chain.""" + plan = TreePlan.from_order(range(4), structure="balanced", top_arity=2) + operator = TreeMPO.from_terms( + plan, + { + (0,): np.diag([1.0, 2.0]), + (1, 2): np.arange(16.0).reshape(2, 2, 2, 2), + }, + compress=False, + ) + + assert isinstance(operator, qtn.TensorNetworkGenOperator) + assert operator.sites == (0, 1, 2, 3) + assert operator.nsites == operator.nqubits == 4 + assert operator.site_tag(2) == "I2" + assert operator.upper_ind(2) == "k2" + assert operator.lower_ind(2) == "b2" + assert tuple(operator.gen_sites_present()) == operator.sites + assert operator.to_dense().shape == (16, 16) + assert isinstance(operator.H, TreeMPO) + assert isinstance(operator.copy(), qtn.TensorNetworkGenOperator) + assert isinstance(qtn.TensorNetworkGenOperator(operator), qtn.TensorNetworkGenOperator) + + root = plan.root + child = plan.children[root][0] + assert operator.node_tag(root) in operator.tags + assert child in operator.neighbors(root) + assert operator.bond(root, child) in operator.inner_inds() + assert operator.validate() is operator + + canonical = operator.canonize_around(f"N{root}") + assert isinstance(canonical, TreeMPO) + assert canonical is not operator + assert operator.compress_between(root, child, max_bond=16) is operator + assert operator.canonize_around_(root) is operator + + +def test_tree_mpo_is_mpo_twin_over_tree_geometry(): + """The mature TreeMPO surface mirrors useful Quimb MPO operations.""" + plan = TreePlan.from_order(range(4), structure="balanced", top_arity=2) + dense = np.arange(256.0).reshape(16, 16) + operator = TreeMPO.from_dense(plan, dense) + + assert operator.L == operator.nsites == 4 + assert operator.site_ind(2) == operator.upper_ind(2) == "k2" + assert operator.validate() is operator + np.testing.assert_allclose(operator.to_dense(), dense, rtol=1e-11, atol=1e-11) + + identity = operator.identity() + np.testing.assert_allclose(identity.to_dense(), np.eye(16)) + np.testing.assert_allclose( + operator.add_MPO(identity).to_dense(), dense + np.eye(16), + ) + np.testing.assert_allclose( + operator.add_MPO(identity, negate=True).to_dense(), dense - np.eye(16), + ) + assert operator.amplitude([0, 0, 0, 0]) == pytest.approx(dense[0, 0]) + + selected = operator.select_sites((0, 1)) + assert isinstance(selected, TreeMPO) + assert len(selected.tree_networks[0].tensors) == 2 + + root = plan.root + child = plan.children[root][0] + values = operator.singular_values(root, child) + assert values.ndim == 1 + canonical = operator.canonize_between(root, child) + assert canonical.is_subtree_canonical_form((root, child)) + assert canonical is not operator + assert operator.copy(conj=True).to_dense().shape == (16, 16) + np.testing.assert_allclose( + operator.copy(transpose=True).to_dense(), dense.T, rtol=1e-11, atol=1e-11, + ) + + +def test_tree_mpo_from_fill_and_random_state_helpers(): + """TreeMPO exposes the corresponding construction and state helpers.""" + plan = TreePlan.from_order(range(4), structure="balanced", top_arity=2) + operator = TreeMPO.from_fill_fn( + lambda shape: np.ones(shape), plan, bond_dim=2, + ) + assert operator.validate() is operator + random_operator = TreeMPO.rand(plan, bond_dim=2, seed=7) + assert random_operator.validate() is random_operator + state = random_operator.rand_state(2, seed=7) + assert state.plan is plan From 9e708920d590179259514d455672d78efe0688a2 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Tue, 4 Aug 2026 09:44:48 -0600 Subject: [PATCH 64/70] Support mixed-charge native TreeMPO sums --- .github/skills/tree-optimizer/SKILL.md | 5 + docs/api/optimizers/tree.md | 6 +- src/pepsy/optimizers/tree/layout.py | 4 +- src/pepsy/optimizers/tree/operators.py | 168 +++++++++++++++++++++++-- src/pepsy/tensors/symmetric.py | 5 + tests/test_tree_mpo.py | 96 ++++++++++++++ 6 files changed, 274 insertions(+), 10 deletions(-) diff --git a/.github/skills/tree-optimizer/SKILL.md b/.github/skills/tree-optimizer/SKILL.md index 1ec09dc..93c2888 100644 --- a/.github/skills/tree-optimizer/SKILL.md +++ b/.github/skills/tree-optimizer/SKILL.md @@ -421,6 +421,11 @@ and `bond` provide the branched geometry. It cannot inherit the chain-only `MatrixProductOperator` implementation because a tree has no left/right ordering; `chain_mpo` remains the separate chain-compatible representation. +Mixed native operator charges are represented as one public `TreeMPO` with one +homogeneous Symmray tree network per charge in `tree_networks`. Use +`charge_sectors=True` only when separate `TreeMPO` objects are explicitly +needed. + For native fermionic Hamiltonians, one-, two-, and higher-site neutral terms are fused and factorized from their native Symmray operator tensor over the TreePlan Steiner subtree, then amalgamated into one charge-aware direct-sum diff --git a/docs/api/optimizers/tree.md b/docs/api/optimizers/tree.md index 0152ab3..09e7a2d 100644 --- a/docs/api/optimizers/tree.md +++ b/docs/api/optimizers/tree.md @@ -149,7 +149,11 @@ does not create a hyperedge for the normal Hamiltonian path. The resulting TTNO can be canonicalized and compressed with `tree_operator.canonicalize()` and `tree_operator.compress(cutoff=..., max_bond=...)`; no Jordan--Wigner -conversion is used. Structured observables can use a smaller compact TTNO. +conversion is used. Nonzero or mixed operator charges remain separate +homogeneous native networks inside the same public `TreeMPO`, so callers do +not need `charge_sectors=True` just to construct one operator object; +`charge_sectors=True` remains available when separate objects are preferred. +Structured observables can use a smaller compact TTNO. Pass `fermionic=False` only for dense ordinary/Jordan--Wigner-compatible terms. Existing `OneDMap` lattice maps remain unchanged and should continue to be used for regular 2D/3D coordinate layouts. diff --git a/src/pepsy/optimizers/tree/layout.py b/src/pepsy/optimizers/tree/layout.py index bb877a5..95e3575 100644 --- a/src/pepsy/optimizers/tree/layout.py +++ b/src/pepsy/optimizers/tree/layout.py @@ -1086,7 +1086,9 @@ def build_tree_operator(self, hamiltonian, **kwargs): The returned object keeps the optional chain MPO available as ``.chain_mpo`` and exposes the TreePlan-routed representation through ``.tree_networks`` and ``.expectation``. With - ``charge_sectors=True`` the result is ``{charge: TreeMPO}``. + mixed native charges, one public ``TreeMPO`` contains one homogeneous + network per charge. ``charge_sectors=True`` remains available when + separate sector objects are specifically desired. """ from .operators import build_tree_operator diff --git a/src/pepsy/optimizers/tree/operators.py b/src/pepsy/optimizers/tree/operators.py index dd93776..cec7543 100644 --- a/src/pepsy/optimizers/tree/operators.py +++ b/src/pepsy/optimizers/tree/operators.py @@ -755,7 +755,7 @@ def cyclic(self): def to_dense(self, *inds_seq, to_qarray=False, **contract_opts): """Contract the complete operator, summing internal term networks.""" - if len(self.tree_networks) == 1: + if len(self.tree_networks) == 1 and not self.fermionic: return qtn.TensorNetworkGenOperator.to_dense( self, *inds_seq, @@ -766,6 +766,24 @@ def to_dense(self, *inds_seq, to_qarray=False, **contract_opts): inds_seq = (self.upper_inds_present, self.lower_inds_present) values = [] for network in self.tree_networks: + if self.fermionic: + # Symmray's block-sparse contraction assumes a neutral scalar + # when it closes all internal legs. A charged operator has + # nonzero open physical charge, so densify each local block + # first and contract the ordinary tree network. This is only + # the explicit ``to_dense`` escape hatch; native expectation + # and compression remain graded and factorized. + dense_tensors = [] + for tensor in network: + data = tensor.data + if hasattr(data, "to_dense"): + data = data.to_dense() + dense_tensors.append(qtn.Tensor( + data, + inds=tensor.inds, + tags=tensor.tags, + )) + network = qtn.TensorNetwork(dense_tensors) view = qtn.TensorNetworkGenOperator( network, virtual=True, @@ -2078,12 +2096,6 @@ def _native_tree_term_network( if physical_map and isinstance(physical_map[0], tuple) else 0 ) - term_charge = getattr(term, "charge", zero) - if term_charge != zero: - raise ValueError( - "a single native TTNO must be neutral; use charge_sectors=True " - "for a charged operator sum." - ) endpoint_nodes = tuple(plan.node_of_qubit[site] for site in support) if len(support) == 1: @@ -3084,6 +3096,100 @@ def _dense_tree_tensor_network_for_term(plan, operator, support, *, dtype=None): return network +def _terms_by_operator_charge(terms): + """Group native terms without mixing their Symmray operator charges.""" + grouped = {} + for where, term in terms.items(): + grouped.setdefault(getattr(term, "charge", 0), {})[where] = term + return grouped + + +def _charge_is_zero(charge): + """Return whether an Abelian scalar or tuple charge is neutral.""" + if isinstance(charge, tuple): + return all(value == 0 for value in charge) + return charge == 0 + + +def _build_mixed_charge_tree_operator( + plan, + hamiltonian, + *, + max_bond=None, + cutoff=1e-12, + compress=True, + dtype=None, + fermionic=True, + to_backend=None, +): + """Build one public ``TreeMPO`` from separate native charge networks.""" + from ...tensors.symmetric import ( + SymHamiltonian, + _apply_to_tensor_network_arrays, + ) + + networks = [] + for charge, sector_terms in _terms_by_operator_charge( + hamiltonian.terms + ).items(): + if not _charge_is_zero(charge): + # A nonzero-charge TTNO cannot be amalgamated into a neutral + # direct-sum tensor: the charge belongs to one open operator + # boundary tensor. Keep each charged term as its own homogeneous + # network and let the public TreeMPO sum them. + for where, term in sector_terms.items(): + network = _native_tree_term_network( + plan, + term, + _term_support(where), + symmetry=hamiltonian.symmetry, + cutoff=cutoff, + dtype=dtype, + ) + networks.append(_normalize_native_term_edge_orientation( + network, + plan, + symmetry=hamiltonian.symmetry, + dtype=dtype, + )) + continue + sector_hamiltonian = SymHamiltonian.from_terms( + hamiltonian.model, + hamiltonian.symmetry, + sector_terms, + parameters=hamiltonian.parameters, + ) + sector_operator = _build_tree_operator( + plan, + sector_hamiltonian, + cutoff=cutoff, + max_bond=max_bond, + compress=False, + dtype=dtype, + fermionic=fermionic, + ) + if isinstance(sector_operator, TreeMPO): + networks.extend(sector_operator.tree_networks) + else: + networks.append(sector_operator) + + operator = TreeMPO( + plan, + tuple(networks), + terms=hamiltonian.terms, + backend="symmray" if fermionic else "dense", + fermionic=fermionic, + symmetry=hamiltonian.symmetry, + compressed=False, + ) + if compress: + operator.compress(max_bond=max_bond, cutoff=cutoff) + if to_backend is not None: + for network in operator.tree_networks: + _apply_to_tensor_network_arrays(network, to_backend) + return operator + + def _build_tree_operator( plan, hamiltonian, @@ -3098,6 +3204,30 @@ def _build_tree_operator( symmetry = hamiltonian.symmetry terms = hamiltonian.terms + if any( + not _charge_is_zero(charge) + for charge in _terms_by_operator_charge(terms) + ): + # A charged TTNO carries its operator charge on an open boundary + # tensor. Keep charged terms as separate homogeneous networks rather + # than forcing them into the neutral direct-sum construction below. + return tuple( + _normalize_native_term_edge_orientation( + _native_tree_term_network( + plan, + term, + _term_support(where), + symmetry=symmetry, + cutoff=cutoff, + dtype=dtype, + ), + plan, + symmetry=symmetry, + dtype=dtype, + ) + for where, term in terms.items() + ) + # The full staggered eta correlator is a symmetric rank-one pair table. # Compile it before falling back to one actual tree contraction per term; # this keeps p_eta_stag2 at a four-state tree bond for arbitrary N. @@ -3376,8 +3506,30 @@ def build_tree_operator( The compatibility :func:`tree_mpo` builder returns the ordinary linear chain MPO and attaches this tree operator to it. This function returns only the tree-native operator, keeping the two tensor-network geometries - explicit. With ``charge_sectors=True`` it returns ``{charge: TreeMPO}``. + explicit. Mixed native charges are combined into one ``TreeMPO`` with one + homogeneous Symmray network per charge. With ``charge_sectors=True`` it + instead returns ``{charge: TreeMPO}`` for callers that need separate + sector objects. """ + if fermionic and not charge_sectors: + from ...tensors.symmetric import SymHamiltonian + + if not isinstance(hamiltonian, SymHamiltonian): + raise TypeError("hamiltonian must be a SymHamiltonian instance.") + if any( + not _charge_is_zero(charge) + for charge in _terms_by_operator_charge(hamiltonian.terms) + ): + return _build_mixed_charge_tree_operator( + plan, + hamiltonian, + max_bond=max_bond, + cutoff=cutoff, + compress=compress, + dtype=dtype, + fermionic=fermionic, + to_backend=to_backend, + ) built = tree_mpo( plan, hamiltonian, diff --git a/src/pepsy/tensors/symmetric.py b/src/pepsy/tensors/symmetric.py index 9d7e187..93e49ef 100644 --- a/src/pepsy/tensors/symmetric.py +++ b/src/pepsy/tensors/symmetric.py @@ -10372,6 +10372,11 @@ def build_tree_operator( Native ``fermionic=True`` keeps Symmray's graded tensors intact for U1, U1U1, and other supported symmetries. + Mixed operator charges are exposed as one public ``TreeMPO`` whose + internal ``tree_networks`` keep one homogeneous native network per + charge. Pass ``charge_sectors=True`` only when separate sector + objects are specifically desired. + This is the canonical ``Fermion`` tree-operator entry point. ``to_tree_mpo`` and ``build_tree_mpo`` remain compatibility aliases. """ diff --git a/tests/test_tree_mpo.py b/tests/test_tree_mpo.py index 271fae6..52b19aa 100644 --- a/tests/test_tree_mpo.py +++ b/tests/test_tree_mpo.py @@ -252,6 +252,102 @@ def test_native_tree_mpo_amalgamates_higher_rank_term(): np.testing.assert_allclose(operator.expectation(state), direct) +@pytest.mark.parametrize("symmetry", ["U1", "U1U1"]) +def test_tree_operator_combines_mixed_native_charges(symmetry): + """Mixed native charges form one compressible public TreeMPO sum.""" + fermion = pepsy.Fermion( + spinful=True, symmetry=symmetry, dtype="complex128", + ) + neutral = fermion.hopping_operator() + charged = fermion.operator_term( + [(1.0, (((1), "double"), ((2), "annihilate_up")))], + sites=(1, 2), + label="mixed_tree_charge", + ) + reference = fermion.hamiltonian([(0, 1)], t=0.0, U=0.0, mu=0.0) + hamiltonian = type(reference).from_terms( + reference.model, + reference.symmetry, + {(0, 1): neutral, (1, 2): charged}, + parameters=reference.parameters, + ) + plan = TreePlan.from_order(range(4), structure="balanced", top_arity=2) + + operator = fermion.build_tree_operator( + hamiltonian=hamiltonian, + tree=plan, + compress=False, + ) + + assert isinstance(operator, TreeMPO) + assert operator.chain_mpo is None + assert len(operator.tree_networks) == 2 + assert all( + type(tensor.data).__name__.endswith("FermionicArray") + for network in operator.tree_networks + for tensor in network + ) + + explicit_sectors = fermion.build_tree_operator( + hamiltonian=hamiltonian, + tree=plan, + compress=False, + charge_sectors=True, + ) + assert set(explicit_sectors) == {fermion.zero_charge, charged.charge} + assert all(isinstance(sector, TreeMPO) for sector in explicit_sectors.values()) + + neutral_hamiltonian = type(reference).from_terms( + reference.model, + reference.symmetry, + {(0, 1): neutral}, + parameters=reference.parameters, + ) + charged_hamiltonian = type(reference).from_terms( + reference.model, + reference.symmetry, + {(1, 2): charged}, + parameters=reference.parameters, + ) + neutral_operator = fermion.build_tree_operator( + hamiltonian=neutral_hamiltonian, tree=plan, compress=False, + ) + charged_operator = fermion.build_tree_operator( + hamiltonian=charged_hamiltonian, tree=plan, compress=False, + ) + np.testing.assert_allclose( + operator.to_dense(), + neutral_operator.to_dense() + charged_operator.to_dense(), + ) + + leaf_charges = ( + { + site: fermion.local_fock_state((0, 0), site=site)[0] + for site in range(4) + } + if symmetry == "U1" else + {site: (0, 0) for site in range(4)} + ) + state = pepsy.TreeTensorNetwork.from_symmray_plan( + plan, + symmetry=symmetry, + physical_sectors=fermion.physical_sectors, + leaf_charges=leaf_charges, + bond_dim=2, + fermionic=True, + seed=7, + dtype="complex128", + ) + np.testing.assert_allclose( + operator.expectation(state), + state.expectation_mpo_exact(operator, range(4)), + ) + + operator.canonicalize().compress(max_bond=16, cutoff=1e-12) + assert operator.compressed is True + assert isinstance(operator.pepsy_compression_report, list) + + @pytest.mark.parametrize("symmetry", ["U1", "U1U1"]) def test_fermion_tree_mpo_class_is_native_and_keeps_chain_representation(symmetry): """Fermion exposes the class API for both native Symmray symmetries.""" From 1bfc06f3bb992c472574ce07c48ac4d1b9f63b61 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Tue, 4 Aug 2026 12:10:30 -0600 Subject: [PATCH 65/70] Harden native GPU tree QR fallbacks --- src/pepsy/optimizers/tree/operators.py | 41 +++-- src/pepsy/optimizers/tree/ttn.py | 237 ++++++++++++++++++++++--- tests/test_optimize_tree.py | 158 +++++++++++++++++ 3 files changed, 396 insertions(+), 40 deletions(-) diff --git a/src/pepsy/optimizers/tree/operators.py b/src/pepsy/optimizers/tree/operators.py index cec7543..126fb68 100644 --- a/src/pepsy/optimizers/tree/operators.py +++ b/src/pepsy/optimizers/tree/operators.py @@ -37,6 +37,17 @@ __all__ = ["TreeMPO", "build_tree_operator", "tree_mpo"] +def _as_numpy(data, *, dtype=None): + """Convert a dense backend array to host NumPy construction data.""" + if hasattr(data, "detach"): + data = data.detach() + if hasattr(data, "cpu"): + data = data.cpu() + if hasattr(data, "get"): + data = data.get() + return np.asarray(data, dtype=dtype) + + def _tree_plan_signature(plan): """Return a stable structural signature for a tree-MPO annotation.""" return ( @@ -1500,8 +1511,8 @@ def _operator_native_channels( raise ValueError("could not split a native two-site operator.") left = left.unfuse(0).transpose((2, 0, 1)) right = right.unfuse(1) - left_data = np.asarray(left.to_dense(), dtype=dtype) - right_data = np.asarray(right.to_dense(), dtype=dtype) + left_data = _as_numpy(left.to_dense(), dtype=dtype) + right_data = _as_numpy(right.to_dense(), dtype=dtype) physical_map = _expanded_index_charges(left.indices[1]) if _expanded_index_charges(left.indices[2]) != physical_map: raise ValueError("native operator factors have mismatched physical maps.") @@ -1611,7 +1622,7 @@ def _combined_tree_operator( if len(support) == 1: data = ( _dense_operator_array(term, dtype=dtype) - if not fermionic else np.asarray(term.to_dense(), dtype=dtype) + if not fermionic else _as_numpy(term.to_dense(), dtype=dtype) ) if data.ndim != 2 or data.shape[0] != data.shape[1]: raise ValueError("one-site operators must be square matrices.") @@ -2192,7 +2203,7 @@ def rebuild_with_axis(data, maps, duals, dense): # A physical TreePlan root can lie on the active Steiner # subtree without being an endpoint. Its operator action is # the identity, so add that even physical pair explicitly. - dense = np.asarray(data.to_dense(), dtype=dtype or complex) + dense = _as_numpy(data.to_dense(), dtype=dtype or complex) dense = np.einsum( "ab,...->ab...", np.eye(len(physical_map), dtype=dense.dtype), @@ -2220,7 +2231,7 @@ def rebuild_with_axis(data, maps, duals, dense): if index in existing: continue dense = np.expand_dims( - np.asarray(data.to_dense(), dtype=dtype or complex), + _as_numpy(data.to_dense(), dtype=dtype or complex), axis=-1, ) maps = [ @@ -2324,7 +2335,7 @@ def _normalize_native_term_edge_orientation(network, plan, *, symmetry, dtype=No permutation.append(old_positions[old_charge][position]) used[old_charge] = position + 1 dense = np.take( - np.asarray(tensor.data.to_dense(), dtype=dtype or complex), + _as_numpy(tensor.data.to_dense(), dtype=dtype or complex), permutation, axis=axis, ) @@ -2470,7 +2481,7 @@ def edge_name(node, neighbor): data = np.zeros(shape, dtype=dtype or complex) for term_index, network in enumerate(term_networks): tensor = network[f"N{node}"].transpose(*desired) - local = np.asarray(tensor.data.to_dense(), dtype=data.dtype) + local = _as_numpy(tensor.data.to_dense(), dtype=data.dtype) slices = [] if qubit is not None: slices.extend((slice(None), slice(None))) @@ -2718,7 +2729,7 @@ def _native_from_dense( def _pair_coefficient_factors(terms, nsite): """Factor an off-diagonal symmetric coefficient table, if possible.""" first_support, first_term = next(iter(terms.items())) - first_matrix = np.asarray(first_term.to_dense()).reshape( + first_matrix = _as_numpy(first_term.to_dense()).reshape( (first_term.shape[0] * first_term.shape[2],) * 2 ) table = np.zeros((nsite, nsite), dtype=complex) @@ -2726,7 +2737,7 @@ def _pair_coefficient_factors(terms, nsite): support = _term_support(where) if len(support) != 2 or support[0] >= support[1]: return None - matrix = np.asarray(term.to_dense()).reshape( + matrix = _as_numpy(term.to_dense()).reshape( (term.shape[0] * term.shape[2],) * 2 ) denominator = np.vdot(first_matrix, first_matrix) @@ -2790,8 +2801,8 @@ def _pair_endpoint_automaton( return None left = left.unfuse(0).transpose((2, 0, 1)) right = right.unfuse(1) - left_data = np.asarray(left.to_dense(), dtype=dtype or complex)[0] - right_data = np.asarray(right.to_dense(), dtype=dtype or complex)[0] + left_data = _as_numpy(left.to_dense(), dtype=dtype or complex)[0] + right_data = _as_numpy(right.to_dense(), dtype=dtype or complex)[0] physical_map = _expanded_index_charges(left.indices[1]) physical_dim = len(physical_map) zero = 0 if symmetry in {"U1", "Z2"} else (0, 0) @@ -2897,8 +2908,8 @@ def _pair_chain_mpo( return None left = left.unfuse(0).transpose((2, 0, 1)) right = right.unfuse(1) - left_data = np.asarray(left.to_dense(), dtype=dtype or complex)[0] - right_data = np.asarray(right.to_dense(), dtype=dtype or complex)[0] + left_data = _as_numpy(left.to_dense(), dtype=dtype or complex)[0] + right_data = _as_numpy(right.to_dense(), dtype=dtype or complex)[0] physical_map = _expanded_index_charges(left.indices[1]) physical_dim = len(physical_map) zero = 0 if symmetry in {"U1", "Z2"} else (0, 0) @@ -3003,7 +3014,7 @@ def _tree_tensor_network_for_term( zero = getattr(term, "zero_charge", None) if zero is None: zero = 0 if symmetry in {"U1", "Z2"} else (0, 0) - operator_dtype = np.dtype(dtype or np.asarray(term.to_dense()).dtype) + operator_dtype = np.dtype(dtype or _as_numpy(term.to_dense()).dtype) # The term's native indices are ordered as all upper physical legs, # followed by all lower physical legs. Keep that ordering intact while @@ -3048,7 +3059,7 @@ def _dense_operator_array(operator, *, dtype=None): operator = operator.to_dense() elif hasattr(operator, "data"): operator = operator.data - return np.asarray(operator, dtype=dtype) + return _as_numpy(operator, dtype=dtype) def _dense_tree_tensor_network_for_term(plan, operator, support, *, dtype=None): diff --git a/src/pepsy/optimizers/tree/ttn.py b/src/pepsy/optimizers/tree/ttn.py index e2811fa..067c1bd 100644 --- a/src/pepsy/optimizers/tree/ttn.py +++ b/src/pepsy/optimizers/tree/ttn.py @@ -54,22 +54,141 @@ __all__ = ["TreeTensorNetwork"] +def _native_rank_safe_qr(array, backend): + """Factor a finite complex64 block with safe backend QR.""" + if backend == "torch": + import torch as xp # pylint: disable=import-outside-toplevel + elif backend == "cupy": + import cupy as xp # pylint: disable=import-outside-toplevel + else: # pragma: no cover - only native GPU backends call this helper + raise ValueError(f"unsupported native QR backend {backend!r}") + + rows, cols = array.shape + rank_limit = min(rows, cols) + r_factor = array.clone() + q_factor = xp.eye(rows, dtype=array.dtype) + + for index in range(rank_limit): + vector = r_factor[index:, index] + magnitude = vector.abs() + scale = magnitude.amax() + # Construct the reflector from a max-normalized vector. This avoids + # the subnormal norm division that makes Torch's complex64 QR fail. + valid_scale = xp.isfinite(scale) & (scale > 0) + safe_scale = xp.where(valid_scale, scale, xp.ones_like(scale)) + normalized = vector / safe_scale + norm = xp.sqrt(xp.sum(normalized.abs() ** 2)) + valid_norm = valid_scale & xp.isfinite(norm) & (norm > 0) + safe_norm = xp.where(valid_norm, norm, xp.ones_like(norm)) + first = normalized[0] + first_abs = first.abs() + safe_first_abs = xp.where( + first_abs > 0, first_abs, xp.ones_like(first_abs), + ) + phase = xp.where( + first_abs > 0, + first / safe_first_abs, + xp.ones_like(first), + ) + alpha = -phase * safe_norm + reflector = normalized.clone() + reflector[0] = reflector[0] - alpha + reflector_norm = xp.sum(reflector.conj() * reflector).real + valid_reflector = ( + valid_norm + & xp.isfinite(reflector_norm) + & (reflector_norm > 0) + ) + safe_reflector_norm = xp.where( + valid_reflector, reflector_norm, xp.ones_like(reflector_norm), + ) + beta = xp.where( + valid_reflector, + 2 / safe_reflector_norm, + xp.zeros_like(reflector_norm), + ) + + trailing = r_factor[index:, index:] + r_factor[index:, index:] = trailing - reflector[:, None] * ( + beta * (reflector.conj() @ trailing) + )[None, :] + + q_trailing = q_factor[:, index:] + q_factor[:, index:] = q_trailing - (q_trailing @ reflector)[:, None] * ( + beta * reflector.conj() + )[None, :] + + return q_factor[:, :rank_limit], r_factor[:rank_limit, :] + + +def _native_factors_finite(factors, backend): + """Check native QR factors without converting GPU arrays to NumPy.""" + if factors is None: + return False + for factor in factors: + if factor is None: + continue + if backend == "torch": + finite = factor.isfinite().all().item() + elif backend == "cupy": + import cupy as cp # pylint: disable=import-outside-toplevel + + finite = cp.isfinite(factor).all().item() + else: + finite = np.isfinite(factor).all() + if not bool(finite): + return False + return True + + +def _native_cast_complex128(array, backend): + """Promote one native GPU block without moving it off device.""" + if backend == "torch": + import torch # pylint: disable=import-outside-toplevel + + return array.to(dtype=torch.complex128) + if backend == "cupy": + import cupy as cp # pylint: disable=import-outside-toplevel + + return array.astype(cp.complex128, copy=False) + raise ValueError(f"unsupported native GPU backend {backend!r}") + + +def _native_cast_like(array, reference, backend): + """Cast one native factor back to the original device and dtype.""" + if backend == "torch": + return array.to(dtype=reference.dtype, device=reference.device) + if backend == "cupy": + return array.astype(reference.dtype, copy=False) + raise ValueError(f"unsupported native GPU backend {backend!r}") + + +def _native_cast_factors(factors, reference, backend): + """Cast non-empty QR factors back to the original native dtype.""" + return tuple( + None if factor is None else _native_cast_like(factor, reference, backend) + for factor in factors + ) + + def _native_qr_block_scaled(array, **kwargs): - """QR one native charge block after a reversible power-of-two scaling. - - Torch's complex64 QR can return NaNs for a rank-deficient block whose - entries are small (around ``1e-9``) even though the block is finite. Native - Symmray QR is blockwise, so scaling each block independently is exact: the - isometric factor is unchanged and the triangular factor is divided by the - same scalar afterwards. Power-of-two scaling avoids introducing an extra - rounding step into the block values. + """QR one native charge block with a reversible scaling fallback. + + Torch's complex64 QR can return NaNs for a rank-deficient block even when + its largest entry is moderate, if other entries in the same block are + many orders of magnitude smaller. Healthy blocks keep Torch's native QR + path unchanged. Only a failed finite block is retried after a reversible + power-of-two scaling, with a native rank-safe fallback for structural + rank deficiency. """ opts = dict(kwargs) opts.pop("method", None) opts.pop("fn", None) - def torch_qr(x, qr_opts): - """Run the common native Torch QR block without composed dispatch.""" + def native_qr( + x, qr_opts, *, rank_safe=False, allow_failure=False, + ): + """Run one native backend QR block without composed dispatch.""" absorb = qr_opts.get("absorb", "right") left_like = absorb in { -1, "left", "Us,VH", "lfactor", "Us", @@ -80,7 +199,20 @@ def torch_qr(x, qr_opts): } if left_like: x = ar.do("transpose", x, (1, 0)) - q, r = ar.do("linalg.qr", x, **qr_kwargs) + try: + q, r = ar.do("linalg.qr", x, **qr_kwargs) + except Exception: + if not allow_failure: + raise + q, r = None, None + if rank_safe and not _native_factors_finite((q, r), backend): + # Native GPU QR can still emit NaNs for finite, structurally + # rank-deficient blocks after scaling. Use a backend-native + # fallback that avoids division by subnormal norms and completes + # the orthonormal basis explicitly for zero sectors. + q, r = _native_rank_safe_qr(x, backend) + if q is None or r is None: + return None if left_like: left = ar.do("transpose", r, (1, 0)) right = ar.do("transpose", q, (1, 0)) @@ -103,33 +235,88 @@ def torch_qr(x, qr_opts): # dispatch layer from every block while retaining exactly the same # reduced QR and the same ``stabilized=False`` policy. use_torch_qr = backend == "torch" + use_native_gpu_qr = backend == "cupy" or ( + backend == "torch" + and getattr(getattr(array, "device", None), "type", None) == "cuda" + ) if ar.get_dtype_name(array) != "complex64": - if use_torch_qr: - return torch_qr(array, opts) + if use_torch_qr or use_native_gpu_qr: + return native_qr(array, opts) return _quimb_qr_stabilized(array, **opts) - if backend == "torch": + if use_native_gpu_qr: + # Keep the normal requested-dtype GPU QR path unchanged. Only a + # genuinely nonfinite result pays for the same-device double retry. + direct = native_qr(array, opts, allow_failure=True) + if _native_factors_finite(direct, backend): + return direct + + high = _native_cast_complex128(array, backend) + high_result = native_qr(high, opts, allow_failure=True) + if _native_factors_finite(high_result, backend): + return _native_cast_factors(high_result, array, backend) + + if backend == "torch": + block_max = float(array.detach().abs().amax().item()) + else: + block_max = to_float(ar.do("max", ar.do("abs", array))) + elif use_torch_qr: + # Preserve Torch's normal QR exactly for healthy blocks. Only the + # exceptional nonfinite result pays for a scaled retry and, if needed, + # the rank-safe complex64 fallback. + direct = native_qr(array, opts) + if _native_factors_finite(direct, backend): + return direct + block_max = float(array.detach().abs().amax().item()) else: block_max = to_float(ar.do("max", ar.do("abs", array))) if not np.isfinite(block_max) or block_max == 0.0: # Preserve the original failure behaviour for non-finite input, while # allowing genuinely empty structural sectors through unchanged. - if use_torch_qr: - return torch_qr(array, opts) - return _quimb_qr_stabilized(array, **opts) - - # Values above this scale are not affected by the low-norm complex64 QR - # failure and avoid an unnecessary multiply/divide pair. - if block_max >= 2.0**-8: - if use_torch_qr: - return torch_qr(array, opts) + if use_torch_qr or use_native_gpu_qr: + return direct return _quimb_qr_stabilized(array, **opts) + # Normalize the exceptional finite block before retrying QR. A threshold + # based only on ``block_max`` is insufficient: the failing native block + # has max magnitude ~9e-3 but also contains entries ~1e-32 and structural + # zeros. Use a power of two so the scaling is exactly reversible in the + # complex64 representation. _, exponent = np.frexp(block_max) scale = float(np.ldexp(1.0, -int(exponent))) - if use_torch_qr: - left, singular_values, right = torch_qr(array * scale, opts) + if use_native_gpu_qr: + high_scaled = native_qr(high * scale, opts, allow_failure=True) + if _native_factors_finite(high_scaled, backend): + left, singular_values, right = _native_cast_factors( + high_scaled, array, backend, + ) + else: + fallback = native_qr( + array * scale, + opts, + rank_safe=True, + allow_failure=True, + ) + if not _native_factors_finite(fallback, backend): + raise RuntimeError( + "native GPU QR failed in requested and complex128 " + "dtypes, including the rank-safe fallback." + ) + left, singular_values, right = fallback + elif use_torch_qr: + fallback = native_qr( + array * scale, + opts, + rank_safe=True, + allow_failure=True, + ) + if not _native_factors_finite(fallback, backend): + raise RuntimeError( + "native Torch QR failed in requested dtype and its " + "rank-safe fallback." + ) + left, singular_values, right = fallback else: left, singular_values, right = _quimb_qr_stabilized( array * scale, **opts, diff --git a/tests/test_optimize_tree.py b/tests/test_optimize_tree.py index a66b95d..23f0a44 100644 --- a/tests/test_optimize_tree.py +++ b/tests/test_optimize_tree.py @@ -5540,6 +5540,164 @@ def test_native_complex64_qr_scales_low_norm_rank_deficient_block(): ) +def test_native_complex64_qr_leaves_healthy_block_on_native_path(monkeypatch): + """Healthy complex64 blocks use one unmodified native Torch QR call.""" + torch = pytest.importorskip("torch") + import pepsy.optimizers.tree.ttn as ttn_module + + generator = torch.Generator().manual_seed(4108) + block = torch.randn((6, 4), generator=generator).to(torch.complex64) + qr_calls = [] + original_do = ttn_module.ar.do + + def record_qr_call(fn, x, *args, **kwargs): + if fn == "linalg.qr": + qr_calls.append(x.detach().clone()) + return original_do(fn, x, *args, **kwargs) + + monkeypatch.setattr(ttn_module.ar, "do", record_qr_call) + q, _, r = _native_qr_block_scaled( + block, + method="qr", + absorb="right", + stabilized=False, + ) + expected_q, expected_r = torch.linalg.qr(block) + + assert len(qr_calls) == 1 + torch.testing.assert_close(qr_calls[0], block) + torch.testing.assert_close(q, expected_q) + torch.testing.assert_close(r, expected_r) + + +def test_native_complex64_qr_scales_dynamic_range_block(monkeypatch): + """Native QR scales moderate-norm blocks with tiny charge entries.""" + torch = pytest.importorskip("torch") + import pepsy.optimizers.tree.ttn as ttn_module + + block = torch.zeros((4, 10), dtype=torch.complex64) + for index, magnitude in enumerate((8.9e-3, 8.9e-11, 8.9e-25, 8.9e-41)): + block[index, index] = complex(magnitude, magnitude) + + qr_input_maxes = [] + original_do = ttn_module.ar.do + + def record_qr_input(fn, x, *args, **kwargs): + if fn == "linalg.qr": + qr_input_maxes.append(float(x.abs().amax().item())) + return original_do(fn, x, *args, **kwargs) + + monkeypatch.setattr(ttn_module.ar, "do", record_qr_input) + q, _, r = _native_qr_block_scaled( + block, + method="qr", + absorb="right", + stabilized=False, + ) + + assert len(qr_input_maxes) == 2 + assert qr_input_maxes[0] == pytest.approx(8.9e-3 * 2**0.5, rel=1e-6) + assert 0.5 <= qr_input_maxes[1] < 1.0 + assert torch.isfinite(q).all() + assert torch.isfinite(r).all() + torch.testing.assert_close( + q @ r, + block, + rtol=2e-4, + atol=1e-12, + ) + + +def test_native_complex64_qr_handles_structurally_rank_deficient_block(): + """Native QR keeps structural rank-deficient blocks in complex64.""" + torch = pytest.importorskip("torch") + block = torch.zeros((4, 10), dtype=torch.complex64) + for index, magnitude in enumerate((2e-2, 2e-10, 2e-24, 2e-40)): + block[index, index] = complex(magnitude, magnitude) + + q, _, r = _native_qr_block_scaled( + block, + method="qr", + absorb="right", + stabilized=False, + ) + + assert torch.isfinite(q).all() + assert torch.isfinite(r).all() + assert q.dtype == torch.complex64 + assert r.dtype == torch.complex64 + torch.testing.assert_close( + q @ r, + block, + rtol=2e-4, + atol=1e-12, + ) + + +@pytest.mark.parametrize("backend_name", ["torch", "cupy"]) +def test_native_complex64_gpu_qr_retries_same_device_double_precision( + monkeypatch, backend_name, +): + """Failed GPU QR uses an optimized same-device complex128 retry.""" + if backend_name == "torch": + torch = pytest.importorskip("torch") + if not torch.cuda.is_available(): + pytest.skip("CUDA is unavailable") + backend_module = torch + block = torch.tensor( + [[1.0 + 0.0j, 2.0 + 0.0j], [3.0 + 0.0j, 4.0 + 0.0j]], + dtype=torch.complex64, + device="cuda", + ) + dtype32 = torch.complex64 + else: + cupy = pytest.importorskip("cupy") + try: + if cupy.cuda.runtime.getDeviceCount() < 1: + pytest.skip("CUDA is unavailable") + except cupy.cuda.runtime.CUDARuntimeError as exc: + pytest.skip(f"CUDA is unavailable: {exc}") + backend_module = cupy + block = cupy.asarray( + [[1.0 + 0.0j, 2.0 + 0.0j], [3.0 + 0.0j, 4.0 + 0.0j]], + dtype=cupy.complex64, + ) + dtype32 = cupy.complex64 + + import pepsy.optimizers.tree.ttn as ttn_module + + original_do = ttn_module.ar.do + qr_dtypes = [] + + def fail_complex64_qr(fn, value, *args, **kwargs): + result = original_do(fn, value, *args, **kwargs) + if fn == "linalg.qr": + qr_dtypes.append(value.dtype) + if value.dtype == dtype32: + result = tuple( + backend_module.full_like( + factor, complex(float("nan"), float("nan")), + ) + for factor in result + ) + return result + + monkeypatch.setattr(ttn_module.ar, "do", fail_complex64_qr) + q, _, r = _native_qr_block_scaled( + block, + method="qr", + absorb="right", + stabilized=False, + ) + + assert qr_dtypes == [dtype32, backend_module.complex128] + assert q.dtype == dtype32 + assert r.dtype == dtype32 + assert bool(backend_module.isfinite(q).all()) + assert bool(backend_module.isfinite(r).all()) + assert float(backend_module.max(backend_module.abs(q @ r - block))) < 1e-5 + + def test_tree_stable_labels_route_submpo_by_payload_sites(monkeypatch): """Stable logical labels do not disable native structured MPO routing.""" x = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex) From f14a82690b2a7dac3e7b736773c9bd4f3e989823 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Tue, 4 Aug 2026 12:12:31 -0600 Subject: [PATCH 66/70] Stabilize native fermionic CTMRG and Torch VMC --- src/pepsy/boundary/__init__.py | 2 + src/pepsy/boundary/metrics.py | 378 +++++++++++++++++++++++++- src/pepsy/vmc/torch/amplitude.py | 50 +++- tests/test_prepare_boundary_inputs.py | 119 ++++++++ tests/test_symmetric_tensors.py | 76 ++++++ tests/test_vmc_api.py | 48 ++++ 6 files changed, 659 insertions(+), 14 deletions(-) diff --git a/src/pepsy/boundary/__init__.py b/src/pepsy/boundary/__init__.py index 81a5075..9aad363 100644 --- a/src/pepsy/boundary/__init__.py +++ b/src/pepsy/boundary/__init__.py @@ -14,6 +14,7 @@ peps_infidelity, peps_norm, peps_normalize, + quimb_ctmrg_projector_compat, ) from .states import BdyMPS, make_numpy_array_caster from .sweeps import CompBdy @@ -33,6 +34,7 @@ "peps_infidelity", "peps_norm", "peps_normalize", + "quimb_ctmrg_projector_compat", "metrics", "states", "sweeps", diff --git a/src/pepsy/boundary/metrics.py b/src/pepsy/boundary/metrics.py index 051f9d0..2b6d982 100644 --- a/src/pepsy/boundary/metrics.py +++ b/src/pepsy/boundary/metrics.py @@ -2,11 +2,15 @@ from __future__ import annotations +from contextlib import contextmanager +from functools import wraps import inspect import math import warnings from dataclasses import dataclass +import autoray as ar + from ..tensors.validation import _PHYS_OUTER, validate_tensor_network_tags from .states import BdyMPS from .sweeps import CompBdy @@ -16,6 +20,7 @@ "BoundaryContractResult", "contract_boundary", "contract_flat", + "quimb_ctmrg_projector_compat", "peps_normalize", "normalize", "boundary_norm", @@ -37,6 +42,248 @@ ) +@contextmanager +def quimb_ctmrg_projector_compat(): + """Use current-network projectors for Quimb's cyclic CTMRG path. + + Some Quimb versions compute a row of CTMRG projectors from one snapshot + while inserting those projectors into a network that is modified after + each insertion. On cyclic networks with non-uniform effective bond + dimensions this can leave the projector shapes out of sync with the + current tensor network. Projectors are therefore computed from a fresh + copy of the live insertion network, with a fresh fermionic dummy-mode + namespace, and inserted back into that live network. Native zero-charge + projector sectors also receive finite identity blocks. + + The patch is scoped to this context and does not modify the installed + Quimb package or affect boundary-MPS contractions. + """ + try: + from quimb.tensor import tensor_core as qtc + except ImportError: # pragma: no cover - Quimb is an optional dependency + yield + return + + tensor_network = getattr(qtc, "TensorNetwork", None) + original = getattr( + tensor_network, + "insert_compressor_between_regions", + None, + ) + original_inplace = getattr( + tensor_network, + "insert_compressor_between_regions_", + None, + ) + original_oblique = getattr(qtc, "compute_oblique_projectors", None) + original_reduced_factor = getattr( + qtc, + "squared_op_to_reduced_factor", + None, + ) + if tensor_network is None or original is None: + yield + return + + is_fermionic = getattr( + qtc, + "isfermionic", + lambda value: bool(getattr(value, "fermionic", False)), + ) + + @wraps(original) + def insert_current_projector(self, ltags, rtags, *args, insert_into=None, **kwargs): + if insert_into is None: + return original(self, ltags, rtags, *args, **kwargs) + + # Quimb's CTMRG caller passes a calculation copy as ``self`` and the + # live, progressively modified network as ``insert_into``. Recompute + # from a fresh copy of the live network so bond ids and sizes match, + # while separating the squared-environment dummy-mode namespace. + current_kwargs = dict(kwargs) + current_kwargs.pop("inplace", None) + projector_source = _freshen_fermionic_dummy_modes(insert_into) + return original( + projector_source, + ltags, + rtags, + *args, + insert_into=insert_into, + inplace=True, + **current_kwargs, + ) + + @wraps(original_inplace or original) + def insert_current_projector_inplace( + self, + ltags, + rtags, + *args, + insert_into=None, + **kwargs, + ): + # Quimb's arbitrary-geometry compression path calls the underscore + # alias, whose partialmethod otherwise retains the unpatched original + # implementation and recomputes projectors on a stale snapshot. + kwargs.setdefault("inplace", True) + return insert_current_projector( + self, + ltags, + rtags, + *args, + insert_into=insert_into, + **kwargs, + ) + + tensor_network.insert_compressor_between_regions = insert_current_projector + if original_inplace is not None: + tensor_network.insert_compressor_between_regions_ = ( + insert_current_projector_inplace + ) + + if callable(original_oblique): + oblique_globals = original_oblique.__globals__ + + @wraps(original_oblique) + def compute_fermionic_oblique_projectors( + Rl, + Rr, + max_bond=None, + cutoff=0.0, + absorb="both", + cutoff_mode="rsum2", + method="svd", + **compress_opts, + ): + """Avoid zero-sector inverses in native fermionic projectors.""" + if not (is_fermionic(Rl) or is_fermionic(Rr)) or absorb != "both": + return original_oblique( + Rl, + Rr, + max_bond=max_bond, + cutoff=cutoff, + absorb=absorb, + cutoff_mode=cutoff_mode, + method=method, + **compress_opts, + ) + + Ut, st, VHt = oblique_globals["array_split"]( + Rl @ Rr, + max_bond=-1 if max_bond is None else max_bond, + cutoff=cutoff, + absorb=None, + cutoff_mode=cutoff_mode, + method=method, + **compress_opts, + ) + + # Quimb's default projector uses 1 / sqrt(s). For a sparse + # symmetry sector with s == 0, the mathematically appropriate + # pseudoinverse is zero rather than inf or NaN. + inverse_sqrt = st.copy() + zero_tol = 1.0e-12 + + def _inverse_sqrt(block): + out = ar.do("zeros_like", block) + mask = ar.do("abs", block) > zero_tol + out[mask] = 1.0 / ar.do("sqrt", block[mask]) + return out + + inverse_sqrt.apply_to_arrays(_inverse_sqrt) + Pl = Rr @ oblique_globals["rdmul"]( + oblique_globals["dag"](VHt), inverse_sqrt + ) + Pr = oblique_globals["ldmul"]( + inverse_sqrt, + oblique_globals["dag"](Ut), + ) @ Rl + + # If an entire charge block is empty, both projectors above are + # zero in that block. Keep a finite identity block so subsequent + # simple-gauge/canonicalization steps do not divide by zero. + for sector, block in st.get_sector_block_pairs(): + if not bool( + ar.do( + "all", + ar.do("abs", block) <= zero_tol, + ) + ): + continue + for projector in (Pl, Pr): + for key, projector_block in projector.get_sector_block_pairs(): + if sector not in key: + continue + identity = ar.do("zeros_like", projector_block) + for i in range(min(identity.shape)): + identity[i, i] = 1.0 + projector.set_block(key, identity) + + return Pl, Pr + + qtc.compute_oblique_projectors = compute_fermionic_oblique_projectors + + if callable(original_reduced_factor): + reduced_impl = getattr(original_reduced_factor, "__wrapped__", None) + reduced_globals = getattr( + original_reduced_factor, + "__globals__", + getattr(reduced_impl, "__globals__", {}), + ) + reduced_dag = reduced_globals.get("dag") + + @wraps(original_reduced_factor) + def compute_fermionic_reduced_factor( + x2, + dl, + dr, + right=True, + method="eigh", + **reduce_opts, + ): + if not is_fermionic(x2): + return original_reduced_factor( + x2, + dl, + dr, + right=right, + method=method, + **reduce_opts, + ) + + _assert_finite_symmray_blocks(x2, name="squared environment") + if callable(reduced_dag): + # Remove tiny anti-Hermitian roundoff before eigh. This is + # deliberately after the finite check: symmetrization must + # never hide a NaN or Inf produced upstream. + x2 = 0.5 * (x2 + reduced_dag(x2)) + _assert_finite_symmray_blocks( + x2, + name="symmetrized squared environment", + ) + + return original_reduced_factor( + x2, + dl, + dr, + right=right, + method=method, + **reduce_opts, + ) + + qtc.squared_op_to_reduced_factor = compute_fermionic_reduced_factor + try: + yield + finally: + tensor_network.insert_compressor_between_regions = original + if original_inplace is not None: + tensor_network.insert_compressor_between_regions_ = original_inplace + if callable(original_oblique): + qtc.compute_oblique_projectors = original_oblique + if callable(original_reduced_factor): + qtc.squared_op_to_reduced_factor = original_reduced_factor + + @dataclass(frozen=True) class BoundaryContractResult: """Structured result from :func:`contract_boundary`. @@ -315,6 +562,110 @@ def _call_with_accepted_kwargs(fn, **kwargs): return fn(**accepted) +def _ctmrg_stabilization_kwargs( + norm, + *, + reduce_opts=None, + gauge_smudge=None, +): + """Prepare numerically safer CTMRG projector options. + + Symmray environments can be very ill-conditioned, especially after the + squared environment used by Quimb's oblique projector construction. For + those networks, add a small positive shift before the Hermitian + factorization and a small gauge smudge by default. Dense networks retain + Quimb's existing defaults unless the caller explicitly supplies options. + """ + if reduce_opts is None: + reduce_opts = {} + else: + try: + reduce_opts = dict(reduce_opts) + except (TypeError, ValueError) as exc: + raise TypeError("ctmrg_reduce_opts must be a mapping or None.") from exc + + symmray = _uses_symmray_arrays(norm) + if symmray: + reduce_opts.setdefault("method", "eigh") + reduce_opts.setdefault("shift", 1.0e-12) + if gauge_smudge is None: + gauge_smudge = 1.0e-10 + + kwargs = {} + if reduce_opts: + kwargs["reduce_opts"] = reduce_opts + if gauge_smudge is not None: + kwargs["gauge_smudge"] = gauge_smudge + if symmray and gauge_smudge is not None: + # ``gauge_smudge`` only reaches projector construction in Quimb. The + # native fermionic path also needs the same floor during the preceding + # simple-gauge normalization, before any reduced-factor decomposition. + kwargs["canonize_opts"] = {"smudge": gauge_smudge} + return kwargs + + +def _freshen_fermionic_dummy_modes(tn): + """Copy ``tn`` with a fresh namespace for Symmray dummy modes. + + Quimb forms a squared environment by conjugating a calculation copy and + contracting it with the original. A double-layer fermionic network can + already contain matching odd dummy modes, so reusing their labels in the + calculation copy creates duplicate same-dual modes before Symmray can + cancel conjugate pairs. Renaming each tensor's dummy modes gives every + copied mode a unique label, while its conjugate receives the matching + dual label during the squared-environment construction. + """ + if not hasattr(tn, "copy"): + return tn + + source = tn.copy() + namespace = id(source) + for tensor in source: + data = getattr(tensor, "data", None) + dummy_modes = getattr(data, "dummy_modes", ()) + if not dummy_modes or not hasattr(data, "modify"): + continue + + fresh_modes = tuple( + type(mode)( + ("__pepsy_ctmrg_dummy__", namespace, id(tensor), mode_i), + dual=mode.dual, + parity=mode.parity, + ) + for mode_i, mode in enumerate(dummy_modes) + ) + data = data.copy() + data.modify(dummy_modes=fresh_modes) + tensor.modify(data=data) + return source + + +def _assert_finite_symmray_blocks(value, *, name): + """Raise before decomposition if a native block already is non-finite.""" + blocks = getattr(value, "blocks", None) + if blocks is None: + try: + finite = bool(ar.do("all", ar.do("isfinite", value))) + except (TypeError, ValueError): + return + if not finite: + raise FloatingPointError( + f"Non-finite values entered the native CTMRG {name}." + ) + return + + for sector, block in blocks.items(): + try: + finite = bool(ar.do("all", ar.do("isfinite", block))) + except (TypeError, ValueError): + continue + if not finite: + raise FloatingPointError( + "Non-finite values entered the native CTMRG " + f"{name} block for charge sector {sector!r}." + ) + + def _unpack_bdy_handle(handle, name): """Unpack a BdyMPS or ``{"bdy": BdyMPS}`` boundary handle.""" holder = handle if isinstance(handle, dict) else None @@ -359,6 +710,8 @@ def _contract_quimb_double_layer( # pylint: disable=too-many-arguments cutoff, equalize_norms, layer_tags, + ctmrg_reduce_opts=None, + ctmrg_gauge_smudge=None, ): """Contract an already-built double-layer TN with a quimb-style method.""" if method == "exact": @@ -416,9 +769,17 @@ def _contract_quimb_double_layer( # pylint: disable=too-many-arguments progbar=progress, inplace=False, ) + kwargs.update( + _ctmrg_stabilization_kwargs( + norm, + reduce_opts=ctmrg_reduce_opts, + gauge_smudge=ctmrg_gauge_smudge, + ) + ) if layer_tags is not None: kwargs["layer_tags"] = list(layer_tags) - return _call_with_accepted_kwargs(contract_fn, **kwargs) + with quimb_ctmrg_projector_compat(): + return _call_with_accepted_kwargs(contract_fn, **kwargs) if method == "hotrg": contract_fn = getattr(norm, "contract_hotrg", None) @@ -465,6 +826,8 @@ def _contract_peps_double_layer( # pylint: disable=too-many-arguments layer_tags=None, bdy_name="bdy", flat=False, + ctmrg_reduce_opts=None, + ctmrg_gauge_smudge=None, ): """Contract a double-layer PEPS norm/overlap network by the selected method.""" method = _normalize_contraction_method(method) @@ -526,6 +889,8 @@ def _contract_peps_double_layer( # pylint: disable=too-many-arguments cutoff=cutoff, equalize_norms=equalize_norms, layer_tags=layer_tags, + ctmrg_reduce_opts=ctmrg_reduce_opts, + ctmrg_gauge_smudge=ctmrg_gauge_smudge, ) return cost, None @@ -550,6 +915,8 @@ def contract_flat( # pylint: disable=too-many-arguments,too-many-positional-arg cutoff=1.0e-12, equalize_norms=False, layer_tags=None, + ctmrg_reduce_opts=None, + ctmrg_gauge_smudge=None, ): """Contract an already-flat PEPS-like tensor network. @@ -575,6 +942,13 @@ def contract_flat( # pylint: disable=too-many-arguments,too-many-positional-arg input network exposes the corresponding method. strip_exponent : bool, default=False If ``True``, return ``(mantissa, exponent)``. + ctmrg_reduce_opts : mapping | None, default=None + Optional options forwarded to Quimb's squared-environment + factorization for ``method="ctmrg"``. Symmray networks receive + ``method="eigh", shift=1e-12`` by default when this is omitted. + ctmrg_gauge_smudge : float | None, default=None + Relative regularization for CTMRG projector environments. Symmray + networks default to ``1e-10`` when this is omitted. Returns ------- @@ -613,6 +987,8 @@ def contract_flat( # pylint: disable=too-many-arguments,too-many-positional-arg layer_tags=layer_tags, bdy_name="bdy", flat=True, + ctmrg_reduce_opts=ctmrg_reduce_opts, + ctmrg_gauge_smudge=ctmrg_gauge_smudge, ) return _format_scaled_output(cost, strip_exponent=strip_exponent) diff --git a/src/pepsy/vmc/torch/amplitude.py b/src/pepsy/vmc/torch/amplitude.py index fdeb861..88a7257 100644 --- a/src/pepsy/vmc/torch/amplitude.py +++ b/src/pepsy/vmc/torch/amplitude.py @@ -10,6 +10,10 @@ import warnings from ..torch_types import _check_positive_int, _require_torch +from ...boundary.metrics import ( + _ctmrg_stabilization_kwargs, + quimb_ctmrg_projector_compat, +) from ._common import ( _as_contraction_options, _as_long_matrix, @@ -497,6 +501,24 @@ def _final_contraction_options(self, *, strip_exponent=None): options.setdefault("strip_exponent", strip_exponent) return options + def _ctmrg_options(self, tnx): + """Return CTMRG options with the native Symmray safety defaults.""" + options = dict(self.contraction_opts) + defaults = _ctmrg_stabilization_kwargs( + tnx, + reduce_opts=options.get("reduce_opts"), + gauge_smudge=options.get("gauge_smudge"), + ) + for key, value in defaults.items(): + if key != "canonize_opts": + if options.get(key) is None: + options[key] = value + continue + canonize_opts = dict(value) + canonize_opts.update(options.get(key) or {}) + options[key] = canonize_opts + return options + def _contract_remaining(self, tn, *args, final_opts=None): """Close an approximate PEPS contraction with the requested path.""" if final_opts is None: @@ -581,12 +603,13 @@ def _contract_value(self, tnx, reference=None): **self.contraction_opts, ) elif self.contraction == "ctmrg": - value = self._contract_approximate( - tnx.contract_ctmrg, - max_bond=self.chi, - close_final=True, - **self.contraction_opts, - ) + with quimb_ctmrg_projector_compat(): + value = self._contract_approximate( + tnx.contract_ctmrg, + max_bond=self.chi, + close_final=True, + **self._ctmrg_options(tnx), + ) elif self.contraction == "boundary": value = self._contract_approximate( tnx.contract_boundary, @@ -609,13 +632,14 @@ def _contract_log_parts(self, tnx, reference=None): **self.contraction_opts, ) elif self.contraction == "ctmrg": - mantissa, exponent_10 = self._contract_approximate( - tnx.contract_ctmrg, - max_bond=self.chi, - strip_exponent=True, - close_final=True, - **self.contraction_opts, - ) + with quimb_ctmrg_projector_compat(): + mantissa, exponent_10 = self._contract_approximate( + tnx.contract_ctmrg, + max_bond=self.chi, + strip_exponent=True, + close_final=True, + **self._ctmrg_options(tnx), + ) elif self.contraction == "boundary": mantissa, exponent_10 = self._contract_approximate( tnx.contract_boundary, diff --git a/tests/test_prepare_boundary_inputs.py b/tests/test_prepare_boundary_inputs.py index 97ecc0f..7801220 100644 --- a/tests/test_prepare_boundary_inputs.py +++ b/tests/test_prepare_boundary_inputs.py @@ -1491,6 +1491,125 @@ def contract_hotrg(self, **kwargs): assert call_kwargs["final_contract_opts"]["optimize"] == "OPT" +def test_contract_flat_ctmrg_enters_projector_compatibility_scope(monkeypatch): + """The shared flat CTMRG route enables the scoped Quimb workaround.""" + events = [] + + class _Scope: + def __enter__(self): + events.append("enter") + + def __exit__(self, exc_type, exc, tb): + events.append("exit") + + class _FlatTN: + Lx = 2 + Ly = 2 + + def contract_ctmrg(self, **kwargs): + events.append("contract") + return 2.0 + + monkeypatch.setattr( + pepsy.boundary.metrics, + "quimb_ctmrg_projector_compat", + lambda: _Scope(), + ) + + assert pepsy.contract_flat(_FlatTN(), method="ctmrg", chi=4) == 2.0 + assert events == ["enter", "contract", "exit"] + + +def test_contract_flat_ctmrg_forwards_stabilization_options(): + """CTMRG stabilization controls should reach Quimb's projector path.""" + captured = {} + + class _FlatTN: + Lx = 2 + Ly = 2 + + def contract_ctmrg(self, **kwargs): + captured.update(kwargs) + return 2.0 + + reduce_opts = {"method": "cholesky", "shift": 1.0e-9} + out = pepsy.contract_flat( + _FlatTN(), + method="ctmrg", + chi=4, + ctmrg_reduce_opts=reduce_opts, + ctmrg_gauge_smudge=2.0e-8, + ) + + assert out == 2.0 + assert captured["reduce_opts"] == reduce_opts + assert captured["gauge_smudge"] == 2.0e-8 + assert reduce_opts == {"method": "cholesky", "shift": 1.0e-9} + + +def test_contract_flat_ctmrg_adds_symmray_stabilization_defaults(monkeypatch): + """Symmray CTMRG should get a shift and gauge smudge by default.""" + captured = {} + + class _FlatTN: + Lx = 2 + Ly = 2 + + def contract_ctmrg(self, **kwargs): + captured.update(kwargs) + return 2.0 + + monkeypatch.setattr( + pepsy.boundary.metrics, + "_uses_symmray_arrays", + lambda tn: tn is not None, + ) + + assert pepsy.contract_flat(_FlatTN(), method="ctmrg", chi=4) == 2.0 + assert captured["reduce_opts"] == {"method": "eigh", "shift": 1.0e-12} + assert captured["gauge_smudge"] == 1.0e-10 + assert captured["canonize_opts"] == {"smudge": 1.0e-10} + + +def test_quimb_ctmrg_projector_compat_uses_live_insert_target(monkeypatch): + """CTMRG projector insertion should use the current network snapshot.""" + import quimb.tensor.tensor_core as qtc + + calls = [] + def fake_insert(self, ltags, rtags, *args, insert_into=None, **kwargs): + calls.append((self, ltags, rtags, insert_into, kwargs)) + return self + + monkeypatch.setattr( + qtc.TensorNetwork, + "insert_compressor_between_regions", + fake_insert, + ) + + stale = object() + current = object() + with pepsy.boundary.metrics.quimb_ctmrg_projector_compat(): + qtc.TensorNetwork.insert_compressor_between_regions( + stale, + ("L",), + ("R",), + insert_into=current, + ) + qtc.TensorNetwork.insert_compressor_between_regions_( + stale, + ("L",), + ("R",), + insert_into=current, + ) + + assert len(calls) == 2 + for call in calls: + assert call[0] is current + assert call[3] is current + assert call[4]["inplace"] is True + assert qtc.TensorNetwork.insert_compressor_between_regions is fake_insert + + def test_infidelity_accepts_bdy_holder_dicts_and_fills_missing(monkeypatch): """infidelity should accept dict holders and populate missing bdy entries.""" created = [] diff --git a/tests/test_symmetric_tensors.py b/tests/test_symmetric_tensors.py index e8a2095..9f0ced0 100644 --- a/tests/test_symmetric_tensors.py +++ b/tests/test_symmetric_tensors.py @@ -4285,6 +4285,82 @@ def test_sympeps_measure_delegates_to_quimb_boundary_modes(): assert ctmrg_norm == pytest.approx(exact_norm) +def test_native_fermionic_ctmrg_matches_exact_on_small_double_layer(): + """Native fermionic CTMRG should remain finite on a small U1U1 PEPS.""" + site_charge = site_charge_from_occupations( + { + (0, 0): (1, 0), + (0, 1): (0, 1), + (1, 0): (0, 1), + (1, 1): (1, 0), + (2, 0): (1, 0), + (2, 1): (0, 1), + } + ) + state = SymPEPS.random( + 3, + 2, + symmetry="U1U1", + phys_dim=default_physical_sectors(model="fermi_hubbard_u1u1"), + fermionic=True, + site_charge=site_charge, + bond_dim=2, + seed=74, + dtype="complex128", + ) + norm = state.tn.make_norm() + exact = norm.contract(all, optimize="auto-hq") + ctmrg = pepsy.contract_flat( + norm, + chi=2, + method="ctmrg", + progress=False, + cutoff=1.0e-10, + ) + + assert np.isfinite(ctmrg) + assert ctmrg == pytest.approx(exact, rel=1.0e-6, abs=1.0e-12) + + +def test_native_ctmrg_rejects_nonfinite_squared_environment_before_eigh(): + """Non-finite native environment blocks must stop before factorization.""" + import quimb.tensor.tensor_core as qtc + + site_charge = site_charge_from_occupations( + { + (0, 0): (1, 0), + (0, 1): (0, 1), + (1, 0): (1, 0), + (1, 1): (0, 1), + } + ) + state = SymPEPS.random( + 2, + 2, + symmetry="U1U1", + phys_dim=default_physical_sectors(model="fermi_hubbard_u1u1"), + fermionic=True, + site_charge=site_charge, + bond_dim=2, + seed=75, + dtype="complex128", + ) + environment = next(iter(state.peps.tensor_map.values())).data.copy() + sector = next(iter(environment.blocks)) + environment.set_block( + sector, + np.full_like(environment.get_block(sector), np.nan), + ) + + with pepsy.boundary.metrics.quimb_ctmrg_projector_compat(): + with pytest.raises(FloatingPointError, match="squared environment"): + qtc.squared_op_to_reduced_factor( + environment, + environment.shape[0], + environment.shape[0], + ) + + def test_sympeps_gate_stream_runs_pepsy_gate_and_gate_simple(): """SymPEPS gate streams should work with PEPSY gate wrappers.""" state = SymPEPS.for_model( diff --git a/tests/test_vmc_api.py b/tests/test_vmc_api.py index 513c733..332f022 100644 --- a/tests/test_vmc_api.py +++ b/tests/test_vmc_api.py @@ -946,6 +946,54 @@ def test_torch_amplitude_accepts_common_contraction_config(): assert model.chi == 4 +def test_torch_ctmrg_uses_symmray_stabilization_defaults(monkeypatch): + """Torch CTMRG should share the native Symmray safety controls.""" + pytest.importorskip("torch") + import pepsy.boundary.metrics as metrics + from pepsy.vmc.torch.amplitude import TorchPEPSAmplitude + + model = object.__new__(TorchPEPSAmplitude) + model.contraction_opts = {"mode": "direct"} + monkeypatch.setattr(metrics, "_uses_symmray_arrays", lambda value: True) + + options = model._ctmrg_options(object()) + + assert options["mode"] == "direct" + assert options["reduce_opts"] == { + "method": "eigh", + "shift": 1.0e-12, + } + assert options["gauge_smudge"] == 1.0e-10 + assert options["canonize_opts"] == {"smudge": 1.0e-10} + + +def test_torch_ctmrg_preserves_explicit_stabilization_options(monkeypatch): + """Explicit Torch CTMRG controls should override only the defaults.""" + pytest.importorskip("torch") + import pepsy.boundary.metrics as metrics + from pepsy.vmc.torch.amplitude import TorchPEPSAmplitude + + model = object.__new__(TorchPEPSAmplitude) + model.contraction_opts = { + "reduce_opts": {"method": "cholesky", "shift": 1.0e-9}, + "gauge_smudge": 2.0e-8, + "canonize_opts": {"max_bond": 4}, + } + monkeypatch.setattr(metrics, "_uses_symmray_arrays", lambda value: True) + + options = model._ctmrg_options(object()) + + assert options["reduce_opts"] == { + "method": "cholesky", + "shift": 1.0e-9, + } + assert options["gauge_smudge"] == 2.0e-8 + assert options["canonize_opts"] == { + "smudge": 2.0e-8, + "max_bond": 4, + } + + def test_netket_setup_consumes_shared_sampling_config(): nk = pytest.importorskip("netket") from pepsy.vmc.netket import NetKetPEPSVMC From b1c7207cefbda48267aeaf668b01a2fe9c8027fd Mon Sep 17 00:00:00 2001 From: rezaquant Date: Tue, 4 Aug 2026 12:17:56 -0600 Subject: [PATCH 67/70] Stabilize CTMRG compatibility and MPO backend casting --- AGENTS.md | 17 ++++++++ plan.md | 75 ++++++++++++++++++++++++++++++++++ src/pepsy/tensors/symmetric.py | 8 +++- src/pepsy/vmc/netket.py | 51 ++++++++++++----------- 4 files changed, 125 insertions(+), 26 deletions(-) create mode 100644 plan.md diff --git a/AGENTS.md b/AGENTS.md index 967b725..ab8a7d9 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -127,6 +127,23 @@ The tree-native operator API lives in `pepsy.optimizers.tree.operators`: invariants. Emit explicit warnings for intentional compatibility coercions. - Do not vendor upstream internals. +## Cyclic CTMRG compatibility + +- Use `pepsy.boundary.quimb_ctmrg_projector_compat` around Quimb CTMRG calls + for cyclic PEPS/PEPO networks whose effective bond dimensions can vary, + especially native U(1) term-by-term replays. +- This is a scoped compatibility context: it redirects projector insertion to + the current network and restores Quimb's method on exit. It does not modify + installed `site-packages`, alter boundary-MPS contractions, or replace + CTMRG with MPS. +- Keep the workaround at the Pepsy boundary API. Do not copy Quimb's CTMRG + implementation into Pepsy or edit the installed Quimb source. Add a focused + regression test when changing the compatibility behavior. +- The shared CTMRG entry points already apply this context for + `contract_flat(..., method="ctmrg")`, native Torch PEPS VMC models with + `contraction="ctmrg"`, and the NetKet/JAX PEPS amplitude validation path. + Keep exact, HOTRG, and boundary-MPS routes independent of this workaround. + ## Documentation and skills - Keep user-facing API docs under `docs/api/` and concise implementation maps diff --git a/plan.md b/plan.md new file mode 100644 index 0000000..ece5317 --- /dev/null +++ b/plan.md @@ -0,0 +1,75 @@ +# Native fermionic PEPO plan + +Status: deferred design item. No implementation is planned until this work is +prioritized. + +## Objective + +Add a direct two-dimensional native fermionic PEPO builder for sums of local +and two-body fermionic operator terms. The builder must use Symmray graded +tensors and must not introduce Jordan--Wigner strings or dense intermediate +operators. + +## Planned approach + +1. Accept one- and two-site terms as native Symmray fermionic arrays with + `U1`, `U1U1`, or `Z2` symmetry. +2. Factor each one- or two-site term exactly using graded operator-Schmidt + decomposition. Keep every channel; do not truncate or numerically + compress. +3. Build local PEPO tensors from a shared finite-state automaton. Virtual + states represent idle propagation, opened/closed fermionic operators, + charge and spin sectors, and coefficient insertion. +4. Construct tensors directly from Symmray block structure. Charge flow and + graded contraction supply the fermionic signs natively. +5. Use a small constant-channel automaton for local or finite-range terms. +6. Use a Crosswhite--Bacon / Fröwis--Nebendahl--Dür style 2D automaton for + arbitrary long-range two-body coefficients. Share horizontal and vertical + propagation rather than routing one independent path per term. + +## Expected scaling + +- Local and finite-range interactions: constant bond dimension, up to factors + from the operator-Schmidt rank and symmetry channels. +- Generic arbitrary two-body coefficients on an `L x L` lattice: linear + virtual-bond scaling in `L` for the exact 2D automaton, with additional + constant factors for spin and charge sectors. +- Special structured kernels, such as separable or distance-dependent + interactions, can use smaller specialized automata or sums of auxiliary + PEPOs. + +The construction is exact and has optimal scaling for the generic pairwise +family treated by Fröwis et al.; it is not a universal proof of globally +minimal PEPO bond dimension for every Hamiltonian. + +## Implementation sequence + +1. Implement and test native local one-/two-site PEPO terms. +2. Add nearest-neighbor and finite-range term sums with shared channels. +3. Add the exact long-range 2D automaton. +4. Validate `U1U1` hopping, onsite, density, and long-range terms on small + lattices against the existing native snake-MPO PEPO route. +5. Add structured-kernel strategies only after the generic exact builder is + stable. + +## Possible API + +```python +fermion.to_pepo( + terms, + shape=(Lx, Ly), + strategy="native_2d", + long_range="automaton", +) +``` + +## References + +- Crosswhite and Bacon, *Finite automata for caching in matrix product + algorithms*, [arXiv:0708.1221](https://arxiv.org/abs/0708.1221). +- Fröwis, Nebendahl, and Dür, *Tensor operators: constructions and + applications for long-range interaction systems*, + [arXiv:1003.1047](https://arxiv.org/abs/1003.1047). +- O'Rourke, Li, and Chan, *Efficient representation of long-range + interactions in tensor network algorithms*, + [arXiv:1807.08378](https://arxiv.org/abs/1807.08378). diff --git a/src/pepsy/tensors/symmetric.py b/src/pepsy/tensors/symmetric.py index 93e49ef..34c4e26 100644 --- a/src/pepsy/tensors/symmetric.py +++ b/src/pepsy/tensors/symmetric.py @@ -5574,8 +5574,6 @@ def _assemble_symmray_mpo( lower_ind_id=lower_ind_id, site_tag_id=site_tag_id, ) - if to_backend is not None: - _apply_to_tensor_network_arrays(mpo, to_backend) raw_bond = mpo.max_bond() raw_max_bond = 1 if raw_bond is None else int(raw_bond) did_compress = bool(compress and L > 1) @@ -5584,6 +5582,12 @@ def _assemble_symmray_mpo( if max_bond is not None: compress_opts["max_bond"] = int(max_bond) mpo.compress(**compress_opts) + if to_backend is not None: + # Cast after compression so the SVD-based bond truncation runs in the + # stable build precision (e.g. complex128). Converting first and then + # compressing runs the SVD in the target precision, which for a + # near-singular Hamiltonian MPO in complex64 can hit non-finite values. + _apply_to_tensor_network_arrays(mpo, to_backend) requested_max_bond = None if max_bond is None else int(max_bond) final_bond = mpo.max_bond() diff --git a/src/pepsy/vmc/netket.py b/src/pepsy/vmc/netket.py index 00bfb24..a96efdd 100644 --- a/src/pepsy/vmc/netket.py +++ b/src/pepsy/vmc/netket.py @@ -2149,38 +2149,41 @@ def _contract_ctmrg_for_vmc(tn, *, max_bond, cutoff, method_opts): ``chi`` paired with a Z2 block axis). Keep all requested options active and relax only this stopping threshold to Quimb's stable value ``1``. """ + from ..boundary.metrics import quimb_ctmrg_projector_compat + global _FLAT_SYMMRAY_CTMRG_FALLBACK_WARNED kwargs = dict(method_opts) - try: - return tn.contract_ctmrg( - max_bond=max_bond, - cutoff=cutoff, - strip_exponent=True, - **kwargs, - ) - except (AttributeError, TypeError, ValueError): - if ( - kwargs.get("max_separation", 1) == 0 - and _is_flat_symmray_network(tn) - ): - if not _FLAT_SYMMRAY_CTMRG_FALLBACK_WARNED: - warnings.warn( - "Flat Symmray JAX CTMRG does not support the " - "max_separation=0 intermediate axis path; retrying " - "with max_separation=1. The requested sequence, chi, " - "and canonization options remain active.", - RuntimeWarning, - stacklevel=3, - ) - _FLAT_SYMMRAY_CTMRG_FALLBACK_WARNED = True - kwargs["max_separation"] = 1 + with quimb_ctmrg_projector_compat(): + try: return tn.contract_ctmrg( max_bond=max_bond, cutoff=cutoff, strip_exponent=True, **kwargs, ) - raise + except (AttributeError, TypeError, ValueError): + if ( + kwargs.get("max_separation", 1) == 0 + and _is_flat_symmray_network(tn) + ): + if not _FLAT_SYMMRAY_CTMRG_FALLBACK_WARNED: + warnings.warn( + "Flat Symmray JAX CTMRG does not support the " + "max_separation=0 intermediate axis path; retrying " + "with max_separation=1. The requested sequence, chi, " + "and canonization options remain active.", + RuntimeWarning, + stacklevel=3, + ) + _FLAT_SYMMRAY_CTMRG_FALLBACK_WARNED = True + kwargs["max_separation"] = 1 + return tn.contract_ctmrg( + max_bond=max_bond, + cutoff=cutoff, + strip_exponent=True, + **kwargs, + ) + raise def _resolve_netket_contraction(contraction, chi, cutoff, contraction_opts): From 645ef1d7be7ab9ee9fb79f9967232742334fb3b1 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Tue, 4 Aug 2026 12:42:41 -0600 Subject: [PATCH 68/70] Handle CuPy arrays in dense NumPy conversion --- src/pepsy/tensors/symmetric.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/pepsy/tensors/symmetric.py b/src/pepsy/tensors/symmetric.py index 34c4e26..94e3a85 100644 --- a/src/pepsy/tensors/symmetric.py +++ b/src/pepsy/tensors/symmetric.py @@ -5171,6 +5171,9 @@ def _dense_numpy(value, *, dtype=None): cpu = getattr(value, "cpu", None) if callable(cpu): value = cpu() + # CuPy arrays reject implicit host conversion; move to host explicitly. + if type(value).__module__.split(".", 1)[0] == "cupy": + value = value.get() return np.asarray(value, dtype=dtype) From 0b69cb068b689a8b34d4f11929a546655038128c Mon Sep 17 00:00:00 2001 From: rezaquant Date: Wed, 5 Aug 2026 09:36:01 -0600 Subject: [PATCH 69/70] Cache VMC transition amplitudes --- src/pepsy/backends/config.py | 12 +- src/pepsy/backends/convert.py | 39 +-- src/pepsy/boundary/metrics.py | 16 +- src/pepsy/bp/_symmray.py | 9 +- src/pepsy/bp/gauges.py | 13 +- src/pepsy/bp/series.py | 6 +- src/pepsy/bp/weights.py | 24 +- src/pepsy/operators/gates.py | 10 +- src/pepsy/operators/hamiltonians.py | 5 +- src/pepsy/optimizers/mpo/optimizer.py | 2 +- src/pepsy/optimizers/mps/layout.py | 18 +- src/pepsy/optimizers/mps/optimizer.py | 13 +- .../stabilizer_tn/mps_stab_optimizer.py | 11 +- .../optimizers/stabilizer_tn/operators.py | 3 +- src/pepsy/optimizers/sweep/optimizer.py | 7 +- src/pepsy/optimizers/sym_dmrg.py | 11 +- src/pepsy/optimizers/tree/operators.py | 10 +- .../optimizers/tree_stabilizer/optimizer.py | 10 +- src/pepsy/sampling/samplers.py | 47 ++-- src/pepsy/sampling/tree.py | 6 +- src/pepsy/tensors/symmetric.py | 37 +-- src/pepsy/vmc/__init__.py | 2 + src/pepsy/vmc/netket.py | 15 +- src/pepsy/vmc/torch/__init__.py | 10 +- src/pepsy/vmc/torch/_core.py | 2 + src/pepsy/vmc/torch/_graded.py | 3 +- src/pepsy/vmc/torch/amplitude.py | 15 +- src/pepsy/vmc/torch/cache.py | 229 ++++++++++++++++++ src/pepsy/vmc/torch/connections.py | 103 +++++++- src/pepsy/vmc/torch/driver.py | 165 +++++++++++-- src/pepsy/vmc/torch/fermion.py | 24 +- src/pepsy/vmc/torch/importance.py | 23 +- tests/test_vmc_api.py | 30 +++ tests/test_vmc_transition_plan.py | 74 ++++++ 34 files changed, 784 insertions(+), 220 deletions(-) create mode 100644 src/pepsy/vmc/torch/cache.py create mode 100644 tests/test_vmc_transition_plan.py diff --git a/src/pepsy/backends/config.py b/src/pepsy/backends/config.py index 54f809d..c7573e6 100644 --- a/src/pepsy/backends/config.py +++ b/src/pepsy/backends/config.py @@ -306,10 +306,14 @@ def _resolve_device(dev): target_dtype = jax.dtypes.canonicalize_dtype(jnp.dtype(dtype)) def cast_array(x, device=target_device, dtype=target_dtype): - # Coerce non-JAX inputs (incl. torch tensors) to a numpy-compatible - # form first so jnp.asarray accepts them on any backend. - if torch is not None and isinstance(x, torch.Tensor): - x = x.detach().cpu().numpy() + # Coerce non-JAX inputs to a NumPy-compatible form through Autoray so + # Torch, CuPy, and other registered backends share one host boundary. + try: + if ar.infer_backend(x) != "jax": + x = ar.to_numpy(x) + except Exception: + # Let JAX handle custom array-likes that Autoray cannot infer. + pass arr = jnp.asarray(x, dtype=dtype) if device is not None: arr = jax.device_put(arr, device) diff --git a/src/pepsy/backends/convert.py b/src/pepsy/backends/convert.py index af2bd31..b194ef6 100644 --- a/src/pepsy/backends/convert.py +++ b/src/pepsy/backends/convert.py @@ -26,15 +26,15 @@ def _backend_scalar(value): if shape != (): raise TypeError(f"Expected a scalar-like value, got shape {shape}.") - obj = value - for method_name in ("detach", "cpu"): - method = getattr(obj, method_name, None) - if callable(method): - obj = method() - - get = getattr(obj, "get", None) - if callable(get) and shape is not None: - obj = get() + if shape is not None: + try: + obj = ar.to_numpy(value) + except Exception: + # Keep supporting duck-typed scalar wrappers with ``item`` but no + # registered Autoray backend. + obj = value + else: + obj = value item = getattr(obj, "item", None) if callable(item) and not isinstance(obj, _SCALAR_TYPES): @@ -56,9 +56,8 @@ def to_float(value, *, real=True): """Convert a scalar-like backend value to a Python ``float``. The input can be a Python scalar, NumPy scalar or scalar array, or a - scalar-like backend tensor. Torch-style values are detached and moved to - CPU before extracting ``.item()``; CuPy-style values are host-transferred - with ``.get()`` when available. Non-scalar arrays raise ``TypeError``. + scalar-like backend tensor. Autoray converts backend scalar arrays to host + NumPy before extracting ``.item()``. Non-scalar arrays raise ``TypeError``. Parameters ---------- @@ -261,7 +260,7 @@ def _build_to_numpy(sample_data, dtype_name, *, cast_complex_to_real=False): dtype = _NUMPY_DTYPE_MAP[dtype_name] def _to_numpy(x, dtype=dtype, cast_complex_to_real=cast_complex_to_real): - arr = np.asarray(x) + arr = np.asarray(ar.to_numpy(x)) if cast_complex_to_real and np.issubdtype(dtype, np.floating) and np.iscomplexobj(arr): arr = arr.real target_dtype = dtype @@ -389,13 +388,15 @@ def _to_jax( device=device, cast_complex_to_real=cast_complex_to_real, ): - # Torch tensors need explicit host conversion before jnp.asarray. try: - import torch # pylint: disable=import-outside-toplevel - except ImportError: # pragma: no cover - optional dependency - torch = None - if torch is not None and isinstance(x, torch.Tensor): - x = x.detach().cpu().numpy() + # JAX accepts NumPy inputs directly, while Autoray provides the + # explicit host boundary for Torch/CuPy and other array backends. + if ar.infer_backend(x) != "jax": + x = ar.to_numpy(x) + except Exception: + # Preserve support for custom array-likes which JAX can consume + # even though Autoray cannot infer their namespace. + pass arr = jnp.asarray(x) diff --git a/src/pepsy/boundary/metrics.py b/src/pepsy/boundary/metrics.py index 2b6d982..e597373 100644 --- a/src/pepsy/boundary/metrics.py +++ b/src/pepsy/boundary/metrics.py @@ -335,14 +335,16 @@ def _uses_symmray_arrays(tn): def _to_python_scalar(value): """Convert backend scalar-like objects (torch/numpy) to python scalar.""" - obj = value - if hasattr(obj, "detach"): - obj = obj.detach() - if hasattr(obj, "cpu"): - obj = obj.cpu() - if hasattr(obj, "item") and not isinstance(obj, (int, float, complex, bool)): + if isinstance(value, (int, float, complex, bool)): + return value + try: + obj = ar.to_numpy(value) + except Exception: + obj = value + item = getattr(obj, "item", None) + if callable(item): try: - obj = obj.item() + return item() except (ValueError, RuntimeError): # backend-specific .item() failures pass return obj diff --git a/src/pepsy/bp/_symmray.py b/src/pepsy/bp/_symmray.py index 928cd53..c6d6e3c 100644 --- a/src/pepsy/bp/_symmray.py +++ b/src/pepsy/bp/_symmray.py @@ -11,6 +11,7 @@ from collections.abc import Mapping +import autoray as ar import numpy as np @@ -107,8 +108,8 @@ def dense_message_tree(messages): def to_dense(value): """Materialize one small Symmray value for a local operator calculation.""" if hasattr(value, "to_dense"): - return np.asarray(value.to_dense()) - return np.asarray(value) + value = value.to_dense() + return np.asarray(ar.to_numpy(value)) def dense_index_map(chargemap): @@ -275,7 +276,9 @@ def rank4_operator_from_dense(tn, index, operator, *, layout="pne"): _bond_endpoint_data(tn, index) ) dimension = int(left_data.shape[left_data.indices.index(left_index)]) - dense = np.asarray(operator) + if hasattr(operator, "to_dense"): + operator = operator.to_dense() + dense = np.asarray(ar.to_numpy(operator)) if dense.shape == (dimension * dimension, dimension * dimension): dense = dense.reshape(dimension, dimension, dimension, dimension) elif dense.shape != (dimension, dimension, dimension, dimension): diff --git a/src/pepsy/bp/gauges.py b/src/pepsy/bp/gauges.py index d26fd51..d638e27 100644 --- a/src/pepsy/bp/gauges.py +++ b/src/pepsy/bp/gauges.py @@ -105,17 +105,12 @@ def _copy_array(x): def _as_numpy(x): if hasattr(x, "to_dense"): - return np.asarray(x.to_dense()) - try: - return np.asarray(ar.to_numpy(x)) - except Exception: - return np.asarray(x) + x = x.to_dense() + return np.asarray(ar.to_numpy(x)) def _gauge_values_numpy(gauge): """Materialize a gauge vector for validation, including Symmray vectors.""" - if hasattr(gauge, "to_dense"): - return np.asarray(gauge.to_dense()) return _as_numpy(gauge) @@ -126,8 +121,8 @@ def _is_symmray_array(value) -> bool: def _symmray_dense_matrix(matrix): """Convert one small Symmray message to a host dense matrix.""" if hasattr(matrix, "to_dense"): - return np.asarray(matrix.to_dense()) - return np.asarray(matrix) + matrix = matrix.to_dense() + return np.asarray(ar.to_numpy(matrix)) def _symmray_align_message_to_bond(tn, ix, tid, message): diff --git a/src/pepsy/bp/series.py b/src/pepsy/bp/series.py index ff86921..07272e9 100644 --- a/src/pepsy/bp/series.py +++ b/src/pepsy/bp/series.py @@ -7374,8 +7374,10 @@ def adaptive_open_loop_series( differences.append(None) continue try: - difference = float(np.max(np.abs(np.asarray(value) - np.asarray(values[-2])))) - scale = max(1.0, float(np.max(np.abs(np.asarray(value))))) + value_array = np.asarray(ar.to_numpy(value)) + previous_array = np.asarray(ar.to_numpy(values[-2])) + difference = float(np.max(np.abs(value_array - previous_array))) + scale = max(1.0, float(np.max(np.abs(value_array)))) except (TypeError, ValueError): difference = None scale = 1.0 diff --git a/src/pepsy/bp/weights.py b/src/pepsy/bp/weights.py index fc1604d..7d1684f 100644 --- a/src/pepsy/bp/weights.py +++ b/src/pepsy/bp/weights.py @@ -17,6 +17,7 @@ from dataclasses import dataclass, field from typing import Any +import autoray as ar import numpy as np from ._symmray import dense_bp_tn as _dense_bp_tn @@ -41,7 +42,7 @@ def _apply_left(data, axis, matrix): def _multiply_axis(data, axis, weights): shape = [1] * data.ndim shape[axis] = len(weights) - return data * np.asarray(weights).reshape(shape) + return data * np.asarray(ar.to_numpy(weights)).reshape(shape) def _validate_network(tn): @@ -70,7 +71,7 @@ def _initial_weights(tn, weights): for index in tn.inner_inds(): dimension = _dimension(tn.tensor_map[next(iter(tn.ind_map[index]))], index) value = supplied.get(index, np.ones(dimension, dtype=float)) - value = np.asarray(value) + value = np.asarray(ar.to_numpy(value)) if value.ndim == 2: if value.shape != (dimension, dimension): raise ValueError( @@ -92,7 +93,10 @@ def _initial_weights(tn, weights): def _weighted_tensor(tensor, weights, exclude): - data = np.asarray(tensor.data) + data = tensor.data + if hasattr(data, "to_dense"): + data = data.to_dense() + data = np.asarray(ar.to_numpy(data)) for axis, index in enumerate(tensor.inds): if index != exclude: data = _multiply_axis(data, axis, weights[index]) @@ -172,8 +176,12 @@ def _bond_update(tn, index, weights, *, alpha, eps): @ right_u.conj().T ) - left_data = _apply_right(np.asarray(tensor_left.data), left_axis, left_gauge) - right_data = _apply_left(np.asarray(tensor_right.data), right_axis, right_gauge) + left_data = _apply_right( + np.asarray(ar.to_numpy(tensor_left.data)), left_axis, left_gauge + ) + right_data = _apply_left( + np.asarray(ar.to_numpy(tensor_right.data)), right_axis, right_gauge + ) # The normalized weight differs from S_C**alpha by ``scale``. The # compensating scalar keeps the represented network exactly unchanged. left_data = left_data * scale @@ -307,7 +315,11 @@ def weight_pass( _, right = tuple(network.ind_map[index]) tensor = network.tensor_map[right] axis = tensor.inds.index(index) - tensor.modify(data=_apply_left(np.asarray(tensor.data), axis, np.diag(value))) + tensor.modify( + data=_apply_left( + np.asarray(ar.to_numpy(tensor.data)), axis, np.diag(value) + ) + ) return WeightPassingResult( network=network, diff --git a/src/pepsy/operators/gates.py b/src/pepsy/operators/gates.py index a0ba780..f0ec48f 100644 --- a/src/pepsy/operators/gates.py +++ b/src/pepsy/operators/gates.py @@ -1254,7 +1254,9 @@ def _symmray_dense_gate_from_site_maps( else 0 ) - dense_gate = np.asarray(gate) + if hasattr(gate, "to_dense"): + gate = gate.to_dense() + dense_gate = np.asarray(ar.to_numpy(gate)) if inferred_converter is not None: dense_gate = inferred_converter(dense_gate) @@ -2164,14 +2166,14 @@ def gate_loop_cluster( if any(hasattr(tensor.data, "to_dense") for tensor in tn_work.tensor_map.values()): for tensor in tn_work.tensor_map.values(): if hasattr(tensor.data, "to_dense"): - tensor.modify(data=np.asarray(tensor.data.to_dense())) + tensor.modify(data=np.asarray(ar.to_numpy(tensor.data.to_dense()))) for index, gauge in tuple(gauges.items()): if hasattr(gauge, "to_dense"): - gauges[index] = np.asarray(gauge.to_dense()) + gauges[index] = np.asarray(ar.to_numpy(gauge.to_dense())) for gate_payload, where_payload, which_payload in entries: if hasattr(gate_payload, "to_dense"): - gate_payload = np.asarray(gate_payload.to_dense()) + gate_payload = np.asarray(ar.to_numpy(gate_payload.to_dense())) if _is_explicit_index_where(where_payload): raise ValueError( diff --git a/src/pepsy/operators/hamiltonians.py b/src/pepsy/operators/hamiltonians.py index 38739b2..6eef600 100644 --- a/src/pepsy/operators/hamiltonians.py +++ b/src/pepsy/operators/hamiltonians.py @@ -7,6 +7,7 @@ from numbers import Integral +import autoray as ar import numpy as np import quimb import quimb.tensor as qtn @@ -348,7 +349,9 @@ def _site_tensor(op, site, L): @staticmethod def _as_matrix(op): data = getattr(op, "data", op) - return np.asarray(data) + if hasattr(data, "to_dense"): + data = data.to_dense() + return np.asarray(ar.to_numpy(data)) def _coerce_op(self, op, *, phys_dim, dtype): if callable(op) and not hasattr(op, "shape"): diff --git a/src/pepsy/optimizers/mpo/optimizer.py b/src/pepsy/optimizers/mpo/optimizer.py index af22cc6..68de350 100644 --- a/src/pepsy/optimizers/mpo/optimizer.py +++ b/src/pepsy/optimizers/mpo/optimizer.py @@ -398,7 +398,7 @@ def _fermionic_gate_to_bosonic(cls, p, gate, where, ind_id): dense = gate.to_dense() except AttributeError: dense = gate - dense = np.asarray(dense) + dense = np.asarray(ar.to_numpy(dense)) where = tuple(where) physical_maps = [ cls._symmray_physical_map(p, site, ind_id) for site in where diff --git a/src/pepsy/optimizers/mps/layout.py b/src/pepsy/optimizers/mps/layout.py index 6f3a796..235ca3d 100644 --- a/src/pepsy/optimizers/mps/layout.py +++ b/src/pepsy/optimizers/mps/layout.py @@ -7,6 +7,7 @@ from numbers import Integral import os +import autoray as ar import numpy as np from ...operators.gates import _normalize_gate_entries @@ -216,10 +217,14 @@ def _payload_angle(payload): else: return None - if isinstance(value, (tuple, list, np.ndarray)): - if len(value) != 1: + try: + value_array = np.asarray(ar.to_numpy(value)) + except Exception: + value_array = None + if value_array is not None: + if value_array.size != 1: return None - value = value[0] + value = value_array.reshape(-1)[0] try: angle = abs(float(value)) except (TypeError, ValueError): @@ -240,7 +245,8 @@ def _operator_schmidt_weight(payload, support, *, schmidt_max_dim=4): if len(support) != 2: return None try: - array = np.asarray(payload) + raw = payload.to_dense() if hasattr(payload, "to_dense") else payload + array = np.asarray(ar.to_numpy(raw)) except Exception: return None if array.size == 0 or not np.issubdtype(array.dtype, np.number): @@ -339,8 +345,10 @@ def _operator_schmidt_rank_info( } raw = getattr(payload, "data", payload) + if hasattr(raw, "to_dense"): + raw = raw.to_dense() try: - array = np.asarray(raw) + array = np.asarray(ar.to_numpy(raw)) except Exception: return {"rank": default_bound, "exact": False, "reason": "opaque"} if array.size == 0 or not np.issubdtype(array.dtype, np.number): diff --git a/src/pepsy/optimizers/mps/optimizer.py b/src/pepsy/optimizers/mps/optimizer.py index 20c4525..ef313d4 100644 --- a/src/pepsy/optimizers/mps/optimizer.py +++ b/src/pepsy/optimizers/mps/optimizer.py @@ -371,7 +371,7 @@ def _parse_control_tuple(name, entry, default_axis=None): if len(entry) < 3: raise ValueError("cap event must be ('cap', where, vec[, absorb]).") where = _normalize_control_where(entry[1], single=True) - vec = np.asarray(entry[2], dtype=complex).ravel() + vec = np.asarray(ar.to_numpy(entry[2]), dtype=complex).ravel() absorb = _normalize_absorb(entry[3]) if len(entry) > 3 else "left" return "cap", {"vec": vec, "absorb": absorb}, where if name == "reset": @@ -1357,7 +1357,10 @@ def _mps_data_is_finite(p): """Return whether dense numeric tensor data contains only finite values.""" for tensor in getattr(p, "tensors", ()): try: - data = np.asarray(tensor.data) + data = tensor.data + if hasattr(data, "to_dense"): + data = data.to_dense() + data = np.asarray(ar.to_numpy(data)) except Exception: continue if not np.issubdtype(data.dtype, np.number): @@ -1735,7 +1738,7 @@ def remap_sample(self, config): self.logical_site(position): value for position, value in config.items() } - config = np.asarray(config) + config = np.asarray(ar.to_numpy(config)) if config.ndim == 0 or config.shape[-1] != len(self.logical_order): raise ValueError( "sample configuration must have MPS length as its final " @@ -2960,7 +2963,7 @@ def _to_state_backend(self, array): """Return ``array`` cast to the backend and dtype owned by ``self.p``.""" like = self._state_backend_like() if like is None: - return np.asarray(array, dtype=complex) + return np.asarray(ar.to_numpy(array), dtype=complex) target_signature = _array_backend_signature(like) source_signature = _array_backend_signature(array) if source_signature == target_signature: @@ -3695,7 +3698,7 @@ def _gate_target_norm_from_expectation(self, p, gate, where): to_dense = getattr(gate, "to_dense", None) if not callable(to_dense): return None - gate = np.asarray(to_dense()) + gate = np.asarray(ar.to_numpy(to_dense())) shape = tuple(int(dim) for dim in gate.shape) if len(shape) != 2: dims = self._infer_gate_dims(gate, where) diff --git a/src/pepsy/optimizers/stabilizer_tn/mps_stab_optimizer.py b/src/pepsy/optimizers/stabilizer_tn/mps_stab_optimizer.py index 61ff136..85f0797 100644 --- a/src/pepsy/optimizers/stabilizer_tn/mps_stab_optimizer.py +++ b/src/pepsy/optimizers/stabilizer_tn/mps_stab_optimizer.py @@ -957,7 +957,7 @@ def _analysis_matrix_kind(cls, entry) -> str: if not (isinstance(entry, (list, tuple)) and len(entry) == 2): return "opaque" try: - gate = np.asarray(entry[0], dtype=complex) + gate = np.asarray(ar.to_numpy(entry[0]), dtype=complex) except (TypeError, ValueError): return "opaque" if gate.ndim != 2 or gate.shape[0] != gate.shape[1]: @@ -1264,7 +1264,7 @@ def _magic_strategy_entry_kind(cls, entry) -> str: if len(entry) != 2: return "opaque" try: - gate = np.asarray(entry[0], dtype=complex) + gate = np.asarray(ar.to_numpy(entry[0]), dtype=complex) dim = gate.shape[0] nq = int(round(math.log2(dim))) if ( @@ -4965,7 +4965,8 @@ def _dense_gate_target_norm(self, gate: np.ndarray, where) -> float: ) if norm_squared <= 0.0: return 0.0 - gram = np.asarray(gate).conj().T @ np.asarray(gate) + gate = np.asarray(ar.to_numpy(gate)) + gram = gate.conj().T @ gate expectation = 0.0 + 0.0j for term_index, (labels, coefficient) in enumerate( pauli_decomposition(gram, k, tol=self.operator_tol), start=1 @@ -5580,7 +5581,7 @@ def _looks_like_stream(gates) -> bool: def _is_unitary(gate: np.ndarray, tol: float = 1e-9) -> bool: """Return whether ``gate`` is unitary within ``tol``.""" - g = np.asarray(gate, dtype=complex) + g = np.asarray(ar.to_numpy(gate), dtype=complex) return np.allclose(g.conj().T @ g, np.eye(g.shape[0]), atol=tol) @@ -5590,7 +5591,7 @@ def _zyz_angles(gate: np.ndarray): Up to a global phase, using the convention ``Rz(a) = exp(-i a/2 Z)`` and ``Ry(t) = exp(-i t/2 Y)``. """ - u = np.asarray(gate, dtype=complex) + u = np.asarray(ar.to_numpy(gate), dtype=complex) det = u[0, 0] * u[1, 1] - u[0, 1] * u[1, 0] u = u / np.sqrt(det) # to SU(2) up to a sign (global phase, irrelevant) c = abs(u[0, 0]) diff --git a/src/pepsy/optimizers/stabilizer_tn/operators.py b/src/pepsy/optimizers/stabilizer_tn/operators.py index 50089e9..bcd2a78 100644 --- a/src/pepsy/optimizers/stabilizer_tn/operators.py +++ b/src/pepsy/optimizers/stabilizer_tn/operators.py @@ -13,6 +13,7 @@ from collections.abc import Mapping from typing import List, Sequence, Tuple +import autoray as ar import numpy as np import quimb.tensor as qtn @@ -52,7 +53,7 @@ def pauli_decomposition( Enumerates ``4**k`` Paulis, so intended for small ``k`` (few-qubit gates). """ - gate = np.asarray(gate) + gate = np.asarray(ar.to_numpy(gate)) dim = 2 ** k if gate.shape != (dim, dim): raise ValueError(f"gate must be {dim}x{dim} for k={k}, got {gate.shape}.") diff --git a/src/pepsy/optimizers/sweep/optimizer.py b/src/pepsy/optimizers/sweep/optimizer.py index f837955..e366f62 100644 --- a/src/pepsy/optimizers/sweep/optimizer.py +++ b/src/pepsy/optimizers/sweep/optimizer.py @@ -1068,12 +1068,7 @@ def _to_float_history(history): """Convert solver history entries into plain Python floats.""" values = [] for entry in history or (): - value = entry - if hasattr(value, "detach"): - value = value.detach() - if hasattr(value, "cpu"): - value = value.cpu() - values.append(float(value)) + values.append(float(ar.to_numpy(entry))) return values @staticmethod diff --git a/src/pepsy/optimizers/sym_dmrg.py b/src/pepsy/optimizers/sym_dmrg.py index 86d4fc5..427981c 100644 --- a/src/pepsy/optimizers/sym_dmrg.py +++ b/src/pepsy/optimizers/sym_dmrg.py @@ -16,6 +16,7 @@ import time import warnings +import autoray as ar import numpy as np from scipy.sparse.linalg import LinearOperator @@ -111,15 +112,7 @@ def _normalize_backend(backend): def _to_numpy(array): - if type(array) is np.ndarray: - return array - if hasattr(array, "detach") and hasattr(array, "cpu"): - array = array.detach().cpu() - if hasattr(array, "numpy"): - return np.asarray(array.numpy()) - if hasattr(array, "get"): - return np.asarray(array.get()) - return np.asarray(array) + return np.asarray(ar.to_numpy(array)) def _dense_data(data): diff --git a/src/pepsy/optimizers/tree/operators.py b/src/pepsy/optimizers/tree/operators.py index 126fb68..be83647 100644 --- a/src/pepsy/optimizers/tree/operators.py +++ b/src/pepsy/optimizers/tree/operators.py @@ -39,13 +39,9 @@ def _as_numpy(data, *, dtype=None): """Convert a dense backend array to host NumPy construction data.""" - if hasattr(data, "detach"): - data = data.detach() - if hasattr(data, "cpu"): - data = data.cpu() - if hasattr(data, "get"): - data = data.get() - return np.asarray(data, dtype=dtype) + if hasattr(data, "to_dense"): + data = data.to_dense() + return np.asarray(ar.to_numpy(data), dtype=dtype) def _tree_plan_signature(plan): diff --git a/src/pepsy/optimizers/tree_stabilizer/optimizer.py b/src/pepsy/optimizers/tree_stabilizer/optimizer.py index 8e254a3..fed0445 100644 --- a/src/pepsy/optimizers/tree_stabilizer/optimizer.py +++ b/src/pepsy/optimizers/tree_stabilizer/optimizer.py @@ -144,7 +144,7 @@ def _looks_like_single_entry(gates): def _is_unitary(gate): """Return whether a dense gate is unitary to the STN tolerance.""" - gate = np.asarray(gate) + gate = np.asarray(ar.to_numpy(gate)) if gate.ndim != 2 or gate.shape[0] != gate.shape[1]: return False return np.allclose( @@ -159,8 +159,8 @@ def _apply_dense_gate(state, gate, where, n): """Apply a small gate to a dense state in logical big-endian order.""" where = tuple(int(q) for q in where) k = len(where) - tensor = np.asarray(state).reshape((2,) * n) - operator = np.asarray(gate).reshape((2,) * (2 * k)) + tensor = np.asarray(ar.to_numpy(state)).reshape((2,) * n) + operator = np.asarray(ar.to_numpy(gate)).reshape((2,) * (2 * k)) out = np.tensordot( operator, tensor, @@ -320,7 +320,7 @@ def _dense_to_tree_state(state, plan, *, max_bond=None, cutoff=0.0, dtype=comple if cutoff < 0.0: raise ValueError("cutoff must be non-negative.") - dense = np.asarray(state, dtype=dtype).reshape(-1) + dense = np.asarray(ar.to_numpy(state), dtype=dtype).reshape(-1) expected_size = 2 ** plan.n if dense.size != expected_size: raise ValueError( @@ -1279,7 +1279,7 @@ def _frame_layout_trace_entry(self, entry, records, *, weight_mode): if not isinstance(head, str): if len(entry) != 2: raise ValueError(f"Unsupported gate stream entry: {entry!r}.") - gate = np.asarray(entry[0]) + gate = np.asarray(ar.to_numpy(entry[0])) where = _normalize_sites(entry[1]) if gate.ndim != 2 or gate.shape[0] != gate.shape[1]: raise ValueError(f"Gate matrix must be square, got {gate.shape}.") diff --git a/src/pepsy/sampling/samplers.py b/src/pepsy/sampling/samplers.py index bd5b2f7..8532bff 100644 --- a/src/pepsy/sampling/samplers.py +++ b/src/pepsy/sampling/samplers.py @@ -6,6 +6,7 @@ import math from typing import Any, Iterable +import autoray as ar import numpy as np from tqdm import tqdm @@ -122,12 +123,9 @@ def _mps_array_backend(array): def _backend_array_to_numpy(array): - backend = _mps_array_backend(array) - if backend == "torch": - return array.detach().cpu().numpy() - if backend == "cupy": - return array.get() - return np.asarray(array) + if hasattr(array, "to_dense"): + array = array.to_dense() + return np.asarray(ar.to_numpy(array)) def _fermion_symmray_occupations(charge, offset, fermion): @@ -274,16 +272,7 @@ def _to_dense_numpy(array): """Convert a dense or Symmray array to a NumPy array for BP sampling.""" if hasattr(array, "to_dense"): array = array.to_dense() - detach = getattr(array, "detach", None) - if callable(detach): - array = detach() - cpu = getattr(array, "cpu", None) - if callable(cpu): - array = cpu() - numpy = getattr(array, "numpy", None) - if callable(numpy): - array = numpy() - return np.asarray(array) + return np.asarray(ar.to_numpy(array)) def _prepare_bp_binary_network(tn, *, site_order=None, encoding=None): @@ -727,11 +716,11 @@ def magnetizations(self, *, to_numpy: bool = False): if backend == "torch": configs = self.configs.to(dtype=self.probs.dtype) out = (1 - 2 * configs).sum(dim=1) / float(self.L) - return out.detach().cpu().numpy() if to_numpy else out + return ar.to_numpy(out) if to_numpy else out if backend == "cupy": configs = self.configs.astype(np.float64, copy=False) out = (1 - 2 * configs).sum(axis=1) / float(self.L) - return out.get() if to_numpy else out + return ar.to_numpy(out) if to_numpy else out configs = np.asarray(self.configs, dtype=float) return (1 - 2 * configs).sum(axis=1) / float(self.L) @@ -962,7 +951,7 @@ def refresh(self, psi=None): # Convert to numpy for quimb sampling compatibility self._psi = psi.copy() self._psi.apply_to_arrays( - lambda x: x.get() if hasattr(x, "get") else np.asarray(x) + lambda x: ar.to_numpy(x) ) return self @@ -1326,7 +1315,7 @@ def _symmray_draw_many(probs, n_draws, state, rng): replacement=True, generator=rng, ) - return choices.detach().cpu().numpy().astype(np.int64, copy=False) + return np.asarray(ar.to_numpy(choices), dtype=np.int64) if backend == "cupy": import cupy as cp # pylint: disable=import-outside-toplevel @@ -1334,7 +1323,7 @@ def _symmray_draw_many(probs, n_draws, state, rng): draws = rng.random(n_draws) choices = cp.searchsorted(cdf, draws, side="right") choices = cp.minimum(choices, int(probs.shape[0]) - 1) - return choices.get().astype(np.int64, copy=False) + return np.asarray(ar.to_numpy(choices), dtype=np.int64) return np.asarray( rng.choice(len(probs), size=n_draws, p=probs), dtype=np.int64, @@ -2337,8 +2326,8 @@ def _torch_sample(site_data, n_samples, seed, *, to_numpy): configs = torch.stack(configs, dim=1) if to_numpy: - configs = configs.detach().cpu().numpy() - probs_total = probs_total.detach().cpu().numpy() + configs = np.asarray(ar.to_numpy(configs)) + probs_total = np.asarray(ar.to_numpy(probs_total)) return configs, probs_total @staticmethod @@ -2380,12 +2369,9 @@ def _array_namespace_sample(site_data, n_samples, seed, *, backend, to_numpy): configs.append(choices) configs = xp.stack(configs, axis=1) - if to_numpy and backend == "cupy": - configs = configs.get() - probs_total = probs_total.get() if to_numpy: - configs = np.asarray(configs) - probs_total = np.asarray(probs_total) + configs = np.asarray(ar.to_numpy(configs)) + probs_total = np.asarray(ar.to_numpy(probs_total)) return configs, probs_total @staticmethod @@ -2850,10 +2836,7 @@ def __init__( # Extract the state vector with correct index ordering vec = self._to_vector(state, L, ind_id) - # Move to CPU if needed (cupy) - if hasattr(vec, "get"): - vec = vec.get() - vec = np.asarray(vec, dtype=complex).ravel() + vec = np.asarray(ar.to_numpy(vec), dtype=complex).ravel() expected_size = 2 ** L if vec.size != expected_size: raise ValueError( diff --git a/src/pepsy/sampling/tree.py b/src/pepsy/sampling/tree.py index 36da03e..9cfff45 100644 --- a/src/pepsy/sampling/tree.py +++ b/src/pepsy/sampling/tree.py @@ -46,6 +46,8 @@ import numpy as np +import autoray as ar + from .samplers import FermionConfigurationEncoding __all__ = [ @@ -316,10 +318,10 @@ def _extract_arrays(self, tn): fermionic = bool(getattr(tn, "fermionic", False)) if fermionic: def to_arr(data): - return np.asarray(data.to_dense()) + return np.asarray(ar.to_numpy(data.to_dense())) else: def to_arr(data): - return np.asarray(data) + return np.asarray(ar.to_numpy(data)) def bond_between(a, b): shared = set(tn.node_tensor(a).inds) & set(tn.node_tensor(b).inds) diff --git a/src/pepsy/tensors/symmetric.py b/src/pepsy/tensors/symmetric.py index 94e3a85..d86e344 100644 --- a/src/pepsy/tensors/symmetric.py +++ b/src/pepsy/tensors/symmetric.py @@ -85,13 +85,7 @@ def _as_python_bool(value): def _to_host_numpy(value): value = _to_dense(value) - detach = getattr(value, "detach", None) - if callable(detach): - value = detach() - cpu = getattr(value, "cpu", None) - if callable(cpu): - value = cpu() - return np.asarray(value) + return np.asarray(ar.to_numpy(value)) def _is_symmray_array(value): @@ -4831,15 +4825,15 @@ def _as_scalar(value): shape = tuple(shape) if shape != (): return value - detach = getattr(value, "detach", None) - if callable(detach): - value = detach() - cpu = getattr(value, "cpu", None) - if callable(cpu): - value = cpu() - item = getattr(value, "item", None) - if callable(item): - return item() + try: + return ar.to_numpy(value).item() + except Exception: + # Preserve the small duck-typed scalar contract for backend + # wrappers that expose ``item`` but are not Autoray-registered. + item = getattr(value, "item", None) + if callable(item): + return item() + raise arr = np.asarray(value) if arr.shape == (): return arr.item() @@ -5165,16 +5159,7 @@ def _coupling_is_active(value): def _dense_numpy(value, *, dtype=None): value = _to_dense(value) - detach = getattr(value, "detach", None) - if callable(detach): - value = detach() - cpu = getattr(value, "cpu", None) - if callable(cpu): - value = cpu() - # CuPy arrays reject implicit host conversion; move to host explicitly. - if type(value).__module__.split(".", 1)[0] == "cupy": - value = value.get() - return np.asarray(value, dtype=dtype) + return np.asarray(ar.to_numpy(value), dtype=dtype) def _is_single_site_identity_hamiltonian(target, local_dim, zero_charge): diff --git a/src/pepsy/vmc/__init__.py b/src/pepsy/vmc/__init__.py index f0a1a45..a6ca7d0 100644 --- a/src/pepsy/vmc/__init__.py +++ b/src/pepsy/vmc/__init__.py @@ -41,6 +41,8 @@ "PackedFermionicPEPS": ".netket", "SpinOrbitalColumns": ".netket", "TorchConnections": ".torch", + "TorchFockTransitionPlan": ".torch", + "TorchAmplitudeCache": ".torch", "TorchAmplitudeBenchmark": ".torch", "TorchAmplitudeBenchmarkRun": ".torch", "TorchFermionVMC": ".torch", diff --git a/src/pepsy/vmc/netket.py b/src/pepsy/vmc/netket.py index a96efdd..fbdb231 100644 --- a/src/pepsy/vmc/netket.py +++ b/src/pepsy/vmc/netket.py @@ -14,6 +14,7 @@ from typing import Any import warnings +import autoray as ar import numpy as np import quimb.tensor as qtn @@ -2296,11 +2297,7 @@ def _uses_flat_symmray_arrays(tn): def _host_array_for_flatten(value): """Copy a Torch/CUDA block to a host NumPy array for Symmray packing.""" - if hasattr(value, "detach"): - value = value.detach() - if hasattr(value, "cpu"): - value = value.cpu() - return np.asarray(value) + return np.asarray(ar.to_numpy(value)) def _z2_flat_padded(data, sr): @@ -4619,13 +4616,7 @@ def _native_term_support(where, *, coordinate_sites): def _native_term_to_numpy(term): """Transfer one small native local term to a host dense matrix.""" dense = term.to_dense() if hasattr(term, "to_dense") else term - detach = getattr(dense, "detach", None) - if callable(detach): - dense = detach() - cpu = getattr(dense, "cpu", None) - if callable(cpu): - dense = cpu() - return np.asarray(dense, dtype=np.complex128) + return np.asarray(ar.to_numpy(dense), dtype=np.complex128) def _project_native_term(matrix, candidates, *, where): diff --git a/src/pepsy/vmc/torch/__init__.py b/src/pepsy/vmc/torch/__init__.py index 09b65e4..1b21729 100644 --- a/src/pepsy/vmc/torch/__init__.py +++ b/src/pepsy/vmc/torch/__init__.py @@ -20,7 +20,13 @@ TorchAmplitudeBenchmarkRun, benchmark_torch_amplitudes, ) -from .connections import TorchConnections, compile_operator_sum_torch, torch_hamiltonian_connections +from .cache import TorchAmplitudeCache +from .connections import ( + TorchConnections, + TorchFockTransitionPlan, + compile_operator_sum_torch, + torch_hamiltonian_connections, +) from .driver import TorchVMCDriver from .fermion import ( TorchFermionVMC, @@ -72,6 +78,8 @@ "TorchAmplitudeBenchmark", "TorchAmplitudeBenchmarkRun", "TorchConnections", + "TorchFockTransitionPlan", + "TorchAmplitudeCache", "TorchMetropolisResult", "TorchImportanceSamples", "TorchMCMCSamples", diff --git a/src/pepsy/vmc/torch/_core.py b/src/pepsy/vmc/torch/_core.py index 64e403b..88d443a 100644 --- a/src/pepsy/vmc/torch/_core.py +++ b/src/pepsy/vmc/torch/_core.py @@ -48,6 +48,7 @@ ) from .connections import ( TorchConnections, + TorchFockTransitionPlan, compile_operator_sum_torch, torch_hamiltonian_connections, ) @@ -122,6 +123,7 @@ "TorchPEPSAmplitude", "TorchPEPSBoundaryAmplitude", "TorchConnections", + "TorchFockTransitionPlan", "TorchMetropolisResult", "TorchImportanceSamples", "TorchMCMCSamples", diff --git a/src/pepsy/vmc/torch/_graded.py b/src/pepsy/vmc/torch/_graded.py index dbc0aad..458643b 100644 --- a/src/pepsy/vmc/torch/_graded.py +++ b/src/pepsy/vmc/torch/_graded.py @@ -7,6 +7,7 @@ from dataclasses import dataclass +import autoray as ar import numpy as np from ..torch_types import _require_torch @@ -55,7 +56,7 @@ def _graded_torch_embed_dense(array, labels, full_maps): """Embed a sparse result into the fixed dense index layout.""" shape = tuple(len(full_maps[label]) for label in labels) if getattr(array, "num_blocks", 0): - dense = np.asarray(array.to_dense()) + dense = np.asarray(ar.to_numpy(array.to_dense())) else: dense = np.zeros(shape, dtype=float) diff --git a/src/pepsy/vmc/torch/amplitude.py b/src/pepsy/vmc/torch/amplitude.py index 88a7257..058eb9c 100644 --- a/src/pepsy/vmc/torch/amplitude.py +++ b/src/pepsy/vmc/torch/amplitude.py @@ -396,14 +396,13 @@ def __init__( and self.symmray_tensor_ids and self.contraction_opts.get("mode") is None ): - # Quimb's default CTMRG projector compressor forms arbitrary- - # geometry oblique projectors. With Symmray blocks backed by - # Torch, that intermediate product can have incompatible dense - # dimensions even though the original block-sparse contraction is - # valid. The direct SVD boundary compressor keeps the contraction - # sector-local and is the compatible default for this case. A - # caller-provided mode remains an explicit override. - self.contraction_opts["mode"] = "direct" + # The scoped native-fermionic CTMRG compatibility layer repairs + # projector insertion and zero sectors. Keep the projector route + # as the default for Torch-backed Symmray: the direct boundary + # compressor can mix NumPy intermediates with Symmray blocks on + # larger native PEPS. A caller-provided mode remains an explicit + # override. + self.contraction_opts["mode"] = "projector" # ``torch.vmap`` can batch the pure tensor contractions for dense and # compatible Symmray PEPS. Keep a per-model fallback for contraction # paths or optional backends that cannot be vmapped. diff --git a/src/pepsy/vmc/torch/cache.py b/src/pepsy/vmc/torch/cache.py new file mode 100644 index 0000000..1cb171f --- /dev/null +++ b/src/pepsy/vmc/torch/cache.py @@ -0,0 +1,229 @@ +"""Reusable no-gradient amplitude caching for Torch VMC measurements.""" + +from __future__ import annotations + +from collections import OrderedDict + +from .amplitude import _call_amplitude_fn, _unique_config_rows +from .connections import TorchConnections +from .results import _torch_sample_provenance +from ._common import _as_long_matrix +from ..torch_types import _require_torch + +__all__ = ["TorchAmplitudeCache"] + + +class TorchAmplitudeCache: + """Bounded configuration-to-amplitude cache with model-state invalidation. + + The cache is deliberately detached and no-gradient only. Its provenance + includes the PEPS object identity, parameter versions, contraction + settings, dtype, and device, so reusing a cache after an optimization or + model move cannot silently return stale amplitudes. + """ + + def __init__(self, max_entries=100_000): + if isinstance(max_entries, bool) or int(max_entries) < 1: + raise ValueError("max_entries must be a positive integer.") + self.max_entries = int(max_entries) + self._values = OrderedDict() + self._signature = None + self._requests = 0 + self._hits = 0 + self._misses = 0 + + @staticmethod + def _signature_for(model): + parameters = getattr(model, "parameters", None) + if callable(parameters): + parameter_signature = tuple( + (str(parameter.dtype), str(parameter.device)) + for parameter in parameters() + ) + else: + parameter_signature = () + return _torch_sample_provenance(model), parameter_signature + + @staticmethod + def _key(config): + return tuple(int(value) for value in config.detach().cpu().tolist()) + + def _sync(self, model): + signature = self._signature_for(model) + if self._signature != signature: + self.clear() + self._signature = signature + + def clear(self): + """Drop cached values and counters while retaining the size limit.""" + self._values.clear() + self._requests = 0 + self._hits = 0 + self._misses = 0 + + def _put(self, key, value): + self._values[key] = value.detach().clone() + self._values.move_to_end(key) + while len(self._values) > self.max_entries: + self._values.popitem(last=False) + + def seed(self, configs, amplitudes, *, model=None): + """Insert already-contracted parent amplitudes into the cache.""" + torch = _require_torch() + configs = _as_long_matrix(configs) + amplitudes = torch.as_tensor(amplitudes, device=configs.device).reshape(-1) + if configs.shape[0] != amplitudes.shape[0]: + raise ValueError("configs and amplitudes must have the same length.") + if model is not None: + self._sync(model) + for config, amplitude in zip(configs, amplitudes): + self._put(self._key(config), amplitude) + + def evaluate(self, model, configs, *, chunk_size=None): + """Evaluate amplitudes, contracting only configurations not cached.""" + torch = _require_torch() + if torch.is_grad_enabled(): + raise RuntimeError("TorchAmplitudeCache is only valid in no-grad mode.") + configs = _as_long_matrix(configs) + self._sync(model) + if configs.shape[0] == 0: + return torch.empty(0, dtype=torch.get_default_dtype(), device=configs.device) + unique_configs, inverse = _unique_config_rows(configs) + if inverse is None: + inverse = torch.zeros(1, dtype=torch.long, device=configs.device) + values = [None] * int(unique_configs.shape[0]) + missing = [] + for index, config in enumerate(unique_configs): + self._requests += 1 + key = self._key(config) + value = self._values.get(key) + if value is None: + missing.append(index) + self._misses += 1 + else: + values[index] = value.to(device=configs.device) + self._hits += 1 + self._values.move_to_end(key) + if missing: + missing_index = torch.as_tensor( + missing, dtype=torch.long, device=configs.device, + ) + computed = _call_amplitude_fn( + model, unique_configs[missing_index], chunk_size=chunk_size, + ) + for offset, index in enumerate(missing): + value = computed[offset].detach() + values[index] = value + self._put(self._key(unique_configs[index]), value) + return torch.stack(values)[inverse] + + def wrap(self, model): + """Return a callable model proxy using this cache.""" + return _CachedAmplitudeModel(model, self) + + def snapshot(self): + """Return lightweight counters suitable for a profile/result JSON.""" + return { + "entries": int(len(self._values)), + "max_entries": int(self.max_entries), + "requests": int(self._requests), + "hits": int(self._hits), + "misses": int(self._misses), + "hit_rate": ( + float(self._hits) / float(self._requests) + if self._requests else 0.0 + ), + } + + +class _CachedAmplitudeModel: + """Delegate PEPS-specific connected work while caching target rows.""" + + def __init__(self, model, cache): + self._model = model + self._cache = cache + + def __getattr__(self, name): + return getattr(self._model, name) + + def __call__(self, configs, *args, **kwargs): + chunk_size = kwargs.pop("chunk_size", None) + if args or kwargs: + return self._model(configs, *args, **kwargs) + return self._cache.evaluate(self._model, configs, chunk_size=chunk_size) + + def connected_amplitudes( + self, + configs, + amplitudes, + connections, + *, + chunk_size=None, + reuse_diagonal=True, + ): + torch = _require_torch() + configs = _as_long_matrix(configs) + amplitudes = torch.as_tensor(amplitudes, device=configs.device).reshape(-1) + self._cache._sync(self._model) + self._cache.seed(configs, amplitudes) + target_configs = _as_long_matrix(connections.configs) + if target_configs.shape[0] == 0: + return torch.empty(0, dtype=amplitudes.dtype, device=configs.device) + + unique_targets, inverse = _unique_config_rows(target_configs) + if inverse is None: + inverse = torch.zeros(1, dtype=torch.long, device=configs.device) + values = [None] * int(unique_targets.shape[0]) + missing = [] + representative = {} + for index, config in enumerate(target_configs): + representative.setdefault(self._cache._key(config), index) + for index, config in enumerate(unique_targets): + self._cache._requests += 1 + key = self._cache._key(config) + value = self._cache._values.get(key) + if value is None: + missing.append(index) + self._cache._misses += 1 + else: + values[index] = value.to( + device=configs.device, dtype=amplitudes.dtype, + ) + self._cache._hits += 1 + self._cache._values.move_to_end(key) + + if missing: + missing_index = torch.as_tensor( + missing, dtype=torch.long, device=configs.device, + ) + representative_ids = torch.as_tensor( + [representative[self._cache._key(config)] for config in unique_targets[missing_index]], + dtype=torch.long, + device=configs.device, + ) + subset = TorchConnections( + configs=unique_targets[missing_index], + coeffs=torch.ones( + len(missing), dtype=amplitudes.dtype, device=configs.device, + ), + batch_ids=connections.batch_ids[representative_ids], + ) + connected = getattr(self._model, "connected_amplitudes", None) + if callable(connected): + computed = connected( + configs, + amplitudes, + subset, + chunk_size=chunk_size, + reuse_diagonal=reuse_diagonal, + ) + else: + computed = _call_amplitude_fn( + self._model, subset.configs, chunk_size=chunk_size, + ) + computed = torch.as_tensor(computed, device=configs.device).reshape(-1) + for offset, index in enumerate(missing): + value = computed[offset].detach() + values[index] = value + self._cache._put(self._cache._key(unique_targets[index]), value) + return torch.stack(values)[inverse].to(dtype=amplitudes.dtype) diff --git a/src/pepsy/vmc/torch/connections.py b/src/pepsy/vmc/torch/connections.py index 948e7c2..5db430a 100644 --- a/src/pepsy/vmc/torch/connections.py +++ b/src/pepsy/vmc/torch/connections.py @@ -8,10 +8,12 @@ from __future__ import annotations +from collections.abc import Mapping from dataclasses import dataclass from numbers import Integral from typing import Any +import autoray as ar import numpy as np from ..torch_types import _require_torch @@ -30,6 +32,98 @@ class TorchConnections: coeffs: Any batch_ids: Any + def to(self, device): + """Return this connection table on ``device``.""" + return TorchConnections( + configs=self.configs.to(device=device), + coeffs=self.coeffs.to(device=device), + batch_ids=self.batch_ids.to(device=device), + ) + + def slice(self, start, stop): + """Slice rows belonging to parent samples ``[start, stop)``.""" + start = int(start) + stop = int(stop) + if start < 0 or stop < start: + raise ValueError("connection slice must satisfy 0 <= start <= stop") + mask = (self.batch_ids >= start) & (self.batch_ids < stop) + return TorchConnections( + configs=self.configs[mask], + coeffs=self.coeffs[mask], + batch_ids=self.batch_ids[mask] - start, + ) + + +@dataclass(frozen=True) +class TorchFockTransitionPlan: + """Reusable parent-to-connected-configuration tables. + + A plan is tied to one ordered parent configuration stream and one named + observable map. It contains no PEPS parameters, so the same plan can be + reused for every PEPS bond dimension or dtype that measures the same + stored proposal stream. :meth:`slice` makes it safe to process a persisted + full-stream plan in sequential chunks. + """ + + configs: Any + connection_map: Mapping[str, TorchConnections] + + def __post_init__(self): + configs = _as_long_matrix(self.configs) + if not isinstance(self.connection_map, Mapping): + raise TypeError("connection_map must be a mapping of observable names.") + normalized = {} + for name, connections in self.connection_map.items(): + if not isinstance(name, str): + raise TypeError("transition-plan observable names must be strings.") + if not isinstance(connections, TorchConnections): + raise TypeError("connection_map values must be TorchConnections.") + normalized[name] = connections + object.__setattr__(self, "configs", configs) + object.__setattr__(self, "connection_map", normalized) + + @property + def observable_names(self): + return tuple(self.connection_map) + + @property + def n_samples(self): + return int(self.configs.shape[0]) + + @property + def n_connections(self): + return sum( + int(connections.configs.shape[0]) + for connections in self.connection_map.values() + ) + + def to(self, device): + """Return the plan on ``device`` without changing its parent order.""" + return TorchFockTransitionPlan( + configs=self.configs.to(device=device), + connection_map={ + name: connections.to(device) + for name, connections in self.connection_map.items() + }, + ) + + def slice(self, start, stop): + """Return the plan for parent samples ``[start, stop)``.""" + start = int(start) + stop = int(stop) + if start < 0 or stop < start or stop > self.n_samples: + raise ValueError( + f"invalid transition-plan slice [{start}, {stop}) for " + f"{self.n_samples} samples" + ) + return TorchFockTransitionPlan( + configs=self.configs[start:stop], + connection_map={ + name: connections.slice(start, stop) + for name, connections in self.connection_map.items() + }, + ) + def _term_items(terms): """Return ``(where, operator)`` pairs from common Hamiltonian containers.""" @@ -111,13 +205,7 @@ def _charge_parity(charge): def _operator_dense_numpy(operator): """Get a detached CPU view of a fixed native operator tensor.""" dense = _term_dense_array(operator) - detach = getattr(dense, "detach", None) - if callable(detach): - dense = detach() - cpu = getattr(dense, "cpu", None) - if callable(cpu): - dense = cpu() - return np.asarray(dense) + return np.asarray(ar.to_numpy(dense)) @dataclass(frozen=True) @@ -719,6 +807,7 @@ def map_site(site): __all__ = [ "TorchConnections", + "TorchFockTransitionPlan", "compile_operator_sum_torch", "torch_hamiltonian_connections", "_driver_terms_connections", diff --git a/src/pepsy/vmc/torch/driver.py b/src/pepsy/vmc/torch/driver.py index 13d9134..e2aa137 100644 --- a/src/pepsy/vmc/torch/driver.py +++ b/src/pepsy/vmc/torch/driver.py @@ -17,7 +17,11 @@ _unique_config_rows, ) from .benchmark import benchmark_torch_amplitudes -from .connections import _driver_terms_connections, compile_operator_sum_torch +from .connections import ( + TorchFockTransitionPlan, + _driver_terms_connections, + compile_operator_sum_torch, +) from .distributed import ( distributed_max_float, distributed_metadata, @@ -647,6 +651,59 @@ def make_connections(self, configs=None, *, terms=None): ) return self.connection_fn(configs, self.graph, **self.connection_kwargs) + def compile_fock_plan(self, configs=None, *, observables=None): + """Compile reusable configuration transitions for a fixed sample stream. + + ``observables`` follows :meth:`measure_samples`: ``None`` means the + driver's configured connection function, while a mapping names native + explicit term sets and may use ``None`` for that configured function. + The returned plan contains no model parameters and can therefore be + reused across PEPS bond dimensions, contraction methods, and dtypes. + """ + configs = self.configs if configs is None else _as_long_matrix(configs) + configs = configs.to(device=_model_device(self.model)) + if observables is None: + observable_items = (("observable", None),) + else: + try: + observable_items = tuple(observables.items()) + except AttributeError as exc: + raise TypeError( + "observables must be a mapping of names to native terms." + ) from exc + if not observable_items: + raise ValueError("observables must contain at least one entry.") + connection_map = { + name: ( + self.make_connections(configs, terms=terms) + if terms is not None + else self.make_connections(configs) + ) + for name, terms in observable_items + } + return TorchFockTransitionPlan(configs, connection_map) + + def proposal_configs( + self, + batch, + *, + fermion=None, + one_d_to_two_d=None, + site_order=None, + occupation_map=None, + ): + """Bridge an already sampled external proposal without contracting it.""" + from .importance import _proposal_batch_configs + + return _proposal_batch_configs( + self, + batch, + fermion=fermion, + site_order=site_order, + occupation_map=occupation_map, + device=_model_device(self.model), + ) + def sample_sweep(self, *, n_sweeps=1, track_proposal_stats=False): """Run one or more Metropolis sweeps and update driver state.""" n_sweeps = _check_positive_int("n_sweeps", n_sweeps) @@ -1209,6 +1266,8 @@ def measure_samples( deduplicate=True, progress=False, distributed=None, + connection_plan=None, + amplitude_cache=None, ): """Measure saved chain samples without running another sampler. @@ -1319,6 +1378,14 @@ def measure_samples( else int(flat_configs.shape[0]) ) + if amplitude_cache is not None: + from .cache import TorchAmplitudeCache + + if not isinstance(amplitude_cache, TorchAmplitudeCache): + raise TypeError( + "amplitude_cache must be a TorchAmplitudeCache or None." + ) + if amplitudes is None and sample_object is not None: amplitudes = getattr(sample_object, "amplitudes", None) if refresh_proposal_amplitudes: @@ -1331,18 +1398,32 @@ def measure_samples( with torch.no_grad(): if deduplicate and unique_parent_count < flat_configs.shape[0]: unique_configs, inverse = _unique_config_rows(flat_configs) - unique_amplitudes = _call_amplitude_fn( - self.model, - unique_configs, - chunk_size=self.chunk_size, - ) + if amplitude_cache is None: + unique_amplitudes = _call_amplitude_fn( + self.model, + unique_configs, + chunk_size=self.chunk_size, + ) + else: + unique_amplitudes = amplitude_cache.evaluate( + self.model, + unique_configs, + chunk_size=self.chunk_size, + ) flat_amplitudes = unique_amplitudes[inverse] else: - flat_amplitudes = _call_amplitude_fn( - self.model, - flat_configs, - chunk_size=self.chunk_size, - ) + if amplitude_cache is None: + flat_amplitudes = _call_amplitude_fn( + self.model, + flat_configs, + chunk_size=self.chunk_size, + ) + else: + flat_amplitudes = amplitude_cache.evaluate( + self.model, + flat_configs, + chunk_size=self.chunk_size, + ) chain_amplitudes = flat_amplitudes.reshape(n_steps, n_chains) else: amplitudes = torch.as_tensor(amplitudes, device=model_device) @@ -1436,14 +1517,40 @@ def set_phase(stage, *, n_connections=None): set_phase("connections") connection_start = time.perf_counter() - connection_map = { - name: ( - self.make_connections(flat_configs, terms=terms) - if terms is not None - else self.make_connections(flat_configs) - ) - for name, terms in observable_items - } + if connection_plan is None: + connection_map = { + name: ( + self.make_connections(flat_configs, terms=terms) + if terms is not None + else self.make_connections(flat_configs) + ) + for name, terms in observable_items + } + else: + if not isinstance(connection_plan, TorchFockTransitionPlan): + raise TypeError( + "connection_plan must be a TorchFockTransitionPlan." + ) + if connection_plan.observable_names != tuple( + name for name, _ in observable_items + ): + raise ValueError( + "connection_plan observable names do not match observables." + ) + if ( + connection_plan.configs.shape != flat_configs.shape + or not torch.equal( + connection_plan.configs.to(device=flat_configs.device), + flat_configs, + ) + ): + raise ValueError( + "connection_plan parent configurations do not match samples." + ) + connection_map = { + name: connections.to(flat_configs.device) + for name, connections in connection_plan.connection_map.items() + } n_connections = sum( int(connections.configs.shape[0]) for connections in connection_map.values() @@ -1454,12 +1561,20 @@ def set_phase(stage, *, n_connections=None): set_phase("target amplitudes", n_connections=n_connections) local_start = time.perf_counter() + measurement_model = self.model + if amplitude_cache is not None: + amplitude_cache.seed( + flat_configs, + flat_amplitudes, + model=self.model, + ) + measurement_model = amplitude_cache.wrap(self.model) with torch.no_grad(): flat_values = _local_energies_from_connection_map( flat_configs, flat_amplitudes, connection_map, - self.model, + measurement_model, chunk_size=self.chunk_size, reuse_diagonal=True, deduplicate_targets=deduplicate, @@ -1504,6 +1619,11 @@ def set_phase(stage, *, n_connections=None): ), "total_seconds": elapsed, "cache": _cache_profile_snapshot(self.model), + **( + {"amplitude_cache": amplitude_cache.snapshot()} + if amplitude_cache is not None + else {} + ), "samples_only": True, "deduplicate": bool(deduplicate), "num_samples": int(flat_configs.shape[0]), @@ -2421,6 +2541,7 @@ def sample_from_proposal( sample_kwargs=None, progress=False, amplitude_floor=0.0, + amplitude_cache=None, ): """Draw reusable PEPS-code samples from an external proposal. @@ -2442,6 +2563,7 @@ def sample_from_proposal( sample_kwargs=sample_kwargs, progress=progress, amplitude_floor=amplitude_floor, + amplitude_cache=amplitude_cache, ) def measure_from_proposal( @@ -2460,6 +2582,7 @@ def measure_from_proposal( amplitude_floor=0.0, profile=False, deduplicate=True, + amplitude_cache=None, ): """Measure from an external MPS, BP, tree, or proposal batch. @@ -2480,6 +2603,7 @@ def measure_from_proposal( sample_kwargs=sample_kwargs, progress=progress, amplitude_floor=amplitude_floor, + amplitude_cache=amplitude_cache, ) return self.measure_samples( samples, @@ -2487,6 +2611,7 @@ def measure_from_proposal( profile=profile, deduplicate=deduplicate, progress=progress, + amplitude_cache=amplitude_cache, ) def importance_energy_estimate( diff --git a/src/pepsy/vmc/torch/fermion.py b/src/pepsy/vmc/torch/fermion.py index c24d677..2d122fa 100644 --- a/src/pepsy/vmc/torch/fermion.py +++ b/src/pepsy/vmc/torch/fermion.py @@ -8,6 +8,8 @@ import time from typing import Any +import autoray as ar + from ..torch_types import FermionSiteEncoding, _check_positive_int, _require_torch from ._common import ( _as_long_matrix, @@ -811,6 +813,7 @@ def sample( sample_kwargs = kwargs.pop("sample_kwargs", None) progress = kwargs.pop("progress", False) amplitude_floor = kwargs.pop("amplitude_floor", 0.0) + amplitude_cache = kwargs.pop("amplitude_cache", None) if kwargs: unexpected = ", ".join(sorted(kwargs)) raise TypeError( @@ -831,6 +834,7 @@ def sample( sample_kwargs=sample_kwargs, progress=progress, amplitude_floor=amplitude_floor, + amplitude_cache=amplitude_cache, ) distributed_runtime, initialization_sampling = self._rank_sharded_sampling_config( sampling, @@ -935,6 +939,8 @@ def measure( deduplicate=True, progress=False, distributed=None, + connection_plan=None, + amplitude_cache=None, _include_energy=False, ): """Measure observables from retained samples without resampling. @@ -958,6 +964,10 @@ def measure( "deduplicate": deduplicate, "progress": progress, } + if connection_plan is not None: + measure_kwargs["connection_plan"] = connection_plan + if amplitude_cache is not None: + measure_kwargs["amplitude_cache"] = amplitude_cache if distributed is not None: measure_kwargs["distributed"] = distributed return self.measure_samples( @@ -1222,13 +1232,7 @@ def _compile_observables(self, observables): def _vmc_result_scalar(value): """Convert a scalar Torch/JAX-like result to a real Python float.""" - detach = getattr(value, "detach", None) - if callable(detach): - value = detach() - cpu = getattr(value, "cpu", None) - if callable(cpu): - value = cpu() - array = np.asarray(value) + array = np.asarray(ar.to_numpy(value)) if array.size != 1: raise ValueError("Expected a scalar VMC result.") return float(np.real(array.reshape(-1)[0])) @@ -1314,6 +1318,8 @@ def measure( samples=None, weights=None, proposal_log_probs=None, + connection_plan=None, + amplitude_cache=None, ): """Measure energy and optional observables from one shared sample set. @@ -1345,6 +1351,8 @@ def measure( observables={"energy": None, **extra_terms}, weights=weights, proposal_log_probs=proposal_log_probs, + connection_plan=connection_plan, + amplitude_cache=amplitude_cache, ) energy = estimates["energy"] else: @@ -1352,6 +1360,8 @@ def measure( native_samples, weights=weights, proposal_log_probs=proposal_log_probs, + connection_plan=connection_plan, + amplitude_cache=amplitude_cache, ) estimates = {"energy": energy} return VMCMeasurement( diff --git a/src/pepsy/vmc/torch/importance.py b/src/pepsy/vmc/torch/importance.py index f4e939e..1a230d3 100644 --- a/src/pepsy/vmc/torch/importance.py +++ b/src/pepsy/vmc/torch/importance.py @@ -358,6 +358,7 @@ def _bridge_samples( progress, sample_kwargs, amplitude_floor, + amplitude_cache=None, ): torch = _require_torch() device = _model_device(driver.model) @@ -385,11 +386,18 @@ def _bridge_samples( if amplitude_floor < 0: raise ValueError("amplitude_floor must be non-negative.") with torch.no_grad(): - amplitudes = _call_amplitude_fn( - driver.model, - configs, - chunk_size=getattr(driver, "chunk_size", None), - ) + if amplitude_cache is None: + amplitudes = _call_amplitude_fn( + driver.model, + configs, + chunk_size=getattr(driver, "chunk_size", None), + ) + else: + amplitudes = amplitude_cache.evaluate( + driver.model, + configs, + chunk_size=getattr(driver, "chunk_size", None), + ) amp_abs = amplitudes.abs() valid = torch.isfinite(amp_abs) & (amp_abs > float(amplitude_floor)) & torch.isfinite(log_q) if not bool(torch.any(valid)): @@ -413,6 +421,7 @@ def sample_from_proposal( sample_kwargs=None, progress=False, amplitude_floor=0.0, + amplitude_cache=None, ): """Draw and bridge reusable samples from an MPS, BP, tree, or proposal. @@ -436,6 +445,7 @@ def sample_from_proposal( progress=progress, sample_kwargs=sample_kwargs, amplitude_floor=amplitude_floor, + amplitude_cache=amplitude_cache, ) elapsed = time.perf_counter() - start n_valid = int(configs.shape[0]) @@ -467,6 +477,7 @@ def measure_from_proposal( amplitude_floor=0.0, profile=False, deduplicate=True, + amplitude_cache=None, ): """Compatibility one-shot wrapper around sample then measure.""" samples = sample_from_proposal( @@ -481,6 +492,7 @@ def measure_from_proposal( sample_kwargs=sample_kwargs, progress=progress, amplitude_floor=amplitude_floor, + amplitude_cache=amplitude_cache, ) return driver.measure_samples( samples, @@ -488,6 +500,7 @@ def measure_from_proposal( profile=profile, deduplicate=deduplicate, progress=progress, + amplitude_cache=amplitude_cache, ) diff --git a/tests/test_vmc_api.py b/tests/test_vmc_api.py index 332f022..34a00d1 100644 --- a/tests/test_vmc_api.py +++ b/tests/test_vmc_api.py @@ -994,6 +994,36 @@ def test_torch_ctmrg_preserves_explicit_stabilization_options(monkeypatch): } +def test_torch_native_symmray_ctmrg_defaults_to_projector_mode(): + """Native Torch Symmray CTMRG should use the guarded projector route.""" + torch = pytest.importorskip("torch") + pytest.importorskip("symmray") + from pepsy.tensors import SymPEPS, site_charge_from_occupations + from pepsy.vmc.torch.amplitude import TorchPEPSAmplitude + + state = SymPEPS.random( + 2, + 2, + symmetry="U1", + phys_dim={0: 1, 1: 2, 2: 1}, + fermionic=True, + site_charge=site_charge_from_occupations( + {(x, y): 1 for x in range(2) for y in range(2)} + ), + bond_dim=2, + seed=194, + dtype="complex128", + ) + model = TorchPEPSAmplitude( + state, + contraction="ctmrg", + chi=2, + dtype=torch.complex128, + ) + + assert model.contraction_opts["mode"] == "projector" + + def test_netket_setup_consumes_shared_sampling_config(): nk = pytest.importorskip("netket") from pepsy.vmc.netket import NetKetPEPSVMC diff --git a/tests/test_vmc_transition_plan.py b/tests/test_vmc_transition_plan.py new file mode 100644 index 0000000..4d1ab84 --- /dev/null +++ b/tests/test_vmc_transition_plan.py @@ -0,0 +1,74 @@ +def test_transition_plan_matches_live_connections_and_slices(): + torch = __import__("torch") + from pepsy.vmc.torch import TorchConnections, TorchVMCDriver + + class Amplitude(torch.nn.Module): + def __init__(self): + super().__init__() + self.scale = torch.nn.Parameter(torch.tensor(1.0)) + + def forward(self, configs): + return self.scale * (configs.to(torch.float64).sum(dim=1) + 1.0) + + def connections(configs, graph): + del graph + flipped = configs.clone() + flipped[:, 0] = 1 - flipped[:, 0] + batch_ids = torch.arange(configs.shape[0], device=configs.device) + return TorchConnections( + configs=torch.cat((configs, flipped)), + coeffs=torch.ones(2 * configs.shape[0], device=configs.device), + batch_ids=torch.cat((batch_ids, batch_ids)), + ) + + configs = torch.tensor([[0, 1], [1, 0], [0, 0]], dtype=torch.long) + driver = TorchVMCDriver(Amplitude(), None, configs, connection_fn=connections) + plan = driver.compile_fock_plan(configs, observables={"obs": None}) + live = driver.measure_samples(configs, observables={"obs": None}) + planned = driver.measure_samples( + configs, observables={"obs": None}, connection_plan=plan, + ) + assert torch.allclose(live["obs"].local_energies, planned["obs"].local_energies) + + sliced = plan.slice(1, 3) + sliced_result = driver.measure_samples( + configs[1:], observables={"obs": None}, connection_plan=sliced, + ) + assert torch.allclose( + planned["obs"].local_energies[:, 1:], sliced_result["obs"].local_energies, + ) + + +def test_amplitude_cache_reuses_targets_and_invalidates_on_parameter_update(): + torch = __import__("torch") + from pepsy.vmc.torch import TorchAmplitudeCache, TorchConnections, TorchVMCDriver + + class Amplitude(torch.nn.Module): + def __init__(self): + super().__init__() + self.scale = torch.nn.Parameter(torch.tensor(1.0)) + + def forward(self, configs): + return self.scale * (configs.to(torch.float64).sum(dim=1) + 1.0) + + configs = torch.tensor([[0, 1], [1, 0], [0, 0]], dtype=torch.long) + driver = TorchVMCDriver( + Amplitude(), None, configs, + connection_fn=lambda rows, graph: TorchConnections( + configs=rows, + coeffs=torch.ones(rows.shape[0]), + batch_ids=torch.arange(rows.shape[0]), + ), + ) + cache = TorchAmplitudeCache(max_entries=32) + with torch.no_grad(): + first = cache.evaluate(driver.model, configs) + second = cache.evaluate(driver.model, configs) + assert torch.equal(first, second) + assert cache.snapshot()["hits"] == 3 + with torch.no_grad(): + driver.model.scale.mul_(2.0) + with torch.no_grad(): + updated = cache.evaluate(driver.model, configs) + assert torch.allclose(updated, 2.0 * first) + assert cache.snapshot()["misses"] == 3 From 78666d2c2b02335ac8336a50513b3d422f59d937 Mon Sep 17 00:00:00 2001 From: rezaquant Date: Wed, 5 Aug 2026 10:26:22 -0600 Subject: [PATCH 70/70] Improve Symmray MPS sampling strategy --- docs/api/sampling/samplers.md | 27 ++- src/pepsy/sampling/samplers.py | 329 +++++++++++++++++++++++++++++++-- tests/test_sampler.py | 101 ++++++++++ 3 files changed, 431 insertions(+), 26 deletions(-) diff --git a/docs/api/sampling/samplers.md b/docs/api/sampling/samplers.md index aac1f3c..2d46810 100644 --- a/docs/api/sampling/samplers.md +++ b/docs/api/sampling/samplers.md @@ -101,21 +101,30 @@ constructing the sampler: sampler = MpsSampler( psi, backend="symmray", - prefix_strategy="auto", # "prefix" or "serial" are also available + strategy="auto", # "prefix", "serial", or "dense" are also available max_prefix_groups=256, + dense_memory_limit="256MiB", ) configs, probs = sampler.sample_arrays(4096, seed=0) print(sampler.symmray_sampling_stats) ``` -`"auto"` shares equal prefixes while they still amortize a boundary: singleton -prefixes are completed serially, and retained groups obey both the active-group -cap and a per-level block-storage budget. `"prefix"` keeps every group allowed -by `max_prefix_groups`, including singletons; `"serial"` retains one boundary -at a time. The statistics report distinct conditional distributions, -candidate contractions, charge-pruned branches, peak active groups, and -serial/adaptive fallbacks. Set `max_prefix_groups=None` to remove the hard -group cap while retaining the `"auto"` reuse decision. +`"auto"` selects the fully batched dense kernel when the batch has at least +`dense_min_samples` shots and the estimated dense MPS view fits within +`dense_memory_limit`; otherwise it shares equal sparse prefixes. `"prefix"` +keeps every group allowed by `max_prefix_groups`, including singletons; +`"serial"` retains one boundary at a time. `"dense"` materializes a private +dense view of the source MPS and uses the backend-native batched conditional +kernel. It is supported for resolved fermionic U1/U1U1 states; parity-collapsed +Z2/Z2Z2 states remain on the sparse charge-aware route. Dense batching can be +substantially faster for moderate bond dimensions and high-entropy batches, +but uses more memory. The statistics report the requested and selected +strategy, dense-memory estimate, conditional distributions, candidate +contractions, charge-pruned branches, peak active groups, and serial/adaptive +fallbacks. Set `max_prefix_groups=None` to remove the hard group cap while +retaining sparse prefix sampling. + +`prefix_strategy=` remains a backward-compatible alias for `strategy=`. For comparable throughput measurements, call the public sampling APIs from an external benchmark harness. For fermionic `Z2`/`Z2Z2` inputs, do not treat a naive dense expansion of graded virtual legs as a state-preserving conversion; diff --git a/src/pepsy/sampling/samplers.py b/src/pepsy/sampling/samplers.py index 8532bff..2efeac6 100644 --- a/src/pepsy/sampling/samplers.py +++ b/src/pepsy/sampling/samplers.py @@ -98,6 +98,9 @@ def _normalize_symmray_prefix_strategy(strategy): "shared_prefix": "prefix", "serial": "serial", "one_by_one": "serial", + "dense": "dense", + "dense_batch": "dense", + "batched_dense": "dense", } try: return aliases[key] @@ -109,6 +112,43 @@ def _normalize_symmray_prefix_strategy(strategy): ) from exc +def _normalize_dense_memory_limit(limit): + """Normalize a dense sampling memory budget to bytes.""" + if limit is None: + return None + if isinstance(limit, (int, np.integer)): + limit = int(limit) + else: + text = str(limit).strip().upper().replace(" ", "") + if text in {"NONE", "UNBOUNDED", "INF", "INFINITY"}: + return None + units = ( + ("GIB", 1024**3), + ("GB", 1000**3), + ("MIB", 1024**2), + ("MB", 1000**2), + ("KIB", 1024), + ("KB", 1000), + ("B", 1), + ) + multiplier = 1 + for suffix, factor in units: + if text.endswith(suffix): + text = text[:-len(suffix)] + multiplier = factor + break + try: + limit = int(float(text) * multiplier) + except ValueError as exc: + raise TypeError( + "dense_memory_limit must be bytes, a size such as '256MiB', " + "or None." + ) from exc + if limit < 1: + raise ValueError("dense_memory_limit must be positive or None.") + return int(limit) + + def _mps_array_backend(array): module = type(array).__module__.split(".", 1)[0] if module == "torch": @@ -765,16 +805,29 @@ class MpsSampler: Opt into ``torch.compile`` for repeated, device-resident, unseeded Torch inference batches. Unsupported compiler environments and calls that need eager-only behavior fall back to eager sampling. - prefix_strategy : {"auto", "prefix", "serial"}, default="auto" + strategy : {"auto", "prefix", "serial", "dense"}, optional + Preferred name for the Symmray sampling strategy. ``None`` leaves + ``prefix_strategy`` in control for backward compatibility. + prefix_strategy : {"auto", "prefix", "serial", "dense"}, default="auto" Symmray batch-sampling strategy. ``"prefix"`` shares a normalized block-sparse boundary between equal sampled prefixes; ``"serial"`` uses one independent left-to-right sweep per shot. ``"auto"`` uses prefix sharing until ``max_prefix_groups`` is reached, then finishes the remaining branches serially with bounded memory. + ``"dense"`` creates a temporary dense view of the source MPS and + uses the backend-native fully batched sampler. ``"auto"`` selects + dense batching when the sample count and memory budget permit it. + Dense batching can use more memory than the sparse routes. max_prefix_groups : int or None, default=256 Maximum active Symmray prefix groups before the ``"auto"`` strategy switches the remaining suffixes to serial sampling. ``None`` permits all distinct prefixes. This has no effect on dense MPS backends. + dense_memory_limit : int, str, or None, default="256MiB" + Maximum estimated dense MPS storage allowed by ``strategy="auto"`` or + ``strategy="dense"``. Strings such as ``"256MiB"`` and ``"1GB"`` are + accepted. ``None`` disables the guard. + dense_min_samples : int, default=1024 + Minimum batch size for ``strategy="auto"`` to select dense batching. fermion : pepsy.tensors.Fermion, optional Fermionic physical-space convention associated with this sampler. When supplied, :meth:`sample_batch` attaches its symmetry-aware @@ -801,8 +854,11 @@ def __init__( *, backend: str | None = "quimb", torch_compile: bool = False, + strategy: str | None = None, prefix_strategy: str = "auto", max_prefix_groups: int | None = 256, + dense_memory_limit: int | str | None = 256 * 1024**2, + dense_min_samples: int = 1024, fermion=None, ): if one_d_to_two_d is None: @@ -825,6 +881,12 @@ def __init__( if not isinstance(torch_compile, (bool, np.bool_)): raise TypeError("torch_compile must be a boolean.") self.torch_compile = bool(torch_compile) + if strategy is not None: + if prefix_strategy not in (None, "auto"): + raise ValueError( + "Pass either strategy= or prefix_strategy=, not both." + ) + prefix_strategy = strategy self.prefix_strategy = _normalize_symmray_prefix_strategy(prefix_strategy) if max_prefix_groups is not None: if not isinstance(max_prefix_groups, (int, np.integer)): @@ -835,6 +897,15 @@ def __init__( ) max_prefix_groups = int(max_prefix_groups) self.max_prefix_groups = max_prefix_groups + self.dense_memory_limit = _normalize_dense_memory_limit(dense_memory_limit) + if not isinstance(dense_min_samples, (int, np.integer)): + raise TypeError("dense_min_samples must be a positive integer.") + if int(dense_min_samples) < 1: + raise ValueError("dense_min_samples must be a positive integer.") + self.dense_min_samples = int(dense_min_samples) + # ``strategy`` is the preferred public spelling; retain the old + # attribute for callers that inspect prefix_strategy directly. + self.strategy = self.prefix_strategy self.fermion = fermion self.resolved_backend = None self._source_psi = None @@ -1191,10 +1262,13 @@ def _prepare_symmray_state(psi): return { "sr": sr, "mps": canonical, + "source_mps": psi, "sites": tuple(sites), "physical_code_maps": tuple(physical_code_maps), "array_backend": array_backend, "template": template, + "dense_site_data": None, + "dense_code_maps": None, } @staticmethod @@ -1472,6 +1546,208 @@ def _symmray_candidates(cls, state, site, boundary): weights.append(cls._symmray_weight(candidate, state)) return site_state, local_codes, candidates, weights + @staticmethod + def _dense_array_nbytes(array): + """Estimate dense storage for a backend array without copying it.""" + nbytes = getattr(array, "nbytes", None) + if nbytes is not None: + return int(nbytes) + try: + return int(array.numel()) * int(array.element_size()) + except (AttributeError, TypeError, ValueError): + return int(np.asarray(array).nbytes) + + @classmethod + def _symmray_estimate_dense_site_bytes(cls, state): + """Estimate dense MPS storage from shapes without materializing it.""" + template = state["template"] + if hasattr(template, "element_size"): + itemsize = int(template.element_size()) + else: + itemsize = int(np.dtype(getattr(template, "dtype", template)).itemsize) + total = 0 + source_mps = state["source_mps"] + for site in range(len(state["sites"])): + array = cls._site_array_lr_phys_r(source_mps, site) + total += int(np.prod(array.shape)) * itemsize + return int(total) + + def _resolve_symmray_sampling_strategy(self, n_samples): + """Resolve the requested strategy before any dense allocation.""" + requested = self.prefix_strategy + state = self._require_symmray_state() + symmetry = str(state["sites"][0]["data"].symmetry).upper() + dense_supported = symmetry in {"U1", "U1U1"} + if not dense_supported: + if requested == "dense": + raise ValueError( + "Dense Symmray sampling is supported only for resolved " + f"U1/U1U1 states, not symmetry={symmetry!r}. Use " + "strategy='prefix' for charge-aware sampling." + ) + if requested == "auto": + return "auto", "auto_sparse_unsupported_symmetry", None + return requested, "explicit_sparse", None + estimated_bytes = self._symmray_estimate_dense_site_bytes(state) + if requested == "dense": + if ( + self.dense_memory_limit is not None + and estimated_bytes > self.dense_memory_limit + ): + raise ValueError( + "Dense Symmray sampling requires an estimated " + f"{estimated_bytes} bytes, above the configured limit of " + f"{self.dense_memory_limit} bytes. Increase " + "dense_memory_limit or use strategy='prefix'." + ) + return "dense", "explicit_dense", estimated_bytes + if requested == "auto": + if ( + int(n_samples) >= self.dense_min_samples + and ( + self.dense_memory_limit is None + or estimated_bytes <= self.dense_memory_limit + ) + ): + return "dense", "auto_dense_within_budget", estimated_bytes + return "auto", "auto_sparse_fallback", estimated_bytes + return requested, "explicit_sparse", estimated_bytes + + @classmethod + def _symmray_dense_site_data(cls, state): + """Prepare cached dense site operators for explicit dense batching. + + This route is deliberately opt-in. It keeps the source and canonical + Symmray states intact, materializing only a private sampling view so + the dense native sampler can contract every shot in one backend batch. + """ + cached = state.get("dense_site_data") + if cached is not None: + return cached + + arrays = [] + source_mps = state["source_mps"] + for site in range(len(state["sites"])): + # Use the source MPS rather than the private canonical copy here. + # Symmray's fermionic bond orientations can have different dense + # positional layouts on dual virtual legs even though sparse + # charge-aware contractions remain valid. The source chain has + # matching virtual dimensions, so its dense view is unambiguous. + array = cls._site_array_lr_phys_r(source_mps, site) + if hasattr(array, "to_dense"): + array = array.to_dense() + arrays.append(array) + + backends = {_mps_array_backend(array) for array in arrays} + if len(backends) != 1: + raise ValueError( + "Dense Symmray sampling requires one common dense backend; " + f"got {sorted(backends)!r}." + ) + backend = next(iter(backends)) + if backend == "torch": + site_data = cls._torch_site_ops(tuple(arrays)) + elif backend in {"numpy", "cupy"}: + site_data = cls._array_namespace_site_ops( + tuple(arrays), + backend=backend, + ) + else: + raise ValueError( + "Dense Symmray sampling produced unsupported arrays " + f"with backend {backend!r}." + ) + + dense_bytes = sum(cls._dense_array_nbytes(array) for array in arrays) + code_maps = tuple( + tuple(range(int(array.shape[1]))) + for array in arrays + ) + cached = (backend, site_data, int(dense_bytes)) + state["dense_site_data"] = cached + state["dense_code_maps"] = code_maps + return cached + + @staticmethod + def _symmray_map_dense_configs(configs, state): + """Map canonical dense physical choices back to source code labels.""" + backend = state["array_backend"] + if backend == "torch": + import torch # pylint: disable=import-outside-toplevel + + mapped = torch.empty_like(configs) + for site, code_map in enumerate(state["dense_code_maps"]): + lookup = torch.as_tensor( + code_map, + dtype=torch.long, + device=configs.device, + ) + mapped[:, site] = lookup[configs[:, site]] + return mapped + if backend == "cupy": + import cupy as cp # pylint: disable=import-outside-toplevel + + mapped = cp.empty_like(configs) + for site, code_map in enumerate(state["dense_code_maps"]): + lookup = cp.asarray(code_map, dtype=cp.int64) + mapped[:, site] = lookup[configs[:, site]] + return mapped + + mapped = np.empty_like(configs) + for site, code_map in enumerate(state["dense_code_maps"]): + mapped[:, site] = np.asarray(code_map, dtype=np.int64)[ + configs[:, site] + ] + return mapped + + @classmethod + def _symmray_sample_arrays_dense( + cls, + state, + n_samples, + seed, + *, + to_numpy, + ): + """Sample a Symmray MPS with the dense native batched kernels.""" + backend, site_data, dense_bytes = cls._symmray_dense_site_data(state) + if backend == "torch": + canonical_configs, probabilities = cls._torch_sample( + site_data, + int(n_samples), + seed, + to_numpy=False, + ) + else: + canonical_configs, probabilities = cls._array_namespace_sample( + site_data, + int(n_samples), + seed, + backend=backend, + to_numpy=False, + ) + configs = cls._symmray_map_dense_configs(canonical_configs, state) + stats = { + "strategy": "dense", + "n_samples": int(n_samples), + "conditional_evaluations": len(state["sites"]), + "candidate_contractions": sum( + len(site_state["codes"]) for site_state in state["sites"] + ), + "static_pruned_branches": 0, + "charge_pruned_branches": 0, + "cached_local_slices": False, + "max_active_prefix_groups": 1, + "serial_fallback": False, + "adaptive_serial_fallback": False, + "dense_site_bytes": int(dense_bytes), + "dense_batch_width": int(n_samples), + } + if to_numpy: + configs = _backend_array_to_numpy(configs) + probabilities = _backend_array_to_numpy(probabilities) + return configs, probabilities, stats + @staticmethod def _symmray_note_candidates(stats, site_state, local_codes): """Record the sparse branch work avoided by cache/pruning.""" @@ -1820,6 +2096,13 @@ def _symmray_sample_arrays( max_prefix_groups, to_numpy, ): + if strategy == "dense": + return cls._symmray_sample_arrays_dense( + state, + n_samples, + seed, + to_numpy=to_numpy, + ) if strategy == "serial": return cls._symmray_sample_arrays_serial( state, @@ -2683,6 +2966,32 @@ def sample_arrays( if not isinstance(track_grad, (bool, np.bool_)): raise TypeError("track_grad must be a boolean.") if self._symmray_state is not None: + def sample_symmray(): + strategy, selection, estimated_bytes = ( + self._resolve_symmray_sampling_strategy(int(n_samples)) + ) + configs, probs, stats = self._symmray_sample_arrays( + self._symmray_state, + int(n_samples), + seed, + strategy=strategy, + max_prefix_groups=self.max_prefix_groups, + to_numpy=to_numpy, + ) + stats.update( + { + "requested_strategy": self.prefix_strategy, + "strategy_selection": selection, + "estimated_dense_site_bytes": ( + None + if estimated_bytes is None + else int(estimated_bytes) + ), + "dense_memory_limit_bytes": self.dense_memory_limit, + } + ) + return configs, probs, stats + if ( self._symmray_state["array_backend"] == "torch" and not track_grad @@ -2690,24 +2999,10 @@ def sample_arrays( import torch # pylint: disable=import-outside-toplevel with torch.no_grad(): - configs, probs, stats = self._symmray_sample_arrays( - self._symmray_state, - int(n_samples), - seed, - strategy=self.prefix_strategy, - max_prefix_groups=self.max_prefix_groups, - to_numpy=to_numpy, - ) + configs, probs, stats = sample_symmray() self._last_symmray_sampling_stats = stats return configs, probs - configs, probs, stats = self._symmray_sample_arrays( - self._symmray_state, - int(n_samples), - seed, - strategy=self.prefix_strategy, - max_prefix_groups=self.max_prefix_groups, - to_numpy=to_numpy, - ) + configs, probs, stats = sample_symmray() self._last_symmray_sampling_stats = stats return configs, probs if self._native_arrays is not None: diff --git a/tests/test_sampler.py b/tests/test_sampler.py index 4a7e58f..dafa5bc 100644 --- a/tests/test_sampler.py +++ b/tests/test_sampler.py @@ -1476,6 +1476,42 @@ def test_mps_sampler_symmray_prefix_controls_bound_high_entropy_batches(): ) +@pytest.mark.parametrize("symmetry", ("U1", "U1U1")) +def test_mps_sampler_symmray_dense_strategy_uses_batched_native_kernel(symmetry): + """The explicit dense strategy batches all Symmray shots safely.""" + pytest.importorskip("symmray") + from pepsy.tensors import Fermion, SymMPS + + fermion = Fermion(spinful=True, symmetry=symmetry) + psi = SymMPS.random( + 4, + symmetry=symmetry, + fermionic=True, + phys_dim=fermion.physical_sectors, + site_charge=fermion.half_filled_site_charge(4), + bond_dim=4, + seed=29, + dtype="complex128", + ).mps + sampler = sampler_mod.MpsSampler( + psi, + backend="symmray", + prefix_strategy="dense", + ) + + configs, sampled_probs = sampler.sample_arrays(64, seed=7) + stats = sampler.symmray_sampling_stats + + assert stats["strategy"] == "dense" + assert stats["conditional_evaluations"] == psi.L + assert stats["dense_site_bytes"] > 0 + np.testing.assert_allclose( + sampler.probabilities(configs), + sampled_probs, + atol=1e-12, + ) + + def test_mps_sampler_rejects_invalid_symmray_prefix_controls(): """Prefix controls should be explicit before any MPS preprocessing.""" psi = qtn.MPS_computational_state("0") @@ -1486,6 +1522,71 @@ def test_mps_sampler_rejects_invalid_symmray_prefix_controls(): sampler_mod.MpsSampler(psi, max_prefix_groups=0) +def test_mps_sampler_auto_selects_dense_within_memory_budget(): + """Auto strategy should use dense batching for a sufficiently large batch.""" + pytest.importorskip("symmray") + from pepsy.tensors import Fermion, SymMPS + + fermion = Fermion(spinful=True, symmetry="U1") + psi = SymMPS.random( + 4, + symmetry="U1", + fermionic=True, + phys_dim=fermion.physical_sectors, + site_charge=fermion.half_filled_site_charge(4), + bond_dim=4, + seed=43, + dtype="complex128", + ).mps + sampler = sampler_mod.MpsSampler( + psi, + strategy="auto", + dense_min_samples=32, + dense_memory_limit="1GiB", + ) + + configs, probs = sampler.sample_arrays(64, seed=5) + stats = sampler.symmray_sampling_stats + + assert stats["requested_strategy"] == "auto" + assert stats["strategy_selection"] == "auto_dense_within_budget" + assert stats["strategy"] == "dense" + np.testing.assert_allclose(sampler.probabilities(configs), probs, atol=1e-12) + + +def test_mps_sampler_rejects_conflicting_strategy_aliases_and_dense_budget(): + """The new strategy alias and dense guard should fail clearly.""" + psi = qtn.MPS_computational_state("0") + with pytest.raises(ValueError, match="either strategy"): + sampler_mod.MpsSampler( + psi, + strategy="dense", + prefix_strategy="serial", + ) + + pytest.importorskip("symmray") + from pepsy.tensors import Fermion, SymMPS + + fermion = Fermion(spinful=True, symmetry="U1") + symm_psi = SymMPS.random( + 3, + symmetry="U1", + fermionic=True, + phys_dim=fermion.physical_sectors, + site_charge=fermion.half_filled_site_charge(3), + bond_dim=3, + seed=47, + dtype="complex128", + ).mps + guarded = sampler_mod.MpsSampler( + symm_psi, + strategy="dense", + dense_memory_limit="1B", + ) + with pytest.raises(ValueError, match="above the configured limit"): + guarded.sample_arrays(8, seed=1) + + def test_mps_sampler_symmray_torch_nonfermionic_blocks_stay_on_torch(): """Generic Symmray Torch blocks preserve device-resident outputs.""" pytest.importorskip("symmray")